#!/usr/bin/env python3
"""Extract human pose landmarks from a video (or single image) with MediaPipe Pose
(Tasks API, mediapipe >= 0.10.30 — the legacy `solutions` API no longer exists).

Outputs one JSON document with per-sample 2D landmarks (normalized image coords)
and 3D world landmarks (meters, hip-centered) — the world landmarks are what you
want for pose comparison because they are camera-scale invariant.

Sampling: --fps N samples the video at N Hz (default 10). --times "0.0,0.5,1.0"
samples at explicit timestamps instead (e.g. a musical beat grid). --overlay
writes the sampled frames with the skeleton drawn to an .mp4 for eyeball checks.

Works on real footage and (usually) on stylized humanoid game characters, but
detection confidence drops on non-photoreal rigs — lower --min-conf and try
--model heavy if frames come back undetected.
"""
import argparse, json, os, sys

import cv2
import numpy as np
import mediapipe as mp
from mediapipe.tasks import python as mp_tasks
from mediapipe.tasks.python import vision

SKILL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# The 33 pose landmarks, in model output order (fixed contract of the pose model).
POSE_LANDMARKS = [
    "nose", "left_eye_inner", "left_eye", "left_eye_outer", "right_eye_inner",
    "right_eye", "right_eye_outer", "left_ear", "right_ear", "mouth_left",
    "mouth_right", "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
    "left_wrist", "right_wrist", "left_pinky", "right_pinky", "left_index",
    "right_index", "left_thumb", "right_thumb", "left_hip", "right_hip",
    "left_knee", "right_knee", "left_ankle", "right_ankle", "left_heel",
    "right_heel", "left_foot_index", "right_foot_index",
]

# Skeleton edges for overlay drawing (standard MediaPipe pose connections, torso+limbs).
CONNECTIONS = [
    (11, 12), (11, 13), (13, 15), (12, 14), (14, 16),          # shoulders + arms
    (11, 23), (12, 24), (23, 24),                              # torso
    (23, 25), (25, 27), (27, 29), (27, 31),                    # left leg
    (24, 26), (26, 28), (28, 30), (28, 32),                    # right leg
    (15, 17), (15, 19), (15, 21), (16, 18), (16, 20), (16, 22),  # hands
]


def lms_to_list(lms):
    return [
        {"name": POSE_LANDMARKS[i],
         "x": round(p.x, 5), "y": round(p.y, 5), "z": round(p.z, 5),
         "visibility": round(p.visibility or 0.0, 3)}
        for i, p in enumerate(lms)
    ]


