import json, math, random
import numpy as np
import torch

SEED = 7
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)


def rot(a):
    c, s = torch.cos(a), torch.sin(a)
    z = torch.zeros_like(a)
    return torch.stack([torch.stack([c, -s], -1), torch.stack([s, c], -1)], -2)


def residual(theta, lv, lt, q, y):
    # theta = [p_x,p_y,x_x,x_y,psi]
    p, x, psi = theta[:2], theta[2:4], theta[4]
    R = rot(psi)
    r = x + torch.einsum('ij,tj->ti', R, lv) - q
    s = p - x - torch.einsum('ij,tj->ti', R, lt) - y
    return torch.cat([r.reshape(-1), s.reshape(-1)])


def jacobian(theta, lv, lt, q, y):
    return torch.autograd.functional.jacobian(
        lambda z: residual(z, lv, lt, q, y), theta
    ).detach().numpy()


def make_window(d, theta_true, T=2, noise=0.0, rng=None):
    rng = np.random.default_rng(SEED) if rng is None else rng
    lv = np.zeros((T, 2)); lv[:, 0] = np.linspace(0.0, d, T)
    lt = np.tile(np.array([0.4, -0.25]), (T, 1))
    p, x, psi = theta_true[:2], theta_true[2:4], theta_true[4]
    c, s = np.cos(psi), np.sin(psi)
    R = np.array([[c, -s], [s, c]])
    q = x + lv @ R.T
    y = p - x - lt @ R.T
    if noise:
        q += rng.normal(0, noise, q.shape)
        y += rng.normal(0, noise, y.shape)
    return tuple(torch.tensor(a) for a in (lv, lt, q, y))


def sigma_for(d, T=2):
    true = np.array([.7, -.2, .15, .1, .35])
    lv, lt, q, y = make_window(d, true, T=T)
    th = torch.tensor(true, requires_grad=True)
    sv = np.linalg.svd(jacobian(th, lv, lt, q, y), compute_uv=False)
    return float(sv[-1]), float(sv[0]), float(np.linalg.matrix_rank(jacobian(th, lv, lt, q, y), tol=1e-9))


def optimize(gated, d, T=4, steps=220, noise=.025, seed=0):
    rng = np.random.default_rng(seed)
    true = np.array([.7, -.2, .15, .1, .35])
    lv, lt, q, y = make_window(d, true, T=T, noise=noise, rng=rng)
    # Start with a substantially wrong calibration. Optimize only theta, isolating the gate mechanism.
    theta = torch.tensor([-.5, .6, .8, -.7, -1.0], requires_grad=True)
    opt = torch.optim.Adam([theta], lr=.035)
    # threshold calibrated from the well-excited d=1 window
    tau = sigma_for(1.0, T=T)[0] * .65
    for _ in range(steps):
        rr = residual(theta, lv, lt, q, y)
        consistency = (rr * rr).mean()
        if gated:
            J = torch.autograd.functional.jacobian(lambda z: residual(z,lv,lt,q,y), theta, create_graph=False)
            sv = torch.linalg.svdvals(J)
            sig = sv[-1].detach()
            gate = torch.clamp(sig / tau, max=1.0)
            # A small hinge penalty follows the proposed objective but does not dominate fitting.
            loss = gate * consistency + .02 * torch.relu(torch.tensor(tau) - sv[-1].detach())**2
        else:
            loss = consistency
        opt.zero_grad(); loss.backward(); opt.step()
    err = float(torch.linalg.norm(theta.detach() - torch.tensor(true)))
    return err, float(theta.detach()[4]), float(consistency.detach())


def main():
    true = np.array([.7, -.2, .15, .1, .35])
    # Prediction 1: d=0 leaves yaw unidentifiable (rank deficient).
    rank0 = sigma_for(0.0, T=2)[2]
    rank1 = sigma_for(1.0, T=2)[2]
    # Prediction 2: for small motion, sigma_min is linear in d.
    ds = np.geomspace(.005, 1.0, 9)
    sigmas = np.array([sigma_for(float(d), T=2)[0] for d in ds])
    slope = float(np.polyfit(np.log(ds), np.log(sigmas), 1)[0])
    ratio = float(np.median(sigmas[:5] / ds[:5]))
    # Prediction 3: longer windows improve excitation, with diminishing/no worse trend.
    lengths = [2, 3, 4, 6, 8]
    sigT = [sigma_for(.8, T=T)[0] for T in lengths]
    monotone = all(sigT[i+1] >= sigT[i] - 1e-10 for i in range(len(sigT)-1))

    # Secondary mini-experiment: difficult low-excitation windows versus ordinary training.
    train_ds = [.03, .1, .3, 1.0]
    results = {}
    for d in train_ds:
        b = [optimize(False, d, seed=20+i)[0] for i in range(3)]
        g = [optimize(True, d, seed=20+i)[0] for i in range(3)]
        results[str(d)] = {'sigma_min': sigma_for(d,4)[0],
                           'baseline_error_mean': float(np.mean(b)),
                           'gated_error_mean': float(np.mean(g))}

    report = {
      'seed': SEED,
      'predictions': {
        'zero_motion_rank_deficiency': {'predicted': 'rank drops below 5 at d=0', 'observed_rank_d0': rank0, 'observed_rank_d1': rank1, 'confirmed': rank0 < 5 and rank1 == 5},
        'small_motion_scaling': {'predicted': 'sigma_min proportional to d, log-log slope 1', 'observed_loglog_slope': slope, 'median_sigma_over_d': ratio, 'd_values': ds.tolist(), 'sigma_values': sigmas.tolist(), 'confirmed': abs(slope-1) < .08},
        'window_length': {'predicted': 'sigma_min nondecreasing with T', 'T_values': lengths, 'sigma_values': sigT, 'confirmed': monotone}
      },
      'mini_experiment': results,
      'notes': 'The residual parameterization includes both vehicle and target views. At exactly zero motion, yaw is ambiguous; nonzero motion restores full local rank.'
    }
    with open('results.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()