import json, math, time
from pathlib import Path
import numpy as np


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


def spectral_init(n, edges, phi, weights, ref=0, iters=30):
    C = np.zeros((n, n), dtype=np.complex128)
    for (i, j), p, w in zip(edges, phi, weights):
        C[i, j] += w * np.exp(1j*p)
        C[j, i] += w * np.exp(-1j*p)
    # power iteration, with Rayleigh-sign stabilization
    v = np.ones(n, dtype=np.complex128) / np.sqrt(n)
    for _ in range(iters):
        q = C @ v
        nq = np.linalg.norm(q)
        if nq == 0: break
        v = q / nq
    z = v / np.maximum(np.abs(v), 1e-10)
    theta = np.angle(z)
    return wrap(theta - theta[ref])


def gauge_rmse(est, truth):
    d = wrap(est - truth)
    # both are nominally gauged, but optimize residual global gauge for robustness
    shift = np.angle(np.mean(np.exp(1j*(est-truth))))
    return float(np.sqrt(np.mean(wrap(d-shift)**2)))


def edge_rmse(theta, truth, edges):
    return float(np.sqrt(np.mean([wrap((theta[i]-theta[j])-(truth[i]-truth[j]))**2 for i,j in edges])))


def centered_jacobian_margin(n, edges, weights):
    """Smallest nonzero singular value after removing the global phase gauge."""
    J = np.zeros((len(edges), n))
    for k, (i, j) in enumerate(edges):
        J[k, i] = weights[k]
        J[k, j] = -weights[k]
    P = np.eye(n) - np.ones((n, n)) / n
    sv = np.linalg.svd(J @ P, compute_uv=False)
    nz = sv[sv > 1e-8]
    return float(nz[-1]) if nz.size else 0.0


def gated_refine(init, edges, phi, weights, rho, eta, threshold=0.35):
    """Use spectral estimate directly in the low-eta regime; otherwise refine."""
    if eta <= threshold:
        return init.copy(), False
    return gd_refine(init, edges, phi, weights, lr=0.04, steps=100), True


def gd_refine(init, edges, phi, weights, lr=0.08, steps=100):
    # WLS on circular relative-phase residuals; node 0 fixed as the gauge.
    x = init.copy(); x[0] = 0.0
    for _ in range(steps):
        g = np.zeros_like(x)
        for (i,j), p, w in zip(edges, phi, weights):
            r = wrap(x[i]-x[j]-p)
            # derivative of 1-cos(r), robust and smooth
            s = w*np.sin(r)
            g[i] += s; g[j] -= s
        x[1:] -= lr*g[1:]
        x[0] = 0.0
    return wrap(x)


def trial(seed, n=20, degree=3, noise=0.35):
    rng=np.random.default_rng(seed)
    truth=wrap(rng.normal(0, 1.0, n)); truth -= truth[0]
    # connected ring plus random undirected edges
    edges={(i,(i+1)%n) for i in range(n)}
    target=n*degree//2
    while len(edges)<target:
        i,j=sorted(rng.choice(n,2,replace=False)); edges.add((i,j))
    edges=sorted(edges)
    weights=np.ones(len(edges))
    phi=np.array([wrap(truth[i]-truth[j]+rng.normal(0,noise)) for i,j in edges])
    spec=spectral_init(n,edges,phi,weights)
    rand=wrap(rng.uniform(-np.pi,np.pi,n)); rand -= rand[0]
    spec_loss=edge_rmse(spec,truth,edges); rand_loss=edge_rmse(rand,truth,edges)
    spec_final=gd_refine(spec,edges,phi,weights)
    rand_final=gd_refine(rand,edges,phi,weights)
    rho = centered_jacobian_margin(n, edges, weights)
    residual=np.array([wrap(spec[i]-spec[j]-p) for (i,j),p in zip(edges,phi)])
    mad=float(np.median(np.abs(residual-np.median(residual))))
    eta=mad/rho
    gated, did_refine = gated_refine(spec, edges, phi, weights, rho, eta)

    # robust residual and normalized perturbation proxy
    return dict(rho=rho, eta=eta, gate_refined=did_refine, gated_rmse=gauge_rmse(gated,truth), spec_init_rmse=gauge_rmse(spec,truth), rand_init_rmse=gauge_rmse(rand,truth),
                spec_edge_rmse=spec_loss, rand_edge_rmse=rand_loss,
                spec_final_rmse=gauge_rmse(spec_final,truth), rand_final_rmse=gauge_rmse(rand_final,truth),
                spec_final_edge_rmse=edge_rmse(spec_final,truth,edges), rand_final_edge_rmse=edge_rmse(rand_final,truth,edges),
                mad=mad)


def jacobian_check(n=8):
    # Each edge observes theta_i-theta_j. J has one row per edge.
    edges=[(i,(i+1)%n) for i in range(n)] + [(i,(i+2)%n) for i in range(n)]
    J=np.zeros((len(edges),n))
    for k,(i,j) in enumerate(edges): J[k,i]=1; J[k,j]=-1
    one=np.ones(n)
    gauge_norm=float(np.linalg.norm(J@one))
    P=np.eye(n)-np.ones((n,n))/n
    sv=np.linalg.svd(J@P,compute_uv=False)
    rho=float(sv[-2] if sv.size>1 and sv[-1]<1e-8 else sv[-1])
    # constrained singular values should equal the nonzero singular values of J
    nonzero=np.linalg.svd(J,compute_uv=False); nonzero=nonzero[nonzero>1e-8]
    return dict(gauge_null_norm=gauge_norm, rho=rho, nonzero_min=float(nonzero[-1]),
                rho_matches=float(abs(rho-nonzero[-1])))


def main():
    t=time.time(); jac=jacobian_check()
    allrows=[]
    for noise in (0.15,0.35,0.70,1.10):
        rows=[trial(1000+q,noise=noise) for q in range(40)]
        summary={'noise':noise}
        for key in rows[0]: summary[key+'_mean']=float(np.mean([r[key] for r in rows]))
        summary['spec_better_init_rate']=float(np.mean([r['spec_init_rmse']<r['rand_init_rmse'] for r in rows]))
        summary['spec_better_final_rate']=float(np.mean([r['spec_final_rmse']<r['rand_final_rmse'] for r in rows]))
        allrows.append(summary)
    out={'jacobian_check':jac,'results':allrows,'runtime_sec':time.time()-t}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()