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

SEEDS=tuple(range(8)); SWEEP_SEEDS=(0,1,2,3); EPOCHS=10; BATCH=128
LRS=[1e-3,3e-3,1e-2]

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def baseline_run(lr, seed, return_model=False):
    seed_all(seed); ds=get_dataset('dynamics', seed=seed, n_train=400, n_test=200)
    net=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    net, metric, hist=train_model(net, ds, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *a,**k:None)
    if return_model: return float(metric), net, ds
    return float(metric)

class PICopies(nn.Module):
    def __init__(self, hidden=64, copies=3, kp=0.5, ki=0.05):
        super().__init__(); self.copies=copies; self.hidden=hidden; self.kp=kp; self.ki=ki
        self.cells=nn.ModuleList([nn.GRUCell(3,hidden) for _ in range(copies)])
        self.head=nn.Linear(hidden,1)
        # deliberate replica-specific constant biases (heterogeneity)
        self.bias=nn.Parameter(torch.tensor([[0.03],[-0.02],[0.04]]).repeat(1,hidden), requires_grad=False)
    def forward(self,x, return_states=False):
        seq=x.view(x.shape[0],-1,3); b=x.shape[0]; dev=x.device
        hs=[torch.zeros(b,self.hidden,device=dev) for _ in range(self.copies)]
        z=[torch.zeros_like(hs[0]) for _ in range(self.copies)]
        states=[]
        for t in range(seq.shape[1]):
            raw=[self.cells[i](seq[:,t],hs[i])+self.bias[i] for i in range(self.copies)]
            mean=torch.stack(raw).mean(0)
            # Complete-graph zero-row-sum Laplacian output: mean - replica.
            p=[raw[i]-mean for i in range(self.copies)]
            mean_p=torch.stack(p).mean(0)
            z=[z[i]+self.ki*(raw[i]-mean) for i in range(self.copies)]
            hs=[raw[i]-self.kp*p[i]-z[i] for i in range(self.copies)]
            states.append(torch.stack(hs))
        hbar=torch.stack(hs).mean(0)
        out=self.head(hbar)
        if return_states: return out, torch.stack(states)
        return out

def idea_run(lr, seed, kp=0.5, ki=0.05, return_model=False):
    seed_all(seed); ds=get_dataset('dynamics', seed=seed, n_train=400, n_test=200)
    net=PICopies(kp=kp,ki=ki)
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net=net.to(dev); xtr,ytr=ds['xtr'].to(dev),ds['ytr'].to(dev); xte,yte=ds['xte'].to(dev),ds['yte'].to(dev)
        opt=torch.optim.Adam(net.parameters(),lr=lr); loss_fn=nn.MSELoss()
        for _ in range(EPOCHS):
            net.train(); perm=torch.randperm(len(xtr),device=dev)
            for ix in perm.split(BATCH):
                opt.zero_grad(); loss=loss_fn(net(xtr[ix]),ytr[ix]); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step()
        net.eval()
        with torch.no_grad(): metric=float(loss_fn(net(xte),yte).item())
    except Exception:
        net=net.cpu(); xtr,ytr=ds['xtr'],ds['ytr']; xte,yte=ds['xte'],ds['yte']
        opt=torch.optim.Adam(net.parameters(),lr=lr); loss_fn=nn.MSELoss()
        for _ in range(EPOCHS):
            for ix in torch.randperm(len(xtr)).split(BATCH):
                opt.zero_grad(); loss=loss_fn(net(xtr[ix]),ytr[ix]); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step()
        with torch.no_grad(): metric=float(loss_fn(net(xte),yte).item())
    if return_model:return metric,net,ds
    return metric

def signature():
    # Measured on trained benchmark models, not the analytic toy system.
    vals=[]
    for s in (0,1,2,3):
        m,net,ds=idea_run(3e-3,s,return_model=True)
        dev=next(net.parameters()).device
        with torch.no_grad():
            _,st=net(ds['xte'].to(dev),return_states=True)
            dis=float(st.std(dim=1).mean().cpu())
        vals.append(dis)
    obs=float(np.mean(vals))
    # PI's predicted qualitative signature is contraction of transverse replicas.
    return {'prediction':'integral coupling should yield small transverse hidden disagreement in trained heterogeneous copies',
            'observed_mean_pairwise_hidden_std':obs,'per_seed':vals,
            'threshold':0.20,'confirmed':bool(np.isfinite(obs) and obs<0.20)}

def main():
    grid=[{'lr':v} for v in LRS]
    base=sweep_baseline(lambda c: lambda s: baseline_run(c['lr'],s),grid,seeds=SWEEP_SEEDS)
    trials=[{'cfg':{'lr':lr},'result':evaluate(lambda s,lr=lr: idea_run(lr,s),SEEDS)} for lr in LRS]
    best=min(trials,key=lambda q:q['result']['mean'])
    rep=make_report('dynamics','rnn_small',base,best['result'],{
      'idea_config':best['cfg'],'idea_sweep':trials,
      'mechanism_signature':signature(),
      'protocol_note':'baseline sweep and idea sweep use identical learning-rate union; all task metrics are test MSE'})
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()