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, train_model, make_report, sweep_baseline, evaluate

SEED = 2486
EPOCHS = 8
NTR, NTE = 1200, 400
LRS = [1e-3, 3e-3, 6e-3]
A = 8
D = 64

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

def gale_check():
    rng = np.random.default_rng(SEED)
    rows=[]
    for a,b in [(3,7),(8,32),(12,32)]:
        X=rng.normal(size=(a,b))
        U,S,Vh=np.linalg.svd(X.T, full_matrices=True)
        Y=U[:,a:].T
        rows.append(float(np.linalg.norm(Y@X.T)))
    return {'sizes': [[3,7],[8,32],[12,32]], 'residuals': rows,
            'max_residual': max(rows), 'passed': bool(max(rows)<1e-10)}

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

class GaleTransformer(BaseTransformer):
    def __init__(self, win=32, d=64, a=A, depth=2):
        super().__init__(win,d,depth); self.a=a; self.k=win-a
        self.wx=nn.Linear(d,a,bias=False)
        self.wy=nn.Linear(d,d,bias=False)
        self.dual_out=nn.Linear(d,d,bias=False)
        self.gate=nn.Parameter(torch.tensor(0.0))
    def forward_features(self,x, return_stats=False):
        h=self.encode(x)
        X=self.wx(h).transpose(1,2) # B,a,T; X^T is B,T,a
        # full U of X^T has T columns; last T-a rows of Y span null(X)
        U,S,Vh=torch.linalg.svd(X.transpose(1,2), full_matrices=True)
        Y=U[:,:,self.a:].transpose(1,2) # B,k,T
        dual=Y @ self.wy(h) # B,k,d
        correction=Y.transpose(1,2) @ self.dual_out(dual) # B,T,d
        z=h + torch.sigmoid(self.gate)*correction
        if return_stats:
            # Behavior signature: overlap of learned primary and dual summaries.
            p=X.mean(dim=(1,2)); q=dual.mean(dim=(1,2))
            pc=p-p.mean(); qc=q-q.mean()
            corr=(pc*qc).sum()/(torch.sqrt((pc.square().sum()+1e-12)*(qc.square().sum()+1e-12)))
            residual=(Y @ X.transpose(1,2)).norm(dim=(1,2)).mean()
            return z, {'abs_summary_corr': float(corr.abs().detach()),
                       'null_residual': float(residual.detach()),
                       'gate': float(torch.sigmoid(self.gate).detach())}
        return z
    def forward(self,x): return self.head(self.forward_features(x).reshape(x.shape[0],-1))

def dataset(seed): return get_dataset('sequence', seed, n_train=NTR, n_test=NTE)

def run_one(kind, cfg, seed, want_stats=False):
    seed_all(seed); ds=dataset(seed)
    net=BaseTransformer() if kind=='baseline' else GaleTransformer()
    net, metric, hist=train_model(net, ds, epochs=EPOCHS, lr=float(cfg['lr']), batch=128)
    if net is None: raise RuntimeError('training failed')
    st=None
    if want_stats:
        net.eval()
        with torch.no_grad():
            try: _,st=net.forward_features(ds['xte'], True)
            except RuntimeError:
                # move to CPU if a shared CUDA slot failed during signature collection
                net=net.cpu(); _,st=net.forward_features(ds['xte'].cpu(), True)
    return float(metric), st

def main():
    check=gale_check()
    assert check['passed'], check
    grid=[{'lr':v} for v in LRS]
    def mkbase(cfg): return lambda s: run_one('baseline',cfg,s)[0]
    base=sweep_baseline(mkbase,grid)
    # Same union of learning rates on idea side; best is selected on the same four sweep seeds.
    idea_sweep=[]
    for cfg in grid:
        r=evaluate(lambda s: run_one('idea',cfg,s)[0], seeds=(0,1,2,3))
        idea_sweep.append({'cfg':cfg,'mean':r['mean']})
    best_cfg=min(idea_sweep,key=lambda r:r['mean'])['cfg']
    idea=evaluate(lambda s: run_one('idea',best_cfg,s)[0])
    stats=[]
    for s in range(8): stats.append(run_one('idea',best_cfg,s,True)[1])
    sig={'prediction': 'dual summary should have low linear overlap with primary summary (abs corr < 0.2), while Gale residual is numerical roundoff',
         'predicted_abs_summary_corr_max':0.2,
         'observed_abs_summary_corr_mean':float(np.mean([x['abs_summary_corr'] for x in stats])),
         'observed_abs_summary_corr_per_seed':[x['abs_summary_corr'] for x in stats],
         'observed_null_residual_mean':float(np.mean([x['null_residual'] for x in stats])),
         'observed_gate_mean':float(np.mean([x['gate'] for x in stats])),
         'confirmed':bool(float(np.mean([x['abs_summary_corr'] for x in stats]))<0.2)}
    report=make_report('sequence','transformer_tiny',
        {'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},idea,sig)
    report['idea_sweep']=idea_sweep
    report['core_math_check']=check
    report['setup']={'epochs':EPOCHS,'n_train':NTR,'n_test':NTE,'a':A,'window':32,
                     'track_justification':'sequence is structurally matched because Gale mixing operates across multi-token windows'}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()