#!/usr/bin/env python3
"""
Grader for Task 1 — AR overlay projection.

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

Binary pass/fail. ALL checks must pass for EVERY frame. Prints the first failure.

Hardening: every numeric field is validated to be a finite real number (rejecting
NaN/Inf, booleans, strings) BEFORE any tolerance comparison, and nested objects must
carry exactly their declared key set. Malformed values fail deterministically rather
than slipping through comparisons.
"""
import json
import sys
import math

TOL_VEC = 1e-4      # per-component abs error for 3D vectors and plane (A,B,C,D)
TOL_UNIT = 1e-6     # unit-length tolerance for normals / directions
TOL_DIST = 1e-4     # mm, signed distance
TOL_PX = 1e-2       # px, Euclidean

def fail(msg):
    print(f"FAIL: {msg}")
    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):
    """Return x as a finite float, or fail. Rejects bool/str/None/NaN/Inf."""
    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_num_list(x, n, where):
    """Validate x is a list of exactly n finite numbers; return list[float]."""
    if not isinstance(x, list) or len(x) != n:
        fail(f"{where}: expected a list of {n} numbers, got {x!r}")
    return [as_num(x[i], f"{where}[{i}]") for i in range(n)]

def is_null(x):
    return x is None

def unit_len(vals):
    return math.sqrt(sum(v * v for v in vals))

def main():
    ref_path = sys.argv[1] if len(sys.argv) > 1 else "result.json"
    sub_path = sys.argv[2] if len(sys.argv) > 2 else "submission.json"
    ref = load(ref_path)
    sub = load(sub_path)

    if not isinstance(sub, list):
        fail("submission is not a JSON array")
    if len(sub) != len(ref):
        fail(f"frame count mismatch: ref {len(ref)}, sub {len(sub)}")

    expected_keys = {"id", "plane_tracker", "tip_tracker", "dir_tracker",
                     "signed_distance_mm", "intersection_tracker", "tip_px", "intersection_px"}

    for r, s in zip(ref, sub):
        fid = r["id"]
        if not isinstance(s, dict):
            fail(f"frame (ref id {fid}): entry is not a JSON object")
        if set(s.keys()) != expected_keys:
            fail(f"frame {fid}: key set mismatch. expected {sorted(expected_keys)}, got {sorted(s.keys())}")
        if not isinstance(s["id"], int) or isinstance(s["id"], bool) or s["id"] != r["id"]:
            fail(f"frame order/id mismatch: expected id {r['id']}, got {s['id']!r}")

        # plane_tracker: exact nested key set + finite numeric values
        pr, ps = r["plane_tracker"], s["plane_tracker"]
        if not isinstance(ps, dict) or set(ps.keys()) != {"A", "B", "C", "D"}:
            fail(f"frame {fid}: plane_tracker must be an object with exactly keys A,B,C,D; got {ps!r}")
        for k in ("A", "B", "C", "D"):
            rv = as_num(pr[k], f"frame {fid}: reference plane_tracker['{k}']")
            sv = as_num(ps[k], f"frame {fid}: plane_tracker['{k}']")
            if abs(rv - sv) > TOL_VEC:
                fail(f"frame {fid}: plane_tracker['{k}'] err {abs(rv-sv):.2e} > {TOL_VEC:.0e}")
        if abs(unit_len([as_num(ps[k], f"frame {fid}: plane_tracker['{k}']") for k in ("A", "B", "C")]) - 1.0) > TOL_UNIT:
            fail(f"frame {fid}: plane normal not unit length")

        # 3D vecs (always present, finite, length 3)
        for name, tol in (("tip_tracker", TOL_VEC), ("dir_tracker", TOL_VEC)):
            rv = as_num_list(r[name], 3, f"frame {fid}: reference {name}")
            sv = as_num_list(s[name], 3, f"frame {fid}: {name}")
            for i in range(3):
                if abs(rv[i] - sv[i]) > tol:
                    fail(f"frame {fid}: {name}[{i}] err {abs(rv[i]-sv[i]):.2e} > {tol:.0e}")
        if abs(unit_len(as_num_list(s["dir_tracker"], 3, f"frame {fid}: dir_tracker")) - 1.0) > TOL_UNIT:
            fail(f"frame {fid}: dir_tracker not unit length")

        # signed distance
        rd = as_num(r["signed_distance_mm"], f"frame {fid}: reference signed_distance_mm")
        sd = as_num(s["signed_distance_mm"], f"frame {fid}: signed_distance_mm")
        if abs(rd - sd) > TOL_DIST:
            fail(f"frame {fid}: signed_distance_mm err {abs(rd-sd):.2e} > {TOL_DIST:.0e}")

        # nullable fields: null-agreement, then finite + tolerance
        for name, tol, n, px in (("intersection_tracker", TOL_VEC, 3, False),
                                 ("tip_px", TOL_PX, 2, True),
                                 ("intersection_px", TOL_PX, 2, True)):
            if is_null(r[name]) != is_null(s[name]):
                fail(f"frame {fid}: '{name}' null-mismatch (ref null={is_null(r[name])}, sub null={is_null(s[name])})")
            if r[name] is None:
                continue
            rv = as_num_list(r[name], n, f"frame {fid}: reference {name}")
            sv = as_num_list(s[name], n, f"frame {fid}: {name}")
            if px:
                if math.hypot(rv[0] - sv[0], rv[1] - sv[1]) > tol:
                    fail(f"frame {fid}: '{name}' pixel err {math.hypot(rv[0]-sv[0], rv[1]-sv[1]):.2e} > {tol:.0e}")
            else:
                for i in range(n):
                    if abs(rv[i] - sv[i]) > tol:
                        fail(f"frame {fid}: '{name}[{i}]' err {abs(rv[i]-sv[i]):.2e} > {tol:.0e}")

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

if __name__ == "__main__":
    main()