def draw_skeleton(frame_bgr, lms, t, detected):
    h, w = frame_bgr.shape[:2]
    if detected:
        pts = [(int(p.x * w), int(p.y * h)) for p in lms]
        for a, b in CONNECTIONS:
            cv2.line(frame_bgr, pts[a], pts[b], (0, 220, 255), 2)
        for x, y in pts:
            cv2.circle(frame_bgr, (x, y), 3, (0, 255, 0), -1)
    cv2.putText(frame_bgr, f"t={t:.2f}s {'OK' if detected else 'NO POSE'}",
                (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8,
                (0, 255, 0) if detected else (0, 0, 255), 2)
    return frame_bgr


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("input", help="video file (.mp4/.mov/...) or single image")
    ap.add_argument("--out", required=True, help="output JSON path")
    ap.add_argument("--fps", type=float, default=10.0, help="sample rate in Hz (default 10)")
    ap.add_argument("--times", default=None, help="comma-separated timestamps in seconds (overrides --fps)")
    ap.add_argument("--overlay", default=None, help="optional .mp4 path: sampled frames with skeleton drawn")
    ap.add_argument("--model", default="full", choices=["lite", "full", "heavy"],
                    help="pose model variant (heavy = most accurate; auto-downloads)")
    ap.add_argument("--min-conf", type=float, default=0.5,
                    help="min detection/tracking confidence (lower for stylized characters)")
    args = ap.parse_args()

    model_path = os.path.join(SKILL_DIR, "models", f"pose_landmarker_{args.model}.task")
    if not os.path.exists(model_path):
        url = (f"https://storage.googleapis.com/mediapipe-models/pose_landmarker/"
               f"pose_landmarker_{args.model}/float16/latest/pose_landmarker_{args.model}.task")
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        print(f"downloading model {args.model}...", file=sys.stderr)
        import subprocess
        subprocess.run(["curl", "-sL", "-o", model_path, url], check=True)

    is_image = os.path.splitext(args.input)[1].lower() in (".png", ".jpg", ".jpeg", ".webp", ".bmp")

    options = vision.PoseLandmarkerOptions(
        base_options=mp_tasks.BaseOptions(model_asset_path=model_path),
        running_mode=vision.RunningMode.IMAGE if is_image else vision.RunningMode.VIDEO,
        min_pose_detection_confidence=args.min_conf,
        min_pose_presence_confidence=args.min_conf,
        min_tracking_confidence=args.min_conf,
    )
    landmarker = vision.PoseLandmarker.create_from_options(options)

    frames_out, overlay_frames = [], []

    def process(frame_bgr, t, idx):
        mp_img = mp.Image(image_format=mp.ImageFormat.SRGB,
                          data=cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
        if is_image:
            res = landmarker.detect(mp_img)
        else:
            res = landmarker.detect_for_video(mp_img, int(t * 1000))
        detected = bool(res.pose_landmarks)
        entry = {"t": round(t, 4), "index": idx, "detected": detected}
        if detected:
            entry["landmarks"] = lms_to_list(res.pose_landmarks[0])
            if res.pose_world_landmarks:
                entry["world_landmarks"] = lms_to_list(res.pose_world_landmarks[0])
        frames_out.append(entry)
        if args.overlay is not None:
            overlay_frames.append(draw_skeleton(
                frame_bgr.copy(), res.pose_landmarks[0] if detected else None, t, detected))

    if is_image:
        img = cv2.imread(args.input)
        if img is None:
            sys.exit(f"ERROR: cannot read image {args.input}")
        process(img, 0.0, 0)
        src_fps, duration = 0.0, 0.0
    else:
        cap = cv2.VideoCapture(args.input)
        if not cap.isOpened():
            sys.exit(f"ERROR: cannot open video {args.input}")
        src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
        duration = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) / src_fps if src_fps else 0.0

        if args.times:
            sample_ts = sorted(float(s) for s in args.times.split(",") if s.strip())
        else:
            sample_ts = list(np.arange(0.0, max(duration - 1e-6, 0.0), 1.0 / args.fps))

        for i, t in enumerate(sample_ts):  # ascending — VIDEO mode needs monotonic timestamps
            cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000.0)
            ok, frame = cap.read()
            if not ok:
                frames_out.append({"t": round(t, 4), "index": i, "detected": False, "error": "read failed"})
                continue
            process(frame, t, i)
        cap.release()

    landmarker.close()

    if args.overlay and overlay_frames:
        h, w = overlay_frames[0].shape[:2]
        vw = cv2.VideoWriter(args.overlay, cv2.VideoWriter_fourcc(*"mp4v"),
                             args.fps if not args.times else 4.0, (w, h))
        for f in overlay_frames:
            vw.write(f)
        vw.release()

    detected = sum(1 for f in frames_out if f.get("detected"))
    doc = {
        "source": os.path.abspath(args.input),
        "source_fps": round(src_fps, 3),
        "duration_s": round(duration, 3),
        "samples": len(frames_out),
        "detected": detected,
        "landmark_names": POSE_LANDMARKS,
        "frames": frames_out,
    }
    os.makedirs(os.path.dirname(os.path.abspath(args.out)) or ".", exist_ok=True)
    with open(args.out, "w") as f:
        json.dump(doc, f)
    print(f"OK: {detected}/{len(frames_out)} samples had a detected pose -> {args.out}"
          + (f" (overlay: {args.overlay})" if args.overlay else ""))


if __name__ == "__main__":
    main()