#!/opt/homebrew/bin/python3.11
"""
Compare Apple Vision pose vs MediaPipe pose

Finds:
- Intersection: Poses detected by both
- Apple Vision only: Poses only in Apple Vision
- MediaPipe only: Poses only in MediaPipe

Usage:
    python3 scripts/compare_pose_detections.py --file-uuid <uuid>
"""

import argparse
import json
from pathlib import Path


def load_apple_vision_poses(file_uuid, output_dir):
    """Load Apple Vision pose data from pose.json"""
    pose_path = Path(output_dir) / f"{file_uuid}.pose.json"
    
    if not pose_path.exists():
        return {}
    
    with open(pose_path) as f:
        data = json.load(f)
    
    poses = {}
    for frame in data.get('frames', []):
        frame_num = frame.get('frame', frame.get('frame_number', 0))
        for i, person in enumerate(frame.get('persons', [])):
            pose_key = f"frame_{frame_num}_person_{i}"
            poses[pose_key] = {
                'frame': frame_num,
                'person_idx': i,
                'keypoints': person.get('keypoints', []),
                'source': 'apple_vision'
            }
    
    return poses


def load_mediapipe_poses(file_uuid, output_dir):
    """Load MediaPipe pose data from pose.mediapipe.json"""
    pose_path = Path(output_dir) / f"{file_uuid}.pose.mediapipe.json"
    
    if not pose_path.exists():
        return {}
    
    with open(pose_path) as f:
        data = json.load(f)
    
    poses = {}
    for frame in data.get('frames', []):
        frame_num = frame.get('frame', 0)
        for i, person in enumerate(frame.get('persons', [])):
            pose_key = f"frame_{frame_num}_person_{i}"
            poses[pose_key] = {
                'frame': frame_num,
                'person_idx': i,
                'keypoints': person.get('keypoints', []),
                'source': 'mediapipe'
            }
    
    return poses


def compare_poses(file_uuid, output_dir):
    """Compare Apple Vision vs MediaPipe poses."""
    
    print(f"[compare] Loading pose data for {file_uuid}...")
    
    av_poses = load_apple_vision_poses(file_uuid, output_dir)
    mp_poses = load_mediapipe_poses(file_uuid, output_dir)
    
    print(f"[compare] Apple Vision poses: {len(av_poses)}")
    print(f"[compare] MediaPipe poses: {len(mp_poses)}")
    
    # Find intersection and differences
    av_keys = set(av_poses.keys())
    mp_keys = set(mp_poses.keys())
    
    intersection = av_keys & mp_keys
    av_only = av_keys - mp_keys
    mp_only = mp_keys - av_keys
    
    print(f"\n[compare] === COMPARISON ===")
    print(f"[compare] Intersection (both detected): {len(intersection)}")
    print(f"[compare] Apple Vision only: {len(av_only)}")
    print(f"[compare] MediaPipe only: {len(mp_only)}")
    
    # Analyze intersection - check alignment
    intersection_aligned = 0
    for key in intersection:
        av_pose = av_poses[key]
        mp_pose = mp_poses[key]
        
        # Check if both have face keypoints
        av_face_kps = [kp for kp in av_pose.get('keypoints', []) if kp.get('name') in ['nose', 'left_eye', 'right_eye']]
        mp_face_kps = [kp for kp in mp_pose.get('keypoints', []) if kp.get('name') in ['nose', 'left_eye', 'right_eye']]
        
        if av_face_kps and mp_face_kps:
            intersection_aligned += 1
    
    print(f"\n[compare] Intersection with face keypoints: {intersection_aligned}")
    
    # Frame coverage
    av_frames = set(av_poses[k]['frame'] for k in av_keys)
    mp_frames = set(mp_poses[k]['frame'] for k in mp_keys)
    
    print(f"\n[compare] === FRAME COVERAGE ===")
    print(f"[compare] Apple Vision frames: {len(av_frames)}")
    print(f"[compare] MediaPipe frames: {len(mp_frames)}")
    print(f"[compare] Overlapping frames: {len(av_frames & mp_frames)}")
    
    # Save results
    output_path = Path(output_dir) / f"{file_uuid}.pose_comparison.json"
    with open(output_path, 'w') as f:
        json.dump({
            'apple_vision_count': len(av_poses),
            'mediapipe_count': len(mp_poses),
            'intersection_count': len(intersection),
            'apple_vision_only_count': len(av_only),
            'mediapipe_only_count': len(mp_only),
            'intersection_aligned_count': intersection_aligned,
            'apple_vision_frames': len(av_frames),
            'mediapipe_frames': len(mp_frames),
            'overlapping_frames': len(av_frames & mp_frames),
            'intersection_keys': sorted(list(intersection))[:100],  # Sample
            'apple_vision_only_keys': sorted(list(av_only))[:100],
            'mediapipe_only_keys': sorted(list(mp_only))[:100],
        }, f, indent=2)
    
    print(f"\n[compare] Results saved to: {output_path}")


def main():
    parser = argparse.ArgumentParser(description="Compare pose detections")
    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")
    args = parser.parse_args()
    
    compare_poses(args.file_uuid, args.output_dir)


if __name__ == "__main__":
    main()