import os, sys, json, math, time, random
import numpy as np
import torch
from torch import nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, train_model, sweep_baseline, evaluate, make_report, count_params

EPOCHS=8; BATCH=128; NTRAIN=800; NTEST=400
LRS=[1e-3,3e-3,6e-3]

class SU11Scan(nn.Module):
    """Shared potential/readout architecture; exact or Euler transition is the only difference."""
    def __init__(self, exact=True, pairs=16, h=.12, xi_scale=.07):
        super().__init__(); self.exact=exact; self.pairs=pairs; self.h=h; self.xi_scale=xi_scale
        self.fmap=nn.Linear(3,2*pairs); self.head=nn.Linear(4*pairs,1)
    def forward(self,x,return_state=False):
        B=x.shape[0]; seq=x.view(B,-1,3); dev=x.device
        a=torch.ones(B,self.pairs,device=dev,dtype=torch.complex64); b=torch.zeros_like(a)
        xi=torch.arange(self.pairs,device=dev,dtype=torch.float32)*self.xi_scale
        for k in range(seq.shape[1]):
            raw=self.fmap(seq[:,k]); f=torch.complex(raw[:,:self.pairs],raw[:,self.pairs:])
            f=.7*torch.tanh(f.real)+1j*.7*torch.tanh(f.imag)
            r=torch.abs(f).float(); phase=torch.exp((2j*math.pi*k*xi).to(torch.complex64))[None,:]
            if self.exact:
                q=self.h*r; c=torch.cosh(q).to(torch.complex64)
                s=torch.where(r>1e-7,torch.sinh(q)/r,torch.full_like(r,self.h)).to(torch.complex64)
            else:
                c=torch.ones_like(r).to(torch.complex64); s=torch.full_like(r,self.h).to(torch.complex64)
            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); out=self.head(feat)
        return (out,a,b) if return_state else out

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
def ds(seed): return get_dataset('dynamics',seed,n_train=NTRAIN,n_test=NTEST)
def train_one(kind,lr,seed,capture=False):
    seed_all(seed); d=ds(seed); m=SU11Scan(exact=(kind=='idea'))
    net,metric,_=train_model(m,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
    if net is None: raise RuntimeError('training failed')
    out={'metric':float(metric)}
    if capture:
        with torch.no_grad():
            x=d['xte'].to(next(net.parameters()).device); _,a,b=net(x,True)
            inv=(a.abs()**2-b.abs()**2-1).abs()
            out.update(invariant_mean_error=float(inv.mean().cpu()),invariant_max_error=float(inv.max().cpu()),mean_abs_a=float(a.abs().mean().cpu()),mean_abs_b=float(b.abs().mean().cpu()))
    return out
def metric_fn(kind,lr): return lambda s: train_one(kind,lr,s)['metric']
def main():
    t=time.time(); grid=[{'lr':v} for v in LRS]
    base=sweep_baseline(lambda c:metric_fn('baseline',c['lr']),grid)
    idea_sweep=[]
    for c in grid:
        r=evaluate(metric_fn('idea',c['lr']),seeds=(0,1,2,3)); idea_sweep.append({'cfg':c,'mean':r['mean']})
    best=min(idea_sweep,key=lambda z:z['mean'])['cfg']; idea=evaluate(metric_fn('idea',best['lr']))
    bs=train_one('baseline',base['best_cfg']['lr'],0,True); ins=train_one('idea',best['lr'],0,True)
    sig={'prediction':'Exact exponential preserves |a|^2-|b|^2=1 while Euler drift accumulates.','predicted_exact_invariant_error':0.0,'observed_exact_mean_error':ins['invariant_mean_error'],'observed_exact_max_error':ins['invariant_max_error'],'observed_euler_mean_error':bs['invariant_mean_error'],'observed_euler_max_error':bs['invariant_max_error'],'confirmed':bool(ins['invariant_max_error']<1e-5 and bs['invariant_max_error']>ins['invariant_max_error']*10)}
    rep=make_report('dynamics','su11_scan_matched',base,idea,{'mechanism_signature':sig,'idea_sweep':idea_sweep,'idea_best_cfg':best,'baseline_definition':'Euler discretization of the same learned SU(1,1) generator; exact exponential is the sole intervention.','protocol_note':'Identical trained systems, data, parameters, optimizer, epochs, and learning-rate grid; lower MSE is better.'})
    rep['timing_seconds']=time.time()-t; rep['parameter_counts']={'baseline':count_params(SU11Scan(False)),'idea':count_params(SU11Scan(True))}
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()