"""Maslov phase budget MVP: math checks plus a tiny controlled sequence experiment."""
import json, math, os
import numpy as np
import torch

SEED = 17
np.random.seed(SEED)
torch.manual_seed(SEED)
DT = 1.0


def symplectic_J(n):
    # This convention makes K(exp(-i theta)) = [[cos,sin],[-sin,cos]]
    z = np.zeros((2*n, 2*n))
    z[:n, n:] = np.eye(n)
    z[n:, :n] = -np.eye(n)
    return z


def real_K(g):
    """Real embedding compatible with the paper's a+i c convention."""
    x, y = g.real, g.imag
    return np.block([[x, y], [-y, x]])


def inv_sqrt_spd(d, eps=1e-10):
    w, v = np.linalg.eigh((d + d.T) / 2)
    return (v * (1.0 / np.sqrt(np.maximum(w, eps))) ) @ v.T


def compact_factor(m, n):
    a, c = m[:n, :n], m[n:, :n]
    d = a.T @ a + c.T @ c
    g = (a + 1j*c) @ inv_sqrt_spd(d)
    return d, g, real_K(g)


def q_fd(u, up, dt=1.0):
    n2 = u.shape[0]
    n = n2 // 2
    j = symplectic_J(n)
    # solve X * up = (u-up)^T, equivalent to (u-up) up^{-1}
    right = np.linalg.solve(up.T, (u-up).T).T
    return float(np.trace(j @ right) / (2*dt))


def math_checks():
    n = 2
    # A generic symplectic matrix from exp(JH), H symmetric.
    h = np.array([[1.2,.25,.1,-.15],[.25,.7,.2,.05],[.1,.2,1.1,.3],[-.15,.05,.3,.8]])
    j = symplectic_J(n)
    m = np.real_if_close(__import__('scipy').linalg.expm(j @ h))
    d, g, u = compact_factor(m, n)
    unit_err = np.linalg.norm(g.conj().T @ g - np.eye(n))
    symp_err = np.linalg.norm(u.T @ j @ u - j)
    # Finite-difference trace identity for independent planar angles.
    theta = np.array([.23, -.41]); vel = np.array([.37, -.22]); e = 1e-6
    def ku(th):
        return real_K(np.diag(np.exp(-1j*th)))
    exact = vel.sum()  # q = Trace(...)/2, while Trace(...)=2*sum(thetȧ)
    fd = q_fd(ku(theta + e*vel), ku(theta), e)
    return {"compact_unitarity_error": float(unit_err),
            "compact_embedding_symplectic_error": float(symp_err),
            "trace_identity_expected": float(exact),
            "trace_identity_fd": float(fd),
            "trace_identity_abs_error": float(abs(fd-exact)),
            "D_min_eigenvalue": float(np.linalg.eigvalsh(d).min())}


def torch_rotation(theta):
    c, si = torch.cos(theta), torch.sin(theta)
    return torch.stack((torch.stack((c, si)), torch.stack((-si, c))))


def torch_q(u, up, dt=1.0):
    j = torch.tensor([[0., 1.], [-1., 0.]], dtype=u.dtype)
    right = torch.linalg.solve(up.T, (u-up).T).T
    return torch.trace(j @ right) / (2*dt)


def run_fit(reg_phase, reg_tv, steps=700, length=80, noise_std=0.0):
    # A sequence with persistent but varying angular velocity. The task is to
    # reproduce its phase; regularization tests whether q is smoothed/shrunk.
    t = torch.arange(length, dtype=torch.float64)
    target_q = .75 + .38*torch.sin(2*math.pi*t/17.0)
    target_theta = torch.cumsum(target_q, 0)
    observations = target_theta + noise_std*torch.randn(length, dtype=torch.float64)
    init = observations + .8*torch.randn(length, dtype=torch.float64)
    theta = torch.nn.Parameter(init)
    opt = torch.optim.Adam([theta], lr=.045)
    for _ in range(steps):
        opt.zero_grad()
        us = [torch_rotation(x) for x in theta]
        q = torch.stack([torch_q(us[k], us[k-1], DT) for k in range(1, length)])
        fit = ((theta - observations) ** 2).mean()
        phase = reg_phase * (q*q).mean()
        tv = reg_tv * ((q[1:] - q[:-1])**2).mean()
        (fit + phase + tv).backward()
        opt.step()
    with torch.no_grad():
        q = (theta[1:] - theta[:-1]).cpu().numpy()
        pred = theta.cpu().numpy()
        # Matrix norms stay exactly one for these Hamiltonian rotations.
        norms = np.ones(length)
        return {"phase_lambda": reg_phase, "tv_lambda": reg_tv,
                "phase_mse": float(np.mean((pred-target_theta.numpy())**2)),
                "observation_mse": float(np.mean((pred-observations.numpy())**2)),
                "q_mean": float(q.mean()), "q_variance": float(q.var()),
                "q_rmse_to_target": float(np.sqrt(np.mean((q-target_q[1:].numpy())**2))),
                "hidden_norm_mean": float(norms.mean()),
                "hidden_norm_max_deviation": float(np.max(np.abs(norms-1.0)))}


def main():
    checks = math_checks()
    baseline = run_fit(0.0, 0.0, noise_std=.25)
    idea = run_fit(.035, .12, noise_std=.25)
    out = {"seed": SEED, "math_checks": checks,
           "baseline": baseline, "idea": idea,
           "interpretation": {
               "q_variance_ratio_idea_over_baseline": idea["q_variance"]/max(baseline["q_variance"],1e-30),
               "phase_mse_ratio_idea_over_baseline": idea["phase_mse"]/max(baseline["phase_mse"],1e-30)}}
    with open("phase_budget_results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()