#!/opt/homebrew/bin/python3.11
"""Embed faces from existing detections JSON using CoreML FaceNet."""
import json, os, sys, time
import cv2
import numpy as np
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import coremltools as ct

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

def classify_pose(roll: float, yaw: float) -> str:
    abs_yaw, abs_roll = abs(yaw), abs(roll)
    if abs_yaw < 15 and abs_roll < 15:
        return "frontal"
    elif abs_yaw > 30:
        return "profile_right" if yaw > 0 else "profile_left"
    return "three_quarter"

def extract_embedding(coreml_model, face_img):
    resized = cv2.resize(face_img, (160, 160))
    normalized = (resized.astype(np.float32) / 127.5) - 1.0
    normalized = np.transpose(normalized, (2, 0, 1))
    input_array = np.expand_dims(normalized, axis=0)
    result = coreml_model.predict({"input": input_array})
    emb_key = [k for k in result.keys() if k.startswith("var_")][0]
    return result[emb_key].flatten().tolist()

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Embed faces only")
    parser.add_argument("detections_json")
    parser.add_argument("output_json")
    parser.add_argument("--video", required=True)
    args = parser.parse_args()

    print(f"[EMBED] Loading detections: {args.detections_json}")
    with open(args.detections_json) as f:
        detection_data = json.load(f)

    print(f"[EMBED] Loading CoreML FaceNet: {FACENET_PATH}")
    coreml_model = ct.models.MLModel(FACENET_PATH)

    print(f"[EMBED] Opening video: {args.video}")
    video = cv2.VideoCapture(args.video)
    fps = video.get(cv2.CAP_PROP_FPS)
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))

    face_data = {
        "metadata": {
            "video_path": os.path.abspath(args.video),
            "fps": fps, "width": width, "height": height,
            "sample_interval": detection_data.get("sample_interval", 3),
            "detection_method": "apple_vision",
            "embedding_method": "coreml_facenet",
            "total_frames": total_frames,
        },
        "frames": {}
    }

    frames = detection_data.get("frames", [])
    t0 = time.time()
    embed_count, total_face_count = 0, 0
    batch_size = max(1, len(frames) // 20)

    for idx, frame_info in enumerate(frames):
        frame_num = frame_info["frame"]
        faces = []
        for face in frame_info.get("faces", []):
            total_face_count += 1
            bb = face.get("bbox", face)
            x, y, w, h = bb["x"], bb["y"], bb["width"], bb["height"]
            if w <= 10 or h <= 10:
                continue

            video.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
            ret, frame = video.read()
            if not ret:
                continue

            x1, y1 = max(0, x), max(0, y)
            x2, y2 = min(width, x + w), min(height, y + h)
            if x2 <= x1 or y2 <= y1:
                continue
            face_img = frame[y1:y2, x1:x2]
            if face_img.size == 0:
                continue

            emb = extract_embedding(coreml_model, face_img)
            if emb is not None:
                embed_count += 1

            pose_info = face.get("pose", {})
            pose_angle = classify_pose(
                pose_info.get("roll", 0),
                pose_info.get("yaw", 0)
            )

            faces.append({
                "x": x, "y": y, "width": w, "height": h,
                "confidence": face.get("confidence", 0.5),
                "embedding": emb,
                "pose_angle": {
                    "angle": pose_angle,
                    "roll": pose_info.get("roll", 0),
                    "yaw": pose_info.get("yaw", 0),
                    "pitch": pose_info.get("pitch", 0),
                },
                "landmarks": face.get("landmarks", []),
            })

        face_data["frames"][str(frame_num)] = faces

        if (idx + 1) % batch_size == 0:
            pct = (idx + 1) / len(frames) * 100
            elapsed = time.time() - t0
            eta = (elapsed / (idx + 1)) * (len(frames) - idx - 1) if idx > 0 else 0
            print(f"[EMBED] {pct:.0f}% | {idx+1}/{len(frames)} frames | "
                  f"{embed_count} embeddings | {elapsed:.0f}s elapsed | "
                  f"{eta:.0f}s ETA", flush=True)

    video.release()
    face_data["metadata"]["status"] = "completed"

    print(f"[EMBED] Writing output: {args.output_json}")
    with open(args.output_json, "w") as f:
        json.dump(face_data, f, indent=2)

    elapsed = time.time() - t0
    print(f"[EMBED] Done: {len(frames)} frames, {embed_count}/{total_face_count} embeddings, {elapsed:.0f}s")

if __name__ == "__main__":
    main()