import json
import math
from pathlib import Path
import numpy as np
from scipy.optimize import least_squares


def rot(a):
    c, s = np.cos(a), np.sin(a)
    return np.array([[c, -s], [s, c]])


def wrap(a):
    return (a + np.pi) % (2 * np.pi) - np.pi


def two_view_init(q, ell_v, ell_t, eps=1e-9):
    """Constructive gauge initialization; pair is the largest global displacement."""
    q, ell_v, ell_t = map(np.asarray, (q, ell_v, ell_t))
    best = None
    for a in range(len(q)):
        for b in range(a + 1, len(q)):
            n = np.linalg.norm(q[b] - q[a])
            if best is None or n > best[0]:
                best = (n, a, b)
    _, a, b = best
    dq = q[b] - q[a]
    dl = ell_v[b] - ell_v[a]
    if np.linalg.norm(dl) < eps:
        raise ValueError("unexcited relay displacement")
    psi = math.atan2(dq[1], dq[0]) - math.atan2(dl[1], dl[0])
    psi = wrap(psi)
    R = rot(psi)
    x = q[a] - R @ ell_v[a]
    # Equal robust weights are the noiseless/low-noise special case. Huber weights
    # could be added here without changing the gauge construction.
    p_views = x[None, :] + (R @ ell_t.T).T
    p = np.mean(p_views, axis=0)
    return np.array([x[0], x[1], p[0], p[1], psi]), (a, b)


def make_data(d, psi=1.1, x=np.array([0.8, -0.4]), p=np.array([2.0, 1.3]),
              sigma=0.0, rng=None):
    rng = np.random.default_rng() if rng is None else rng
    # Two views with known global motion d along x. The second view is selected.
    q = np.array([[0., 0.], [d, 0.]])
    Rm = rot(-psi)
    ell_v = np.array([Rm @ (qk - x) for qk in q])
    ell_t = np.array([Rm @ (p - x) for _ in q])
    if sigma:
        q = q + rng.normal(0, sigma, q.shape)
        ell_v = ell_v + rng.normal(0, sigma, ell_v.shape)
        ell_t = ell_t + rng.normal(0, sigma, ell_t.shape)
    return q, ell_v, ell_t


def residual(z, q, ev, et):
    x, p, psi = z[:2], z[2:4], z[4]
    R = rot(psi)
    rv = q - x[None, :] - (R @ ev.T).T
    rt = p[None, :] - x[None, :] - (R @ et.T).T
    return np.concatenate([rv.ravel(), rt.ravel()])


def refine(z0, q, ev, et):
    out = least_squares(residual, z0, args=(q, ev, et), max_nfev=200,
                        xtol=1e-12, ftol=1e-12, gtol=1e-12)
    return out.x, np.linalg.norm(residual(out.x, q, ev, et)), out.nfev


def main():
    rng = np.random.default_rng(2048)
    # Prediction 1: noiseless nonzero displacement gives exact gauge recovery.
    exact = []
    for d in [0.01, 0.1, 1.0, 3.0]:
        q, ev, et = make_data(d, sigma=0, rng=rng)
        z, pair = two_view_init(q, ev, et)
        truth = np.array([.8, -.4, 2., 1.3, 1.1])
        exact.append([d, float(np.linalg.norm(z[:4] - truth[:4])), abs(wrap(z[4]-truth[4])), pair])

    # Prediction 2: with isotropic measurement noise, yaw RMSE scales as 1/d.
    # For two independent noisy vectors, first-order theory predicts sqrt(2)*sigma/d.
    sigma = 0.002
    ds = np.array([0.05, 0.1, 0.2, 0.4, 0.8, 1.6, 3.2])
    trials = 800
    sweep = []
    for d in ds:
        errs = []
        for _ in range(trials):
            q, ev, et = make_data(d, sigma=sigma, rng=rng)
            try:
                z, _ = two_view_init(q, ev, et)
                errs.append(abs(wrap(z[4] - 1.1)))
            except ValueError:
                errs.append(np.pi)
        rmse = float(np.sqrt(np.mean(np.square(errs))))
        predicted = math.sqrt(2) * sigma / d
        sweep.append([float(d), rmse, predicted, rmse / predicted])

    # Prediction 3: at zero relay displacement the angle is not identifiable;
    # numerically, tiny displacement has rapidly increasing error.
    small = []
    for d in [0.0, 1e-5, 1e-4, 1e-3, 1e-2]:
        errs = []
        rejected = 0
        for _ in range(300):
            q, ev, et = make_data(d, sigma=sigma, rng=rng)
            try:
                z, _ = two_view_init(q, ev, et, eps=1e-10)
                errs.append(abs(wrap(z[4] - 1.1)))
            except ValueError:
                rejected += 1
        small.append([d, float(np.sqrt(np.mean(np.square(errs)))) if errs else None, rejected])

    # Secondary mini comparison: nonlinear refinement from random gauge versus initializer.
    q, ev, et = make_data(1.5, sigma=0.01, rng=rng)
    zi, _ = two_view_init(q, ev, et)
    truth = np.array([.8, -.4, 2., 1.3, 1.1])
    init_runs, random_runs = [], []
    for _ in range(40):
        zr = np.array([rng.uniform(-2, 2), rng.uniform(-2, 2), rng.uniform(-2, 4),
                       rng.uniform(-2, 4), rng.uniform(-np.pi, np.pi)])
        zo, lr, nr = refine(zr, q, ev, et)
        _, li, ni = refine(zi, q, ev, et)
        random_runs.append([lr, nr, float(np.linalg.norm(zo[:4]-truth[:4]))])
        init_runs.append([li, ni, float(np.linalg.norm(zi[:4]-truth[:4]))])
    result = {
        "exact_recovery": exact,
        "noise_inverse_displacement": sweep,
        "near_singular": small,
        "refinement": {
            "random_median_residual": float(np.median(np.array(random_runs)[:,0])),
            "initializer_median_residual": float(np.median(np.array(init_runs)[:,0])),
            "random_median_nfev": float(np.median(np.array(random_runs)[:,1])),
            "initializer_median_nfev": float(np.median(np.array(init_runs)[:,1])),
            "random_target_error_median": float(np.median(np.array(random_runs)[:,2])),
            "initializer_target_error": float(np.median(np.array(init_runs)[:,2]))
        }
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()