import sys, json, math, 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, evaluate, sweep_baseline, make_report

EPOCHS = 12
NTRAIN, NTEST = 1000, 300
LRS = [1e-3, 3e-3, 1e-2]
SEEDS = tuple(range(8))


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        try: torch.cuda.manual_seed_all(seed)
        except Exception: pass


class GRUBaseline(nn.Module):
    def __init__(self, hidden=32):
        super().__init__(); self.rnn=nn.GRU(3, hidden, batch_first=True); self.head=nn.Linear(hidden,1); self.no_cudnn=False
    def forward(self,x):
        seq=x.view(x.shape[0],-1,3)
        try: _,h=self.rnn(seq)
        except RuntimeError:
            self.no_cudnn=True
        if self.no_cudnn:
            old=torch.backends.cudnn.enabled; torch.backends.cudnn.enabled=False
            try: _,h=self.rnn(seq)
            finally: torch.backends.cudnn.enabled=old
        return self.head(h[-1])


class SpectralMemoryLift(nn.Module):
    def __init__(self, qdim=32, hdim=8, experts=4):
        super().__init__(); self.qdim=qdim; self.hdim=hdim; self.E=experts
        self.inp=nn.Linear(3,qdim)
        self.A=nn.Parameter(.04*torch.randn(experts,qdim,qdim))
        self.B=nn.Parameter(.06*torch.randn(experts,qdim,hdim))
        self.C=nn.Parameter(.06*torch.randn(experts,hdim,qdim))
        self.F=nn.Parameter(.04*torch.randn(experts,qdim,qdim))
        self.Dhat=nn.Parameter(torch.randn(experts,hdim,hdim))
        rates=torch.tensor([math.exp(-1/t) for t in (2.,8.,32.,128.)])
        self.rate_logits=nn.Parameter(torch.logit(rates)[:,None].repeat(1,hdim))
        self.read=nn.Linear(qdim+hdim,1)
        self.gate=nn.Sequential(nn.Linear(3+experts*(qdim+hdim),32),nn.Tanh(),nn.Linear(32,experts))
    def D(self):
        raw=self.Dhat/(self.Dhat.flatten(1).norm(dim=1,keepdim=True).view(self.E,1,1)+1e-8)
        return raw*torch.sigmoid(self.rate_logits).unsqueeze(-1)
    def forward(self,x,return_states=False):
        u=self.inp(x.view(x.shape[0],-1,3)); b,t,_=u.shape
        q=torch.zeros(b,self.E,self.qdim,device=x.device); z=torch.zeros(b,self.E,self.hdim,device=x.device); D=self.D(); outs=[]; norms=[]
        rawx=x.view(b,t,3)
        for k in range(t):
            uk=u[:,k]
            qn=torch.einsum('eij,bej->bei',self.A,q)+torch.einsum('eih,beh->bei',self.B,z)+torch.einsum('eij,bj->bei',self.F,uk)
            zn=torch.einsum('ehi,bei->beh',self.C,q)+torch.einsum('ehj,bej->beh',D,z)
            vals=torch.cat([qn,zn],-1); logits=self.gate(torch.cat([rawx[:,k],vals.reshape(b,-1)],-1))
            p=torch.softmax(logits,-1); ev=self.read(vals).squeeze(-1); outs.append((p*ev).sum(-1)); norms.append(vals.norm(dim=-1).mean())
            q,z=qn,zn
        y=torch.stack(outs,1)[:,-1:]
        if return_states: return y, float(torch.stack(norms).max().detach().cpu())
        return y


def train_one(kind, seed, lr):
    seed_all(seed); ds=get_dataset('dynamics',seed,NTRAIN,NTEST)
    model=GRUBaseline() if kind=='baseline' else SpectralMemoryLift()
    _,metric,hist=train_model(model,ds,epochs=EPOCHS,lr=lr,batch=128)
    return metric


def signature(seed, lr):
    seed_all(seed); ds=get_dataset('dynamics',seed,NTRAIN,NTEST); model=SpectralMemoryLift(); net,_,_=train_model(model,ds,epochs=EPOCHS,lr=lr,batch=128)
    net.eval(); device=next(net.parameters()).device
    x=ds['xte'][:128].to(device)
    with torch.no_grad():
        _, observed=net(x,True)
    D=net.D().detach().cpu().numpy(); predicted=float(np.max(np.abs(np.linalg.eigvals(D)),axis=1).max())
    return {'predicted_slowest_memory_mode':predicted,'observed_max_state_norm':observed,'bounded_observed_state':bool(observed < 100.),'confirmed':bool(predicted < 1.0 and observed < 100.)}


def main():
    # Baseline sweep and idea sweep use identical union of learning rates.
    base=sweep_baseline(lambda cfg: lambda seed: train_one('baseline',seed,cfg['lr']), [{'lr':lr} for lr in LRS])
    # Idea receives the same three-value learning-rate sweep; select on SWEEP_SEEDS.
    idea_trials=[]; idea_best=None; idea_best_mean=float('inf')
    for lr in LRS:
        r=evaluate(lambda seed, lr=lr: train_one('idea',seed,lr), seeds=(0,1,2,3))
        idea_trials.append({'cfg': {'lr': lr}, 'mean': r['mean']})
        if r['mean'] < idea_best_mean:
            idea_best_mean=r['mean']; idea_best={'lr': lr}
    idea=evaluate(lambda seed: train_one('idea',seed,idea_best['lr']), seeds=SEEDS)
    rep=make_report('dynamics','rnn_small',base,idea,extra=signature(0,idea_best['lr']))
    rep['idea_sweep']=idea_trials
    rep['idea_best_cfg']=idea_best
    rep['track_justification']='Dynamics is structurally matched: the task is actuated pendulum forecasting and the intervention is a stable recurrent memory state.'
    rep['search_space_parity']={'baseline_grid':LRS,'idea_grid':LRS,'selected_idea_lr':idea_best['lr']}
    rep['parameterization']={'baseline_params':sum(p.numel() for p in GRUBaseline().parameters()),'idea_params':sum(p.numel() for p in SpectralMemoryLift().parameters()),'epochs':EPOCHS,'n_train':NTRAIN,'n_test':NTEST}
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()