#!/opt/homebrew/bin/python3.11
"""
Test MediaPipe Pose alignment on face trace key frames

Tests 50 random key frames to determine alignment rate.
Handles multiple poses by matching only the correct one with face bbox.

Usage:
    python3 scripts/test_mediapipe_pose_alignment.py --file-uuid <uuid> --sample-size 50
"""

import argparse
import json
import os
import random
import sys
import time
from pathlib import Path

try:
    import cv2
    import mediapipe as mp
    from mediapipe.tasks.python.vision import PoseLandmarker, PoseLandmarkerOptions
    from mediapipe.tasks.python.core.base_options import BaseOptions
except ImportError as e:
    print(f"Missing dependency: {e}", file=sys.stderr)
    sys.exit(1)


LANDMARK_NAMES = [
    "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",
]


def point_in_bbox(x, y, bbox):
    """Check if point (in pixels) is inside bbox (in pixels)."""
    return (bbox['x'] <= x <= bbox['x'] + bbox['width'] and
            bbox['y'] <= y <= bbox['y'] + bbox['height'])


def get_face_bbox(face_data, trace_id):
    """Get face bbox for a trace from face_traced.json."""
    traces = face_data.get('traces', {})
    frames = face_data.get('frames', {})
    
    if trace_id not in traces:
        return None
    
    trace = traces[trace_id]
    start_frame = str(trace.get('start_frame', 0))
    
    if start_frame not in frames:
        return None
    
    frame_data = frames[start_frame]
    faces = frame_data.get('faces', [])
    
    # Find the face that belongs to this trace
    for face in faces:
        if face.get('trace_id') == int(trace_id):
            return {
                'x': face['x'],
                'y': face['y'],
                'width': face['width'],
                'height': face['height']
            }
    
    # Fallback: use first face
    if faces:
        face = faces[0]
        return {
            'x': face['x'],
            'y': face['y'],
            'width': face['width'],
            'height': face['height']
        }
    
    return None


def check_pose_alignment(pose_landmarks, face_bbox, img_width, img_height):
    """
    Check if pose face keypoints align with face bbox.
    
    Returns: (aligned: bool, face_keypoints: dict)
    """
    # Get face keypoints from pose
    kp_dict = {}
    for idx, landmark in enumerate(pose_landmarks):
        if idx < len(LANDMARK_NAMES):
            name = LANDMARK_NAMES[idx]
            if name in ['nose', 'left_eye', 'right_eye']:
                # Convert normalized coords to pixels
                kp_dict[name] = {
                    'x': landmark.x * img_width,
                    'y': landmark.y * img_height,
                    'confidence': landmark.visibility if hasattr(landmark, 'visibility') else 1.0
                }
    
    # Check if all 3 keypoints are in bbox
    required = ['nose', 'left_eye', 'right_eye']
    if not all(k in kp_dict for k in required):
        return False, kp_dict
    
    for name in required:
        kp = kp_dict[name]
        if not point_in_bbox(kp['x'], kp['y'], face_bbox):
            return False, kp_dict
    
    return True, kp_dict


