"""MotionDNA — extrae secuencias de pose 2D (17 keypoints COCO) de los clips de
music video usando YOLO-pose. Guarda un .npy [T,17,2] (xy normalizado) por clip."""
import os, glob, sys
import numpy as np
import cv2
from ultralytics import YOLO

VIDEOS = os.path.expanduser("~/cosmos-predict2/datasets/boostify_mv_demo/videos")
OUT = os.path.expanduser("~/motiondna/poses")
os.makedirs(OUT, exist_ok=True)

model = YOLO("yolov8n-pose.pt")  # se descarga solo la 1a vez
clips = sorted(glob.glob(os.path.join(VIDEOS, "*.mp4")))
print(f"Clips: {len(clips)}", flush=True)

def biggest_person(res):
    if res.keypoints is None or res.keypoints.xyn is None: return None
    kp = res.keypoints.xyn.cpu().numpy()  # [n,17,2]
    if kp.shape[0] == 0: return None
    if res.boxes is not None and res.boxes.conf is not None and len(res.boxes.conf)==kp.shape[0]:
        idx = int(res.boxes.conf.cpu().numpy().argmax())
    else:
        idx = 0
    return kp[idx]  # [17,2]

total_frames = 0
for ci, path in enumerate(clips):
    cap = cv2.VideoCapture(path)
    seq = []; last = None
    while True:
        ok, frame = cap.read()
        if not ok: break
        res = model.predict(frame, verbose=False, device=0)[0]
        kp = biggest_person(res)
        if kp is None:
            if last is None: continue
            kp = last
        last = kp
        seq.append(kp.astype(np.float32))
    cap.release()
    if len(seq) >= 64:
        arr = np.stack(seq)  # [T,17,2]
        np.save(os.path.join(OUT, os.path.basename(path).replace(".mp4",".npy")), arr)
        total_frames += len(seq)
    print(f"  [{ci+1}/{len(clips)}] {os.path.basename(path)} -> {len(seq)} frames", flush=True)
print(f"LISTO. Secuencias guardadas. Frames totales: {total_frames}", flush=True)