#!/opt/homebrew/bin/python3.11
"""
BVH Exporter v1.11 — Export face traces as 3D BVH motion files

Input: face_traced.json + mediapipe.json
Output: {file_uuid}_trace_{trace_id}_v1.11.bvh per face_trace

Skeleton:
  Head (root, positioned at bbox center)
  └─ Neck
     ├─ LeftShoulder → LeftElbow → LeftHand
     └─ RightShoulder → RightElbow → RightHand
  Spine → Hips

Joint channels: position (Xposition Yposition Zposition) + rotation (Zrotation Xrotation Yrotation)
Uses 2D→3D depth heuristic and head pose yaw/pitch/roll for rotation.
"""

import json, os, sys, argparse, math
from typing import Optional


# BVH skeleton definition
BVH_HEADER = """HIERARCHY
ROOT Head
{{
\tOFFSET 0.0 0.0 0.0
\tCHANNELS 6 Xposition Yposition Zposition Zrotation Xrotation Yrotation
\tJOINT Neck
\t{{
\t\tOFFSET 0.0 -0.15 0.0
\t\tCHANNELS 3 Zrotation Xrotation Yrotation
\t\tJOINT Spine
\t\t{{
\t\t\tOFFSET 0.0 -0.3 0.0
\t\t\tCHANNELS 3 Zrotation Xrotation Yrotation
\t\t\tJOINT Hips
\t\t\t{{
\t\t\t\tOFFSET 0.0 -0.3 0.0
\t\t\t\tCHANNELS 3 Zrotation Xrotation Yrotation
\t\t\t\tEnd Site
\t\t\t\t{{
\t\t\t\t\tOFFSET 0.0 -0.1 0.0
\t\t\t\t}}
\t\t\t}}
\t\t}}
\t\tJOINT LeftShoulder
\t\t{{
\t\t\tOFFSET -0.15 -0.05 0.0
\t\t\tCHANNELS 3 Zrotation Xrotation Yrotation
\t\t\tJOINT LeftElbow
\t\t\t{{
\t\t\t\tOFFSET -0.2 0.0 0.0
\t\t\t\tCHANNELS 3 Zrotation Xrotation Yrotation
\t\t\t\tJOINT LeftHand
\t\t\t\t{{
\t\t\t\t\tOFFSET -0.15 0.0 0.0
\t\t\t\t\tCHANNELS 3 Zrotation Xrotation Yrotation
\t\t\t\t\tEnd Site
\t\t\t\t\t{{
\t\t\t\t\t\tOFFSET -0.05 0.0 0.0
\t\t\t\t\t}}
\t\t\t\t}}
\t\t\t}}
\t\t}}
\t\tJOINT RightShoulder
\t\t{{
\t\t\tOFFSET 0.15 -0.05 0.0
\t\t\tCHANNELS 3 Zrotation Xrotation Yrotation
\t\t\tJOINT RightElbow
\t\t\t{{
\t\t\t\tOFFSET 0.2 0.0 0.0
\t\t\t\tCHANNELS 3 Zrotation Xrotation Yrotation
\t\t\t\tJOINT RightHand
\t\t\t\t{{
\t\t\t\t\tOFFSET 0.15 0.0 0.0
\t\t\t\t\tCHANNELS 3 Zrotation Xrotation Yrotation
\t\t\t\t\tEnd Site
\t\t\t\t\t{{
\t\t\t\t\t\tOFFSET 0.05 0.0 0.0
\t\t\t\t\t}}
\t\t\t\t}}
\t\t\t}}
\t\t}}
\t}}
}}
"""


def depth_heuristic(x: float, y: float, w: float, h: float,
                    img_w: float, img_h: float, frame_w: float = 1.0) -> float:
    """Estimate depth (z) from bbox size: larger = closer"""
    size_ratio = (w * h) / (img_w * img_h)
    return max(-2.0, min(2.0, 2.0 - size_ratio * 100))


