import json, random
from pathlib import Path
import numpy as np
import torch

SEED = 2291
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)

def laplacian(a, b):
    # Directed path 1 -> 2 -> 3, weighted in-degree Laplacian.
    return np.array([[0., 0., 0.], [-a, a, 0.], [0., -b, b]])

def spectrum(a, b):
    return np.linalg.eig(laplacian(a, b))

def gap(a, b):
    ev, _ = spectrum(a, b)
    nz = [z for z in ev if abs(z) > 1e-7]
    return float(min(abs(nz[i]-nz[j])/(1+abs(nz[i])+abs(nz[j]))
                    for i in range(len(nz)) for j in range(i)))

def cond(a, b):
    return float(np.linalg.cond(spectrum(a, b)[1]))

def discriminant_nonzero(a, b):
    # q(s)=(s-a)(s-b), so Disc(q)=(a-b)^2.
    return float((a-b)**2)

def analytic_gap(a,b):
    return abs(a-b)/(1+a+b)

def math_sweep():
    rows=[]; s=1.0
    for eps in [1e-4,3e-4,1e-3,3e-3,1e-2,3e-2,1e-1]:
        a,b=s+eps,s-eps
        rows.append({'eps':eps,'gap':gap(a,b),'analytic_gap':analytic_gap(a,b),
                     'discriminant':discriminant_nonzero(a,b),'condition':cond(a,b)})
    x=np.log([r['eps'] for r in rows[:5]])
    gs=float(np.polyfit(x,np.log([r['gap'] for r in rows[:5]]),1)[0])
    ds=float(np.polyfit(x,np.log([r['discriminant'] for r in rows[:5]]),1)[0])
    cp=[r['condition']*r['eps'] for r in rows[:5]]
    L=laplacian(1.,1.)
    ev=np.linalg.eigvals(L)
    # At a=b, algebraic multiplicity of eigenvalue 1 is two but nullity of
    # L-I is one: this is the exact defective boundary.
    nullity=int(3-np.linalg.matrix_rank(L-np.eye(3),tol=1e-10))
    return {'rows':rows,'observed_gap_log_slope':gs,
            'observed_discriminant_log_slope':ds,'condition_times_epsilon':cp,
            'exact_collision':{'eigenvalues':ev.tolist(),
                'disc_q':discriminant_nonzero(1.,1.),
                'mult_1_algebraic':2,'nullity_L_minus_I':nullity,
                'defective':bool(nullity<2)},
            'predictions':{'gap_slope':1.,'discriminant_slope':2.,
                           'condition_scaling':'condition*epsilon approximately constant'}}

def train(gated, steps=700, gmin=.03, beta=5000.):
    torch.manual_seed(SEED)
    theta=torch.tensor([.4,-.4],dtype=torch.float64,requires_grad=True)
    opt=torch.optim.Adam([theta],lr=.025); omin=1e-3; hist=[]
    for k in range(steps):
        a,b=omin+torch.nn.functional.softplus(theta)
        task=(a-1.)**2+(b-1.)**2
        gg=torch.abs(a-b)/(1.+a+b)
        barrier=beta*torch.relu(torch.tensor(gmin,dtype=torch.float64)-gg)**2 if gated else 0.
        loss=task+barrier
        opt.zero_grad(); loss.backward(); opt.step()
        if k in (0,99,299,699):
            hist.append({'step':k+1,'a':float(a.detach()),'b':float(b.detach()),
                         'task':float(task.detach()),'gap':float(gg.detach()),
                         'barrier':float(barrier.detach()) if gated else 0.})
    af,bf=[float(x.detach()) for x in (a,b)]
    return {'a':af,'b':bf,'gap':gap(af,bf),'condition':cond(af,bf),
            'task':(af-1)**2+(bf-1)**2,'history':hist}

def main():
    math_result=math_sweep()
    base=train(False)
    sweep={str(g):train(True,gmin=g) for g in [1e-4,1e-3,1e-2,3e-2]}
    out={'math':math_result,'training':{'baseline_unconstrained_positive':base,
          'gap_gated_gmin_sweep':sweep,'beta':5000.},'seed':SEED}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()