#!/usr/bin/env python3
"""
Grader for Task 3 — Similarity point-set registration.

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

Verifies R is a valid rotation, the transform achieves the optimal RMS on the data,
the self-reported rms is honest, and (s, R, t) match the reference. Binary pass/fail.

Hardening: scale, R, t, 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 tolerances.
"""
import json
import sys
import math
import numpy as np

TOL_ORTHO = 1e-9
TOL_RMS_SLACK = 1e-4   # submission rms may exceed reference optimum by at most this
TOL_RMS_REPORT = 1e-4  # self-reported rms honesty
TOL_SCALE = 1e-5
TOL_R = 1e-4           # rad geodesic vs reference
TOL_T = 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 as_matrix(x, rows, cols, where):
    if not isinstance(x, list) or len(x) != rows:
        fail(f"{where}: expected {rows}x{cols} matrix, got {x!r}")
    out = []
    for i, row in enumerate(x):
        if not isinstance(row, list) or len(row) != cols:
            fail(f"{where}[{i}]: expected a row of {cols} numbers, got {row!r}")
        out.append([as_num(row[j], f"{where}[{i}][{j}]") for j in range(cols)])
    return np.array(out, float)

def rms_of(s, R, t, src, dst):
    pred = (s * (R @ src.T).T + t)
    return float(np.sqrt(((pred - dst) ** 2).sum(1).mean()))

def geodesic(Ra, Rb):
    c = (np.trace(Ra.T @ Rb) - 1) / 2
    return float(np.arccos(max(-1.0, min(1.0, c))))

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 ("scale", "R", "t", "rms_error"):
        if k not in sub:
            fail(f"submission missing key '{k}'")

    s = as_num(sub["scale"], "scale")
    R = as_matrix(sub["R"], 3, 3, "R")
    t = as_vector(sub["t"], 3, "t")
    reported = as_num(sub["rms_error"], "rms_error")

    if s <= 0:
        fail(f"scale must be > 0, got {s}")
    if np.linalg.norm(R.T @ R - np.eye(3)) > TOL_ORTHO:
        fail("R is not orthonormal")
    if abs(np.linalg.det(R) - 1.0) > TOL_ORTHO:
        fail(f"det(R) != 1 (got {np.linalg.det(R):.6f})")

    src = np.array(scene["src"], float); dst = np.array(scene["dst"], float)
    rms = rms_of(s, R, t, src, dst)
    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 abs(s - float(ref["scale"])) > TOL_SCALE:
        fail(f"scale differs from reference by {abs(s-float(ref['scale'])):.2e}")
    if geodesic(R, np.array(ref["R"], float)) > TOL_R:
        fail(f"rotation differs from reference by {np.degrees(geodesic(R, np.array(ref['R'],float))):.2e} deg")
    if np.linalg.norm(t - np.array(ref["t"], float)) > TOL_T:
        fail(f"translation differs from reference by {np.linalg.norm(t-np.array(ref['t'],float)):.2e} mm")

    print(f"PASS: valid similarity transform, rms {rms:.6f}, matches reference")

if __name__ == "__main__":
    main()