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

SEED=449; TAU=.20; ALPHA=(.5,.5); HIDDEN=64

class SwitchedSSM(nn.Module):
    def __init__(self, input_shape, out_dim, comm_lambda=0.0):
        super().__init__(); self.d=HIDDEN; self.lam=float(comm_lambda)
        eye=torch.eye(self.d)
        self.A=nn.Parameter(-.7*eye[None].repeat(2,1,1)+.03*torch.randn(2,self.d,self.d))
        self.U=nn.Parameter(.08*torch.randn(2,3,self.d)); self.head=nn.Linear(self.d,out_dim)
    def forward(self,x):
        z=x.reshape(x.shape[0],-1,3); h=torch.zeros(x.shape[0],self.d,device=x.device,dtype=x.dtype)
        B=[ALPHA[i]*TAU*self.A[i] for i in range(2)]
        E=[torch.matrix_exp(b) for b in B]
        for t in range(z.shape[1]):
            for i in range(2): h=h@E[i].T+(ALPHA[i]*z[:,t,:])@self.U[i]
        return self.head(h)
    def comm_penalty(self):
        B0=ALPHA[0]*TAU*self.A[0]; B1=ALPHA[1]*TAU*self.A[1]
        return ((B0@B1-B1@B0)**2).sum()

def train_idea(model, ds, epochs=30, lr=.003, batch=128, weight_decay=0.):
    # The intervention is a new loss, hence a local loop; all other settings match train_model.
    try: device='cuda' if torch.cuda.is_available() else 'cpu'; model=model.to(device)
    except Exception: device='cpu'; model=model.to(device)
    x,y=ds['xtr'].to(device),ds['ytr'].to(device); opt=torch.optim.Adam(model.parameters(),lr=lr,weight_decay=weight_decay)
    lossf=nn.CrossEntropyLoss() if ds['task']=='classification' else nn.MSELoss(); hist=[]
    for _ in range(epochs):
        model.train(); order=torch.randperm(len(x),device=device); total=0.
        for ix in order.split(batch):
            pred=model(x[ix]); loss=lossf(pred,y[ix])+model.lam*model.comm_penalty()
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.); opt.step(); total+=float(loss.detach())
        hist.append(total/max(1,(len(x)+batch-1)//batch))
    model.eval()
    with torch.no_grad():
        pred=model(ds['xte'].to(device)); metric=float(lossf(pred,ds['yte'].to(device)).cpu())
    return model,metric,hist

def idea_factory(cfg, seed, ds):
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    return train_idea(SwitchedSSM(ds['input_shape'],ds['out_dim'],cfg['comm_lambda']),ds,epochs=cfg['epochs'],lr=cfg['lr'])

def baseline_factory(cfg, seed, ds):
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    # Matched baseline: exactly the same switched architecture, but no proposed penalty.
    net=SwitchedSSM(ds['input_shape'],ds['out_dim'],comm_lambda=0.0)
    return train_model(net,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,weight_decay=0.)

def eval_side(factory, cfg, seeds):
    vals=[]; models=[]
    for s in seeds:
        ds=get_dataset('dynamics',s,n_train=400,n_test=200)
        out=factory(cfg,s,ds); models.append(out[0]); vals.append(float(out[1]))
    return {'config':cfg,'per_seed':vals,'mean':float(np.mean(vals)),'models':models}

def model_input_sensitivity(models, seeds):
    vals=[]
    for model,s in zip(models,seeds):
        ds=get_dataset('dynamics',s,n_train=8,n_test=8); dev=next(model.parameters()).device
        x=ds['xte'][:4].to(dev).clone().requires_grad_(True)
        model.eval(); out=model(x).sum(); g=torch.autograd.grad(out,x)[0]
        vals.append(float(g.detach().norm()/max(1,g.numel())**0.5))
    return float(np.mean(vals))

def diagnostics(models,seeds):
    comm=[]; gain=[]
    for model,s in zip(models,seeds):
        A=model.A.detach().cpu(); B=[.5*TAU*A[i] for i in range(2)]
        phi=torch.matrix_exp(B[1])@torch.matrix_exp(B[0]); c=torch.linalg.matrix_norm(B[0]@B[1]-B[1]@B[0]).item()**2
        gain.append(float(torch.linalg.svdvals(phi).max())); comm.append(c)
    return {'observed_commutator_mean':float(np.mean(comm)),'observed_cycle_gain_mean':float(np.mean(gain)),'n_models':len(models)}

def main():
    seeds=list(range(8)); grid=[{'lr':lr,'epochs':20} for lr in (.001,.003,.009)]
    # Equal union: baseline evaluates every learning rate considered by the idea.
    base_sweep=[]
    for cfg in grid: base_sweep.append(eval_side(baseline_factory,cfg,seeds[:4]))
    best=min(base_sweep,key=lambda r:r['mean']); best_cfg=best['config']
    base_full=eval_side(baseline_factory,best_cfg,seeds)
    base_block={'best_cfg':best_cfg,'sweep':[{'config':r['config'],'mean':r['mean'],'per_seed':r['per_seed']} for r in base_sweep],'full':{'per_seed':base_full['per_seed'],'mean':base_full['mean']}}
    idea_grid=[dict(best_cfg,comm_lambda=lam) for lam in (.0,.01,.05)]
    idea_runs=[eval_side(lambda c,s,d: idea_factory(c,s,d),cfg,seeds) for cfg in idea_grid]
    idea=min(idea_runs,key=lambda r:r['mean'])
    report=make_report('dynamics','rnn_small',base_block,{'config':idea['config'],'per_seed':idea['per_seed'],'mean':idea['mean']},extra={'mechanism_signature':{'prediction':'commutator regularization lowers trained ordered-cycle noncommutativity and cycle gain','observed':diagnostics(idea['models'],seeds),'baseline_observed':{'cycle_gain_mean':diagnostics(base_full['models'],seeds)['observed_cycle_gain_mean'],'observed_commutator_mean':diagnostics(base_full['models'],seeds)['observed_commutator_mean'],'input_sensitivity_mean':model_input_sensitivity(base_full['models'],seeds),'note':'Matched switched SSM baseline with lambda=0'},'idea_input_sensitivity_mean':model_input_sensitivity(idea['models'],seeds),'confirmed': diagnostics(idea['models'],seeds)['observed_commutator_mean'] < diagnostics(base_full['models'],seeds)['observed_commutator_mean']},'track_justification':'Dynamics is structurally matched to switched latent stability and Lyapunov contraction.'})
    report['idea_sweep']=[{'config':r['config'],'mean':r['mean'],'per_seed':r['per_seed']} for r in idea_runs]
    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=='__main__': main()