import os
import argparse
import pandas as pd
import cv2
import numpy as np

# Skeleton definition
SKELETON = [
    ("left_ankle", "left_knee"), ("left_knee", "left_hip"), 
    ("right_ankle", "right_knee"), ("right_knee", "right_hip"),
    ("left_hip", "right_hip"), ("left_shoulder", "right_shoulder"),
    ("left_shoulder", "left_hip"), ("right_shoulder", "right_hip"),
    ("left_shoulder", "left_elbow"), ("left_elbow", "left_wrist"),
    ("right_shoulder", "right_elbow"), ("right_elbow", "right_wrist"),
    ("left_ankle", "left_heel"), ("left_ankle", "left_big_toe"),
    ("right_ankle", "right_heel"), ("right_ankle", "right_big_toe")
]

def annotate_frames(frames_dir, output_dir, csv_path):
    """Draw keypoints and skeleton on frames."""
    df = pd.read_csv(csv_path)
    os.makedirs(output_dir, exist_ok=True)
    
    frame_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.jpg')])
    
    for _, row in df.iterrows():
        f_idx = int(row['frame'])
        if f_idx >= len(frame_files): continue
        
        frame_name = frame_files[f_idx]
        img = cv2.imread(os.path.join(frames_dir, frame_name))
        
        # Draw Skeleton
        for start_kp, end_kp in SKELETON:
            s_x, s_y = row.get(f"{start_kp}_x"), row.get(f"{start_kp}_y")
            e_x, e_y = row.get(f"{end_kp}_x"), row.get(f"{end_kp}_y")
            if pd.notna(s_x) and pd.notna(e_x):
                cv2.line(img, (int(s_x), int(s_y)), (int(e_x), int(e_y)), (0, 255, 0), 2)
        
        # Draw Keypoints
        for col in df.columns:
            if col.endswith('_x'):
                kp_name = col[:-2]
                x, y = row[col], row[f"{kp_name}_y"]
                if pd.notna(x):
                    cv2.circle(img, (int(x), int(y)), 4, (0, 0, 255), -1)
                    
        cv2.imwrite(os.path.join(output_dir, frame_name), img)
        if f_idx % 100 == 0:
            print(f"Annotated {f_idx} frames")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--frames", required=True, help="Input frames directory")
    parser.add_argument("--csv", required=True, help="Input keypoints CSV")
    parser.add_argument("--output", required=True, help="Output frames directory")
    args = parser.parse_args()
    
    annotate_frames(args.frames, args.output, args.csv)