def extract_trace_frames(face_data: dict, trace_id: int) -> list:
    """Extract frames for a specific trace from face_traced.json"""
    frames = face_data.get("frames", {})
    trace_frames = []
    for fnum_str, frm in sorted(frames.items(), key=lambda x: int(x[0])):
        fnum = int(fnum_str)
        for face in frm.get("faces", []):
            if face.get("trace_id") == trace_id:
                trace_frames.append({
                    "frame": fnum,
                    "timestamp": frm.get("time_seconds", fnum / 30.0),
                    "x": face.get("x", 0),
                    "y": face.get("y", 0),
                    "width": face.get("width", 50),
                    "height": face.get("height", 50),
                    "yaw": face.get("pose_angle", {}).get("yaw", 0),
                    "pitch": face.get("pose_angle", {}).get("pitch", 0),
                    "roll": face.get("pose_angle", {}).get("roll", 0),
                })
                break
    return trace_frames


def generate_motion(trace_frames: list, fps: float,
                    img_w: float = 1920, img_h: float = 1080) -> str:
    """Generate BVH motion data from trace frames"""
    if not trace_frames:
        return ""

    lines = []
    for f in trace_frames:
        # Normalize position to [-1, 1] range
        px = (f["x"] / img_w) * 2 - 1
        py = (f["y"] / img_h) * 2 - 1
        pz = depth_heuristic(f["x"], f["y"], f["width"], f["height"], img_w, img_h)

        yaw = f.get("yaw", 0)
        pitch = f.get("pitch", 0)
        roll = f.get("roll", 0)

        lines.append(f"{px:.4f} {py:.4f} {pz:.4f} {roll:.1f} {pitch:.1f} {yaw:.1f}  "
                      f"0 0 0  "  # Neck (no IK yet)
                      f"0 0 0  "  # Spine
                      f"0 0 0  "  # Hips
                      f"0 0 0  "  # LeftShoulder
                      f"0 0 0  "  # LeftElbow
                      f"0 0 0  "  # LeftHand
                      f"0 0 0  "  # RightShoulder
                      f"0 0 0  "  # RightElbow
                      f"0 0 0")   # RightHand

    n_frames = len(trace_frames)
    frame_time = 1.0 / fps

    motion = (
        f"MOTION\n"
        f"Frames: {n_frames}\n"
        f"Frame Time: {frame_time:.6f}\n"
    )
    motion += "\n".join(lines)
    return motion


def main():
    parser = argparse.ArgumentParser(description="BVH Exporter v1.11")
    parser.add_argument("--file-uuid", required=True)
    parser.add_argument("--trace-id", type=int, default=None,
                        help="Specific trace to export (default: all)")
    parser.add_argument("--face-json", help="Path to face_traced.json")
    parser.add_argument("--output-dir", 
                        default=os.environ.get("MOMENTRY_OUTPUT_DIR", 
                                               "/Users/accusys/momentry/output_dev"))
    args = parser.parse_args()

    OUTPUT_DIR = args.output_dir

    face_json_path = args.face_json or os.path.join(
        OUTPUT_DIR, f"{args.file_uuid}.face_traced.json"
    )
    if not os.path.exists(face_json_path):
        face_json_path = os.path.join(OUTPUT_DIR, f"{args.file_uuid}.face.json")

    if not os.path.exists(face_json_path):
        print(f"[BVH] ❌ face JSON not found: {face_json_path}")
        sys.exit(1)

    with open(face_json_path) as f:
        face_data = json.load(f)

    metadata = face_data.get("metadata", {})
    fps = metadata.get("fps", 30.0)
    width = metadata.get("width", 1920)
    height = metadata.get("height", 1080)

    traces = face_data.get("traces", {})
    if not traces:
        print(f"[BVH] ❌ No traces found")
        sys.exit(1)

    trace_ids = [args.trace_id] if args.trace_id is not None else sorted(
        [int(k) for k in traces.keys()]
    )

    for tid in trace_ids:
        trace_frames = extract_trace_frames(face_data, tid)
        if len(trace_frames) < 5:
            print(f"[BVH] Skip trace {tid}: only {len(trace_frames)} frames")
            continue

        motion = generate_motion(trace_frames, fps, width, height)
        if not motion:
            continue

        bvh_content = BVH_HEADER + "\n" + motion
        out_path = os.path.join(OUTPUT_DIR,
                                 f"{args.file_uuid}_trace_{tid}_v1.11.bvh")
        with open(out_path, "w") as f:
            f.write(bvh_content)

        print(f"[BVH] ✅ Trace {tid}: {len(trace_frames)} frames → {out_path}")

    print(f"[BVH] Done: {len(trace_ids)} traces exported")


if __name__ == "__main__":
    main()
