#!/opt/homebrew/bin/python3.11
"""
Extract embeddings from existing face.json using CoreML FaceNet.

Usage: python3 scripts/extract_embeddings_from_face_json.py <face_json_path> <video_path> <output_path>
"""
import sys
import os
import json
import cv2
import numpy as np
import coremltools as ct

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
FACENET_PATH = os.path.join(SCRIPT_DIR, "..", "models", "facenet512.mlpackage")


def extract_embeddings(face_json_path: str, video_path: str, output_path: str):
    # Load face.json
    with open(face_json_path, 'r') as f:
        face_data = json.load(f)
    
    frames = face_data.get('frames', [])
    if not frames:
        print("No frames in face.json")
        return
    
    # Load CoreML FaceNet
    facenet = os.path.normpath(FACENET_PATH)
    if not os.path.exists(facenet):
        print(f"FaceNet model not found: {facenet}")
        return
    
    coreml_model = ct.models.MLModel(facenet)
    print(f"[EMB] CoreML FaceNet loaded: {facenet}")
    
    # Open video
    video = cv2.VideoCapture(video_path)
    if not video.isOpened():
        print(f"Cannot open video: {video_path}")
        return
    
    fps = video.get(cv2.CAP_PROP_FPS)
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f"[EMB] Video: {fps} fps, {total_frames} frames")
    
    # Extract embeddings for each face
    embed_count = 0
    processed_frames = 0
    
    for frame_entry in frames:
        frame_num = frame_entry.get('frame', 0)
        faces = frame_entry.get('faces', [])
        
        # Seek to frame
        video.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
        ret, img = video.read()
        if not ret:
            continue
        
        processed_frames += 1
        
        for face in faces:
            x = face.get('x', 0)
            y = face.get('y', 0)
            w = face.get('width', 0)
            h = face.get('height', 0)
            
            if w == 0 or h == 0:
                continue
            
            # Crop face
            crop = img[y:y+h, x:x+w]
            if crop.size == 0:
                continue
            
            # Resize to 160x160 (FaceNet input size)
            crop_resized = cv2.resize(crop, (160, 160))
            
            # Convert to RGB
            crop_rgb = cv2.cvtColor(crop_resized, cv2.COLOR_BGR2RGB)
            
            # CoreML embedding
            try:
                input_dict = {'image': crop_rgb}
                output = coreml_model.predict(input_dict)
                emb = output.get('output', output.get('embeddings', None))
                
                if emb is not None:
                    if isinstance(emb, np.ndarray):
                        emb = emb.flatten().tolist()
                    elif isinstance(emb, dict):
                        emb = list(emb.values())[0]
                        if isinstance(emb, np.ndarray):
                            emb = emb.flatten().tolist()
                    
                    face['embedding'] = emb
                    embed_count += 1
            except Exception as e:
                print(f"[EMB] Frame {frame_num} embedding failed: {e}")
        
        if processed_frames % 1000 == 0:
            print(f"[EMB] Processed {processed_frames} frames, {embed_count} embeddings")
    
    video.release()
    
    # Save updated face.json
    face_data['metadata']['total_embeddings'] = embed_count
    
    with open(output_path, 'w') as f:
        json.dump(face_data, f)
    
    print(f"[EMB] Done: {processed_frames} frames, {embed_count} embeddings")


if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: python3 extract_embeddings_from_face_json.py <face_json> <video> <output>")
        sys.exit(1)
    
    extract_embeddings(sys.argv[1], sys.argv[2], sys.argv[3])