import json, math, random, time
import numpy as np
import torch
from torch import nn

SEED = 2374
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def exact_step(a, b, f, x, xi, h):
    r = abs(f)
    q = h*r
    c = np.cosh(q)
    s = h if r == 0 else np.sinh(q)/r
    ph = np.exp(2j*np.pi*x*xi)
    return c*a + s*np.conj(f)*ph*b, c*b + s*f*np.conj(ph)*a

def euler_step(a, b, f, x, xi, h):
    ph = np.exp(2j*np.pi*x*xi)
    return a + h*np.conj(f)*ph*b, b + h*f*np.conj(ph)*a

def verify_math():
    # Random complex potentials and phases; start exactly at the claimed initial state.
    a, b = 1+0j, 0+0j
    ae, be = a, b
    exact_err, euler_err = [], []
    for k in range(1000):
        f = (np.random.randn()+1j*np.random.randn()) * 0.7
        a,b = exact_step(a,b,f,k,0.17,0.025)
        ae,be = euler_step(ae,be,f,k,0.17,0.025)
        exact_err.append(abs(abs(a)**2-abs(b)**2-1))
        euler_err.append(abs(abs(ae)**2-abs(be)**2-1))
    # Also directly check M^2=|f|^2 I numerically.
    f = .31-.77j; x=.43; xi=.29
    ph=np.exp(2j*np.pi*x*xi)
    M=np.array([[0,np.conj(f)*ph],[f*np.conj(ph),0]],complex)
    algebra_err=float(np.max(np.abs(M@M-(abs(f)**2)*np.eye(2))))
    return {
        "exact_max_invariant_error": float(max(exact_err)),
        "exact_final_invariant_error": float(exact_err[-1]),
        "euler_final_invariant_error": float(euler_err[-1]),
        "euler_max_invariant_error": float(max(euler_err)),
        "M2_error": algebra_err,
    }

# Torch version of the same elementwise exponential update.
class SU11Classifier(nn.Module):
    def __init__(self, hidden=2, h=.12):
        super().__init__(); self.h=h; self.hidden=hidden
        self.fmap=nn.Linear(1, 2*hidden)
        self.out=nn.Linear(4*hidden, 2)
    def forward(self, x):
        B,T,_=x.shape
        a=torch.ones(B,self.hidden,device=x.device,dtype=torch.complex64)
        b=torch.zeros_like(a)
        xi=torch.arange(self.hidden,device=x.device,dtype=torch.float32)*.07
        for k in range(T):
            raw=self.fmap(x[:,k])
            f=torch.complex(raw[:,:self.hidden],raw[:,self.hidden:])
            # bounded potentials avoid an irrelevant overflow regime in this tiny test
            f=.7*torch.tanh(f.real)+1j*.7*torch.tanh(f.imag)
            r=torch.abs(f); q=self.h*r
            c=torch.cosh(q.float()).to(torch.complex64)
            s=torch.where(r>1e-7, torch.sinh(q.float())/r.float(), torch.full_like(r,self.h)).to(torch.complex64)
            phase=torch.exp(2j*math.pi*k*xi).to(torch.complex64)[None,:]
            aa=c*a+s*torch.conj(f)*phase*b
            bb=c*b+s*f*torch.conj(phase)*a
            a,b=aa,bb
        feat=torch.cat([a.real,a.imag,b.real,b.imag],1)
        return self.out(feat)

class TanhRNNClassifier(nn.Module):
    def __init__(self, hidden=4):
        super().__init__(); self.cell=nn.RNNCell(1,hidden,nonlinearity='tanh'); self.out=nn.Linear(hidden,2)
    def forward(self,x):
        z=torch.zeros(x.shape[0],self.cell.hidden_size,device=x.device)
        for k in range(x.shape[1]): z=self.cell(x[:,k],z)
        return self.out(z)

def make_data(n, T):
    x=np.random.randn(n,T,1).astype('float32')
    y=(x.sum(axis=1)[:,0]>0).astype('int64')
    return torch.from_numpy(x),torch.from_numpy(y)

def train(model, tr, va, device, steps=350):
    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=0.01); lossfn=nn.CrossEntropyLoss()
    x,y=tr[0].to(device),tr[1].to(device); xv,yv=va[0].to(device),va[1].to(device)
    t0=time.time(); losses=[]
    for i in range(steps):
        opt.zero_grad(); loss=lossfn(model(x),y); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
        if i in (0,steps-1): losses.append(float(loss.detach().cpu()))
    with torch.no_grad():
        vl=float(lossfn(model(xv),yv).cpu()); acc=float((model(xv).argmax(1)==yv).float().mean().cpu())
    return {"initial_train_loss":losses[0],"final_train_loss":losses[-1],"validation_loss":vl,"validation_accuracy":acc,"seconds":time.time()-t0,"parameters":sum(p.numel() for p in model.parameters())}

def train_compare():
    torch.manual_seed(SEED); tr=make_data(128,32); va=make_data(256,32)
    # CPU is the safe default; CUDA errors fall back as required.
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        torch.manual_seed(SEED); su=train(SU11Classifier(),tr,va,device)
        torch.manual_seed(SEED); base=train(TanhRNNClassifier(),tr,va,device)
    except Exception as e:
        device=torch.device('cpu'); torch.manual_seed(SEED); su=train(SU11Classifier(),tr,va,device)
        torch.manual_seed(SEED); base=train(TanhRNNClassifier(),tr,va,device)
        su['fallback_error']=str(e)
    return {"device":str(device),"su11":su,"tanh_rnn":base}

def main():
    out={"math":verify_math(),"training":train_compare()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()