def test_alignment(file_uuid, output_dir, sample_size=None):
    """Test MediaPipe pose alignment on key frames.
    
    Args:
        sample_size: If None, process all key frames. Otherwise, sample N.
    """
    
    # Load face data
    face_json = Path(output_dir) / f"{file_uuid}.face_traced.json"
    if not face_json.exists():
        print(f"[pose_keyframe] Face file not found: {face_json}", file=sys.stderr)
        return
    
    with open(face_json) as f:
        face_data = json.load(f)
    
    # Get all trace IDs with key_frame.jpg
    output_path = Path(output_dir) / file_uuid
    trace_dirs = sorted(output_path.glob("trace_*"))
    
    valid_traces = []
    for trace_dir in trace_dirs:
        key_frame = trace_dir / "key_frame.jpg"
        if key_frame.exists():
            trace_id = trace_dir.name.replace("trace_", "")
            valid_traces.append(trace_id)
    
    print(f"[pose_keyframe] Found {len(valid_traces)} traces with key_frame.jpg")
    
    # Sample or process all
    if sample_size and len(valid_traces) > sample_size:
        sample_traces = random.sample(valid_traces, sample_size)
    else:
        sample_traces = valid_traces
    
    print(f"[pose_keyframe] Processing {len(sample_traces)} key frames...")
    
    # Download model
    model_path = os.path.expanduser("~/.mediapipe/models/pose_landmarker_heavy.task")
    if not os.path.exists(model_path):
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        print(f"[test] Downloading model...")
        import urllib.request
        url = "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/1/pose_landmarker_heavy.task"
        urllib.request.urlretrieve(url, model_path)
    
    # Initialize MediaPipe
    options = PoseLandmarkerOptions(
        base_options=BaseOptions(model_asset_path=model_path),
        running_mode=mp.tasks.vision.RunningMode.IMAGE,
    )
    detector = PoseLandmarker.create_from_options(options)
    
    # Test each sample
    results = []
    
    for i, trace_id in enumerate(sample_traces):
        key_frame_path = output_path / f"trace_{trace_id}" / "key_frame.jpg"
        
        # Get face bbox
        face_bbox = get_face_bbox(face_data, trace_id)
        if not face_bbox:
            results.append({
                'trace_id': trace_id,
                'status': 'no_face_bbox',
                'aligned': False
            })
            continue
        
        # Read image
        img = cv2.imread(str(key_frame_path))
        if img is None:
            results.append({
                'trace_id': trace_id,
                'status': 'cannot_read_image',
                'aligned': False
            })
            continue
        
        img_height, img_width = img.shape[:2]
        
        # Detect pose
        rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        mp_image = mp.Image(mp.ImageFormat.SRGB, rgb_img)
        results_mp = detector.detect(mp_image)
        
        if not results_mp.pose_landmarks:
            results.append({
                'trace_id': trace_id,
                'status': 'no_pose_detected',
                'aligned': False
            })
            continue
        
        # Check each pose for alignment
        aligned_poses = []
        for pose_landmarks in results_mp.pose_landmarks:
            aligned, face_kps = check_pose_alignment(
                pose_landmarks, face_bbox, img_width, img_height
            )
            if aligned:
                # Get all keypoints
                all_keypoints = []
                for idx, landmark in enumerate(pose_landmarks):
                    if idx < len(LANDMARK_NAMES):
                        name = LANDMARK_NAMES[idx]
                        all_keypoints.append({
                            'name': name,
                            'x': landmark.x * img_width,
                            'y': landmark.y * img_height,
                            'confidence': landmark.visibility if hasattr(landmark, 'visibility') else 1.0
                        })
                
                aligned_poses.append({
                    'face_keypoints': face_kps,
                    'keypoints': all_keypoints
                })
        
        if aligned_poses:
            # Get all keypoints from the best aligned pose
            best_pose = aligned_poses[0]
            all_keypoints = []
            for idx, landmark in enumerate(pose_landmarks):
                if idx < len(LANDMARK_NAMES):
                    name = LANDMARK_NAMES[idx]
                    all_keypoints.append({
                        'name': name,
                        'x': landmark.x * img_width,
                        'y': landmark.y * img_height,
                        'confidence': landmark.visibility if hasattr(landmark, 'visibility') else 1.0
                    })
            
            pose_data = {
                'keypoints': all_keypoints,
                'face_keypoints': best_pose['face_keypoints'],
                'num_poses_detected': len(results_mp.pose_landmarks),
                'num_aligned': len(aligned_poses)
            }
            
            results.append({
                'trace_id': trace_id,
                'status': 'aligned',
                'aligned': True,
                'pose_data': pose_data
            })
        else:
            results.append({
                'trace_id': trace_id,
                'status': 'pose_not_aligned',
                'aligned': False,
                'num_poses_detected': len(results_mp.pose_landmarks)
            })
        
        if (i + 1) % 10 == 0:
            print(f"[test] Processed {i+1}/{len(sample_traces)} samples...")
    
    detector.close()
    
    # Calculate statistics
    total = len(results)
    aligned_count = sum(1 for r in results if r.get('aligned'))
    
    status_counts = {}
    for r in results:
        status = r.get('status', 'unknown')
        status_counts[status] = status_counts.get(status, 0) + 1
    
    print(f"\n[pose_keyframe] === RESULTS ===")
    print(f"[pose_keyframe] Total processed: {total}")
    print(f"[pose_keyframe] Aligned: {aligned_count} ({aligned_count/total*100:.1f}%)")
    print(f"[pose_keyframe] Not aligned: {total - aligned_count} ({(total-aligned_count)/total*100:.1f}%)")
    print(f"\n[pose_keyframe] Status breakdown:")
    for status, count in sorted(status_counts.items()):
        print(f"[pose_keyframe]   {status}: {count} ({count/total*100:.1f}%)")
    
    # Save results
    output_file = Path(output_dir) / f"{file_uuid}.pose_keyframe_results.json"
    with open(output_file, 'w') as f:
        json.dump({
            'total_processed': total,
            'aligned': aligned_count,
            'alignment_rate': f"{aligned_count/total*100:.1f}%",
            'status_counts': status_counts,
            'results': results
        }, f, indent=2)
    
    print(f"\n[pose_keyframe] Results saved to: {output_file}")
    
    # Update trace_profile.json for aligned poses
    print(f"\n[pose_keyframe] Updating trace_profile.json for aligned poses...")
    updated_count = 0
    for r in results:
        if r.get('aligned') and r.get('pose_data'):
            trace_id = r['trace_id']
            profile_path = output_path / f"trace_{trace_id}" / "trace_profile.json"
            
            if profile_path.exists():
                with open(profile_path) as f:
                    profile = json.load(f)
                
                profile['pose'] = r['pose_data']
                profile['pose_aligned'] = True
                
                with open(profile_path, 'w') as f:
                    json.dump(profile, f, indent=2)
                
                updated_count += 1
    
    print(f"[pose_keyframe] Updated {updated_count} trace_profile.json files")


def main():
    parser = argparse.ArgumentParser(description="Process MediaPipe pose on key frames")
    parser.add_argument("--file-uuid", "-u", required=True, help="File UUID")
    parser.add_argument("--output-dir", "-o", default="/Users/accusys/momentry/output", help="Output directory")
    parser.add_argument("--sample-size", "-s", type=int, default=None, help="Sample size (default: all)")
    parser.add_argument("--all", "-a", action="store_true", help="Process all key frames")
    args = parser.parse_args()
    
    sample_size = None if args.all else args.sample_size
    test_alignment(args.file_uuid, args.output_dir, sample_size)


if __name__ == "__main__":
    main()