#!/usr/bin/env python3
"""Compare two pose sequences (extract_pose.py JSON) and score their similarity.

Method: each detected frame is reduced to a joint-angle signature (8 angles:
elbows, shoulders, hips, knees) computed from 3D world landmarks — scale- and
translation-invariant, so a phone video and a game capture compare fairly.
The two angle sequences are aligned with dynamic time warping (handles tempo
drift / offset), then scored:

  mean_angle_error_deg   average per-joint angular error along the DTW path
  per_joint_error_deg    which joints diverge most (arms vs legs etc.)
  score_0_100            100 * max(0, 1 - mean_error/90)  (rough but comparable)

Typical use: score a recreated game dance against the source video, per move
segment (--a-range/--b-range trim by time before aligning).
"""
import argparse, json, sys

import numpy as np

# (name, a, b, c) -> angle at b between vectors (a-b) and (c-b), landmark indices
ANGLE_DEFS = [
    ("l_elbow", 11, 13, 15), ("r_elbow", 12, 14, 16),
    ("l_shoulder", 13, 11, 23), ("r_shoulder", 14, 12, 24),
    ("l_hip", 11, 23, 25), ("r_hip", 12, 24, 26),
    ("l_knee", 23, 25, 27), ("r_knee", 24, 26, 28),
]


def load_angles(path, t_min, t_max):
    with open(path) as f:
        doc = json.load(f)
    ts, sigs = [], []
    for fr in doc["frames"]:
        if not fr.get("detected"):
            continue
        t = fr["t"]
        if t < t_min or (t_max is not None and t > t_max):
            continue
        lms = fr.get("world_landmarks") or fr.get("landmarks")
        pts = np.array([[p["x"], p["y"], p["z"]] for p in lms])
        sig = []
        for _, a, b, c in ANGLE_DEFS:
            v1, v2 = pts[a] - pts[b], pts[c] - pts[b]
            n1, n2 = np.linalg.norm(v1), np.linalg.norm(v2)
            if n1 < 1e-8 or n2 < 1e-8:
                sig.append(0.0)
                continue
            cosang = np.clip(np.dot(v1, v2) / (n1 * n2), -1.0, 1.0)
            sig.append(float(np.degrees(np.arccos(cosang))))
        ts.append(t)
        sigs.append(sig)
    if not sigs:
        sys.exit(f"ERROR: no detected frames in range in {path}")
    return np.array(ts), np.array(sigs)


def dtw_path(A, B):
    """O(n*m) DTW on mean-abs-angle-diff cost. Returns (path, cost_matrix)."""
    n, m = len(A), len(B)
    cost = np.mean(np.abs(A[:, None, :] - B[None, :, :]), axis=2)  # n x m, degrees
    acc = np.full((n + 1, m + 1), np.inf)
    acc[0, 0] = 0.0
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            acc[i, j] = cost[i - 1, j - 1] + min(acc[i - 1, j], acc[i, j - 1], acc[i - 1, j - 1])
    path = []
    i, j = n, m
    while i > 0 and j > 0:
        path.append((i - 1, j - 1))
        step = np.argmin([acc[i - 1, j - 1], acc[i - 1, j], acc[i, j - 1]])
        if step == 0: i, j = i - 1, j - 1
        elif step == 1: i -= 1
        else: j -= 1
    path.reverse()
    return path, cost


def parse_range(s):
    if not s:
        return 0.0, None
    lo, _, hi = s.partition(":")
    return float(lo or 0.0), (float(hi) if hi else None)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("pose_a", help="pose JSON (e.g. source video)")
    ap.add_argument("pose_b", help="pose JSON (e.g. recreation)")
    ap.add_argument("--a-range", default=None, help="time window in A, 'start:end' seconds")
    ap.add_argument("--b-range", default=None, help="time window in B, 'start:end' seconds")
    ap.add_argument("--out", default=None, help="optional JSON report path")
    args = ap.parse_args()

    a_lo, a_hi = parse_range(args.a_range)
    b_lo, b_hi = parse_range(args.b_range)
    ts_a, A = load_angles(args.pose_a, a_lo, a_hi)
    ts_b, B = load_angles(args.pose_b, b_lo, b_hi)

    path, cost = dtw_path(A, B)
    path_errs = np.array([np.abs(A[i] - B[j]) for i, j in path])  # steps x joints
    mean_err = float(path_errs.mean())
    per_joint = {ANGLE_DEFS[k][0]: round(float(path_errs[:, k].mean()), 1)
                 for k in range(len(ANGLE_DEFS))}
    score = round(100.0 * max(0.0, 1.0 - mean_err / 90.0), 1)

    report = {
        "a": {"file": args.pose_a, "frames": len(A), "span_s": [float(ts_a[0]), float(ts_a[-1])]},
        "b": {"file": args.pose_b, "frames": len(B), "span_s": [float(ts_b[0]), float(ts_b[-1])]},
        "mean_angle_error_deg": round(mean_err, 2),
        "per_joint_error_deg": dict(sorted(per_joint.items(), key=lambda kv: -kv[1])),
        "score_0_100": score,
        "dtw_path_len": len(path),
    }
    print(json.dumps(report, indent=2))
    if args.out:
        with open(args.out, "w") as f:
            json.dump(report, f, indent=2)


if __name__ == "__main__":
    main()