#!/usr/bin/env python3
"""
Grader for Task 2 — Hand-eye calibration.

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

Method-agnostic: verifies the submitted X is a valid rigid transform and that it
satisfies A_i X = X B_i on the actual data (which pins X uniquely). Also confirms it
matches the reference X. Binary pass/fail.

Hardening: X and max_residual 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 tolerance comparisons.
"""
import json
import sys
import math
import numpy as np

TOL_RESIDUAL = 1e-6   # Frobenius, ||A X - X B||
TOL_ORTHO = 1e-9      # rotation orthonormality / det
TOL_MATCH_R = 1e-4    # rad, geodesic angle vs reference X
TOL_MATCH_T = 1e-4    # mm, translation vs reference X
TOL_REPORTED = 1e-5   # tolerance on the residual value the submission self-reports

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_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 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) or "X" not in sub or "max_residual" not in sub:
        fail("submission must be an object with keys 'X' and 'max_residual'")
    X = as_matrix(sub["X"], 4, 4, "X")
    reported = as_num(sub["max_residual"], "max_residual")

    # valid rigid transform?
    R = X[:3, :3]
    if np.linalg.norm(R.T @ R - np.eye(3)) > TOL_ORTHO:
        fail("X rotation block is not orthonormal")
    if abs(np.linalg.det(R) - 1.0) > TOL_ORTHO:
        fail(f"X rotation det != 1 (got {np.linalg.det(R):.6f})")
    if np.linalg.norm(X[3, :] - np.array([0, 0, 0, 1.0])) > TOL_ORTHO:
        fail("X bottom row is not [0,0,0,1]")

    # satisfies constraint on the data?
    residual = 0.0
    for pr in scene["pairs"]:
        A = np.array(pr["A"], float); B = np.array(pr["B"], float)
        residual = max(residual, float(np.linalg.norm(A @ X - X @ B)))
    if residual > TOL_RESIDUAL:
        fail(f"max ||A X - X B|| = {residual:.2e} > {TOL_RESIDUAL:.0e}")

    # self-reported residual is honest?
    if abs(reported - residual) > TOL_REPORTED:
        fail(f"reported max_residual {reported:.2e} disagrees with actual {residual:.2e}")

    # matches reference X?
    Xr = np.array(ref["X"], float)
    if geodesic(R, Xr[:3, :3]) > TOL_MATCH_R:
        fail(f"rotation differs from reference by {np.degrees(geodesic(R, Xr[:3,:3])):.2e} deg")
    if np.linalg.norm(X[:3, 3] - Xr[:3, 3]) > TOL_MATCH_T:
        fail(f"translation differs from reference by {np.linalg.norm(X[:3,3]-Xr[:3,3]):.2e} mm")

    print(f"PASS: X valid rigid transform, max residual {residual:.2e}, matches reference")

if __name__ == "__main__":
    main()