#!/usr/bin/env python3
"""
Grader for Task 5 — SLERP pose interpolation.

Usage:
    python grader.py [reference.json] [submission.json]
Defaults: result.json, submission.json

Quaternions compared sign-invariantly (geodesic angle). Binary pass/fail.

Hardening: every numeric field is parsed as a finite real number of the exact declared
shape (rejecting NaN/Inf/booleans/strings/ragged arrays) and each entry must carry
exactly its declared key set BEFORE any tolerance comparison, so malformed values fail
deterministically instead of bypassing tolerances.
"""
import json
import sys
import math
import numpy as np

TOL_QUAT_DEG = 1e-4    # geodesic angle between submitted and reference rotation
TOL_UNIT = 1e-6        # quaternion unit-norm
TOL_P = 1e-4           # mm, per component translation
TOL_ANGLE = 1e-3       # deg, angle_from_prev
TOL_T = 1e-9           # query timestamp match

def fail(m):
    print(f"FAIL: {m}"); sys.exit(1)

def _reject_constant(tok):
    # json.load calls this for the non-standard tokens NaN, Infinity, -Infinity.
    raise ValueError(f"non-finite JSON constant '{tok}' is not allowed")

def load(p):
    with open(p) as f:
        try:
            return json.load(f, parse_constant=_reject_constant)
        except ValueError as e:
            fail(f"{p}: invalid JSON ({e})")

def as_num(x, where):
    if isinstance(x, bool) or not isinstance(x, (int, float)):
        fail(f"{where}: expected a real number, got {type(x).__name__} ({x!r})")
    xf = float(x)
    if not math.isfinite(xf):
        fail(f"{where}: value is non-finite ({x!r})")
    return xf

def as_vector(x, n, where):
    if not isinstance(x, list) or len(x) != n:
        fail(f"{where}: expected a list of {n} numbers, got {x!r}")
    return np.array([as_num(x[i], f"{where}[{i}]") for i in range(n)], float)

def geodesic_deg(qa, qb):
    d = abs(float(np.dot(qa / np.linalg.norm(qa), qb / np.linalg.norm(qb))))
    return float(np.degrees(2.0 * np.arccos(min(1.0, d))))

def main():
    ref = load(sys.argv[1] if len(sys.argv) > 1 else "result.json")["interpolated"]
    sub_root = load(sys.argv[2] if len(sys.argv) > 2 else "submission.json")
    if not isinstance(sub_root, dict) or "interpolated" not in sub_root:
        fail("submission must be an object with top-level key 'interpolated'")
    sub = sub_root["interpolated"]

    if not isinstance(sub, list) or len(sub) != len(ref):
        fail(f"query count mismatch: ref {len(ref)}, sub {len(sub) if isinstance(sub,list) else 'not-a-list'}")

    keys = {"t", "q", "p", "angle_from_prev_deg"}
    for i, (r, s) in enumerate(zip(ref, sub)):
        if not isinstance(s, dict) or set(s.keys()) != keys:
            fail(f"query {i}: each entry must be an object with exactly keys {sorted(keys)}; got {s!r}")

        if abs(as_num(s["t"], f"query {i}: t") - float(r["t"])) > TOL_T:
            fail(f"query {i}: timestamp mismatch (expected {r['t']}, got {s['t']!r})")

        q = as_vector(s["q"], 4, f"query {i}: q")
        if abs(np.linalg.norm(q) - 1.0) > TOL_UNIT:
            fail(f"query {i}: q not unit norm (|q|={np.linalg.norm(q):.6f})")
        if q[0] < -TOL_UNIT:
            fail(f"query {i}: quaternion not in canonical form (w={q[0]:.6f} < 0); report w >= 0")
        ang = geodesic_deg(q, np.array(r["q"], float))
        if ang > TOL_QUAT_DEG:
            fail(f"query {i}: rotation differs from reference by {ang:.2e} deg")

        p = as_vector(s["p"], 3, f"query {i}: p")
        if np.max(np.abs(p - np.array(r["p"], float))) > TOL_P:
            fail(f"query {i}: translation err > {TOL_P:.0e}")

        rp, sp = r["angle_from_prev_deg"], s["angle_from_prev_deg"]
        if (rp is None) != (sp is None):
            fail(f"query {i}: angle_from_prev_deg null-mismatch")
        if rp is not None:
            spv = as_num(sp, f"query {i}: angle_from_prev_deg")
            if abs(float(rp) - spv) > TOL_ANGLE:
                fail(f"query {i}: angle_from_prev_deg err {abs(float(rp)-spv):.2e} > {TOL_ANGLE:.0e}")

    print(f"PASS: {len(ref)} queries, all fields finite and within tolerance")

if __name__ == "__main__":
    main()