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

TRACK='sequence'; MODEL='transformer_tiny'; SEEDS=tuple(range(8))
# Small, equal-budget benchmark: all candidate lrs are evaluated by both methods.
NTR,NTE,EPOCHS,BATCH=400,400,5,128
LR_GRID=[1e-3,3e-3,1e-2]
LAMBDA=0.12; RHO=0.5; M=32

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 make_net(d):
    # Same base transformer as bench, with the sigmoid propensity encoder required
    # by the idea; both systems use this identical architecture.
    return SigmoidTransformer(d['input_shape'][0], d['out_dim'])

class SigmoidTransformer(nn.Module):
    def __init__(self, win, out_dim, dim=64):
        super().__init__()
        self.inp=nn.Linear(1,dim)
        self.pos=nn.Parameter(torch.zeros(1,win,dim)); nn.init.normal_(self.pos,std=.02)
        layer=nn.TransformerEncoderLayer(dim,nhead=2,dim_feedforward=128,batch_first=True,dropout=0.)
        self.enc=nn.TransformerEncoder(layer,2)
        self.head=nn.Linear(win*dim,out_dim)
    def forward(self,x):
        h=torch.sigmoid(self.inp(x.unsqueeze(-1))) + self.pos[:,:x.shape[1]]
        return self.head(self.enc(h).reshape(x.shape[0],-1))
    def propensity(self,x):
        return torch.sigmoid(self.inp(x.unsqueeze(-1))).mean(1)

def baseline_run(cfg, seed, keep=False):
    seed_all(seed); d=get_dataset(TRACK,seed,n_train=NTR,n_test=NTE)
    net=make_net(d)
    # train_model is the canonical baseline path; model architecture is shared.
    _, metric, _=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,log=lambda *_:None)
    return float(metric)

def local_run(cfg, seed, collect=False, device_override=None):
    seed_all(seed); d=get_dataset(TRACK,seed,n_train=NTR,n_test=NTE)
    net=make_net(d)
    # Adam trains the non-encoder parameters identically; the intervention is
    # exclusively the online fixed-projection update of inp.weight/bias.
    device=device_override or ('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        net=net.to(device); xtr,ytr=d['xtr'].to(device),d['ytr'].to(device)
        xte,yte=d['xte'].to(device),d['yte'].to(device)
        opt=torch.optim.Adam([p for n,p in net.named_parameters() if not n.startswith('inp.')],lr=cfg['lr'])
        rng=torch.Generator(device=device); rng.manual_seed(seed+991)
        A=torch.randn(M,64,device=device,generator=rng)/np.sqrt(M)
        before=net.inp.weight.detach().clone(); sign_num=sign_den=0.
        lossf=nn.M1Loss() if False else nn.MSELoss()
        for ep in range(EPOCHS):
            net.train(); perm=torch.randperm(len(xtr),device=device)
            for st in range(0,len(xtr),BATCH):
                ix=perm[st:st+BATCH]; pred=net(xtr[ix]); loss=lossf(pred,ytr[ix])
                opt.zero_grad(); loss.backward(); opt.step()
                # Ordered consecutive examples; vectorized local rank-one updates.
                if st == 0 or len(ix) > 1:
                    xa, xb = xtr[ix[:-1]], xtr[ix[1:]]
                    with torch.no_grad():
                        ha = torch.sigmoid(net.inp(xa.unsqueeze(-1))).mean(1)
                        hb = torch.sigmoid(net.inp(xb.unsqueeze(-1))).mean(1)
                        drive = (hb-ha) @ A.t() @ A
                        u = drive - LAMBDA*(ha-RHO)
                        delta = cfg['lr'] * ha*(1-ha)*u
                        net.inp.weight.add_(delta.t() @ xa.mean(1, keepdim=True) / max(1,len(ix)-1))
                        net.inp.bias.add_(delta.mean(0))
                        sign_num += float((delta*drive).sum())
                        sign_den += float(delta.abs().sum())
        net.eval()
        with torch.no_grad(): metric=float(((net(xte)-yte)**2).mean())
        if collect:
            with torch.no_grad():
                h=net.propensity(xte).flatten().cpu().numpy()
            return metric, {'mean_hidden':float(h.mean()),'saturated_fraction':float(np.mean((h<.05)|(h>.95))),
                            'projection_alignment':float(sign_num/(sign_den+1e-12)),'device':device}
        return metric
    except RuntimeError:
        # Explicit one-way GPU fallback; do not recurse if CPU itself fails.
        if device == 'cuda':
            torch.cuda.empty_cache()
            return local_run(cfg, seed, collect, device_override='cpu')
        raise

def main():
    grid=[{'lr':x} for x in LR_GRID]
    base=sweep_baseline(lambda c:(lambda s: baseline_run(c,s)),grid)
    best=base['best_cfg']
    idea_grid=grid
    idea_sweep=[{'cfg':c,'mean':evaluate(lambda s,c=c:local_run(c,s),SEEDS)['mean']} for c in idea_grid]
    best_idea=min(idea_sweep,key=lambda z:z['mean'])['cfg']
    idea=evaluate(lambda s:local_run(best_idea,s),SEEDS)
    sig=[]
    for s in SEEDS:
        _,st=local_run(best_idea,s,True); sig.append(st)
    sigmean={k:float(np.mean([z[k] for z in sig])) for k in sig[0] if k!='device'}
    sigmean['per_seed']=sig; sigmean['predicted_projection_alignment_positive']=True
    sigmean['confirmed']=bool(sigmean['projection_alignment']>0)
    base['idea_grid']=idea_sweep
    rep=make_report(TRACK,MODEL,base,idea,extra=sigmean)
    rep['protocol_notes']={'architecture_match':'identical sigmoid-input transformer; only inp training rule differs',
      'track_justification':'sequence forecast contains ordered multi-token temporal windows',
      'budget':{'n_train':NTR,'n_test':NTE,'epochs':EPOCHS,'batch':BATCH,'lr_union':LR_GRID},
      'custom_track':None}
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()