#!/usr/bin/env python3
"""
Grader for Task 4 — Pivot calibration.

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

Recomputes the least-squares residual for the submitted (p_tip, p_pivot), requires it
to reach the optimum, checks the self-reported rms, and matches the reference values.
Binary pass/fail.

Hardening: p_tip_tool, p_pivot_world, rms_error are parsed as finite real numbers of
the exact declared shape (rejecting NaN/Inf/booleans/strings/ragged arrays) BEFORE any
numeric check, so malformed values fail deterministically instead of bypassing tols.
"""
import json
import sys
import math
import numpy as np

TOL_RMS_SLACK = 1e-4
TOL_RMS_REPORT = 1e-4
TOL_MATCH = 1e-3   # mm, vs reference

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 rms_of(p_tip, p_piv, poses):
    n = len(poses)
    ss = 0.0
    for T in poses:
        T = np.array(T, float)
        r = T[:3, :3] @ p_tip + T[:3, 3] - p_piv
        ss += float(r @ r)
    return np.sqrt(ss / n)

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

    if not isinstance(sub, dict):
        fail("submission must be a JSON object")
    for k in ("p_tip_tool", "p_pivot_world", "rms_error"):
        if k not in sub:
            fail(f"submission missing key '{k}'")
    p_tip = as_vector(sub["p_tip_tool"], 3, "p_tip_tool")
    p_piv = as_vector(sub["p_pivot_world"], 3, "p_pivot_world")
    reported = as_num(sub["rms_error"], "rms_error")

    poses = scene["poses"]
    rms = rms_of(p_tip, p_piv, poses)
    ref_rms = float(ref["rms_error"])
    if rms > ref_rms + TOL_RMS_SLACK:
        fail(f"achieved rms {rms:.6f} exceeds optimum {ref_rms:.6f} by > {TOL_RMS_SLACK:.0e}")
    if abs(reported - rms) > TOL_RMS_REPORT:
        fail(f"reported rms {reported:.6f} disagrees with actual {rms:.6f}")

    if np.linalg.norm(p_tip - np.array(ref["p_tip_tool"], float)) > TOL_MATCH:
        fail(f"p_tip differs from reference by {np.linalg.norm(p_tip-np.array(ref['p_tip_tool'],float)):.2e} mm")
    if np.linalg.norm(p_piv - np.array(ref["p_pivot_world"], float)) > TOL_MATCH:
        fail(f"p_pivot differs from reference by {np.linalg.norm(p_piv-np.array(ref['p_pivot_world'],float)):.2e} mm")

    print(f"PASS: rms {rms:.6f}, tip/pivot match reference within {TOL_MATCH} mm")

if __name__ == "__main__":
    main()