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

# Reproducible toy verification plus a small optimizer experiment.
SEED = 7
np.random.seed(SEED); random.seed(SEED)


def avoided(a, b, g):
    m = (a+b)/2
    d = np.sqrt(((a-b)/2)**2 + g*g)
    return m+d, m-d


def toy_check():
    # Two uncoupled modes cross at theta=0: a,b have slopes +/-s.
    s, base = 0.8, 1.0
    thetas = np.linspace(-2, 2, 4001)
    rows=[]
    for g in [0.02, 0.05, 0.10, 0.20]:
        a = base+s*thetas; b = base-s*thetas
        lp,lm = avoided(a,b,g)
        # Positive eigenvalues in the central region; normalized gap as proposed.
        G=(lp-np.abs(lm))/lp
        i=int(np.argmin(G))
        # xi based on the second eigenvalue, and response d(lambda+)/d(theta).
        xi=1/np.log(lp/np.maximum(np.abs(lm),1e-12))
        # Exact curvature of lambda_plus: s^2*g^2 / ((s*theta)^2+g^2)^(3/2), maximal at resonance.
        response=(s*s*g*g)/((s*thetas)**2+g*g)**1.5
        ir=int(np.argmax(response))
        # Predictions: crossing/response at 0, splitting 2g, xi(0) ~ base/(2g).
        split=lp[i]-lm[i]
        rows.append(dict(g=g, theta_gap=float(thetas[i]), predicted_theta=0.,
                         split=float(split), predicted_split=2*g,
                         xi_at_cross=float(xi[np.argmin(np.abs(thetas))]),
                         predicted_xi=base/(2*g),
                         theta_response=float(thetas[ir])))
    # scaling fit xi vs 1/g, and normalized errors
    gs=np.array([r['g'] for r in rows]); xis=np.array([r['xi_at_cross'] for r in rows])
    slope=float(np.polyfit(1/gs,xis,1)[0])
    # Correlation prediction: a pure subleading mode has C(k)/C(0)=(|lambda1|/lambda0)^k.
    g=.1; a=base+s*thetas; b=base-s*thetas; lp,lm=avoided(a,b,g)
    j=np.argmin(np.abs(thetas)); ratio=abs(lm[j])/lp[j]
    ks=np.arange(1,11); corr=ratio**ks
    fit=float(np.polyfit(ks,np.log(corr),1)[0])
    memory_xi=float(-1/fit)
    return rows, slope, memory_xi


def make_data(n, seq=20):
    x=np.random.rand(n,seq,2).astype('float32')
    y=x.sum((1,2)).astype('float32')
    return x,y

# Torch is imported lazily so the mathematical check also works without CUDA.
def train(use_scheduler, seed=7, steps=500):
    import torch
    torch.manual_seed(seed); np.random.seed(seed)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        class RNN(torch.nn.Module):
            def __init__(self):
                super().__init__(); self.cell=torch.nn.RNNCell(2,24,nonlinearity='tanh'); self.out=torch.nn.Linear(24,1)
            def forward(self,x, collect=False):
                h=torch.zeros(x.shape[0],24,device=x.device); hs=[]
                for t in range(x.shape[1]):
                    h=self.cell(x[:,t],h); hs.append(h)
                return self.out(h).squeeze(1), hs
        model=RNN().to(device)
        opt=torch.optim.Adam(model.parameters(),lr=0.01)
        lr=0.01; records=[]
        for step in range(steps):
            x,y=make_data(64,20); x=torch.tensor(x,device=device); y=torch.tensor(y,device=device)
            pred,hs=model(x,True); loss=((pred-y)**2).mean()
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
            gap=entropy=xi=np.nan
            if use_scheduler and step%10==0 and step>=10:
                # Fit h_{t+1}=T h_t over batch and time, as the proposal suggests.
                A=torch.stack(hs[:-1],1).detach().reshape(-1,24)
                B=torch.stack(hs[1:],1).detach().reshape(-1,24)
                # T maps row states: A @ M ~= B; eigenvalues of M.
                M=torch.linalg.lstsq(A,B,rcond=None).solution
                ev=torch.linalg.eigvals(M).real.abs().sort(descending=True).values
                l0=max(float(ev[0]),1e-8); l1=float(ev[1]) if len(ev)>1 else 0
                gap=(l0-l1)/l0; xi=1/max(math.log(l0/max(l1,1e-8)),1e-8)
                # Projection entropy onto leading right eigenvectors, normalized.
                vals,vecs=torch.linalg.eig(M); ix=torch.argsort(vals.real.abs(),descending=True)[:3]
                V=vecs[:,ix].real; z=A@V; p=(z*z).mean(0); p=p/(p.sum()+1e-8)
                entropy=float(-(p*torch.log(p+1e-8)).sum())
                if gap<0.12 and entropy>0.7:
                    lr=max(lr*0.5,0.0005)
                else:
                    lr=float(np.clip(lr*(gap/0.25),0.0005,0.01))
                for pg in opt.param_groups: pg['lr']=lr
            records.append((float(loss.detach().cpu()), gap, entropy, lr, xi))
        tail=np.array([r[0] for r in records[-50:]])
        return dict(device=device, final=float(records[-1][0]), best=float(np.min([r[0] for r in records])), tail=float(tail.mean()),
                    lr_final=float(lr), detections=int(sum(np.isfinite(r[1]) and r[1]<.12 and r[2]>.7 for r in records)), records=records)
    except Exception as e:
        if device=='cuda':
            # Retry the identical experiment on CPU if CUDA is unavailable/fragile.
            torch.cuda.empty_cache()
            original=torch.cuda.is_available
            torch.cuda.is_available=lambda: False
            try:
                return train(use_scheduler, seed, steps)
            finally:
                torch.cuda.is_available=lambda: original
        raise

if __name__=='__main__':
    rows,slope,memory_xi=toy_check()
    out={'toy':rows,'xi_slope_vs_1_over_g':slope,'memory_xi_numeric':memory_xi,'memory_xi_predicted':rows[2]['xi_at_cross']}
    try:
        out['baseline']=train(False)
        out['idea']=train(True)
    except Exception as e:
        out['training_error']=repr(e)
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))