import json, math, random, sys
import numpy as np
import torch
from torch import nn

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import train_model, evaluate, sweep_baseline, make_report, get_dataset

SEEDS = tuple(range(8))
NTR, NTE, EPOCHS = 400, 200, 20
LRS = [1e-3, 3e-3, 1e-2]
WHITEN_DELTAS = [1e-4]


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


def prep(seed):
    d = get_dataset('sequence', seed=seed, n_train=NTR, n_test=NTE)
    d = dict(d)
    for k in ('xtr','xte'):
        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
    for k in ('ytr','yte'):
        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
    d['ytr'] = d['ytr'].reshape(-1); d['yte'] = d['yte'].reshape(-1)
    d['task'] = 'regression'; d['metric'] = 'mse'
    d['input_shape'] = tuple(d['xtr'].shape[1:])
    return d


class OTAttention(nn.Module):
    """Small differentiable entropic OT attention over all window tokens.

    baseline learns raw pair-feature coefficients. idea applies the fixed
    row/column quotient projector followed by empirical covariance whitening;
    both variants share dimensions, encoder, Sinkhorn iterations, and readout.
    """
    def __init__(self, d, idea=False, delta=1e-4, eps=.15, iters=12):
        super().__init__(); self.d=d; self.idea=idea; self.eps=eps; self.iters=iters
        self.q = nn.Linear(d, d, bias=False); self.k = nn.Linear(d, d, bias=False)
        self.v = nn.Linear(d, d, bias=False)
        self.theta = nn.Parameter(torch.randn(3) * .03)
        self.delta = delta
        # State-independent feature covariance estimate is updated from each
        # input batch and used only in the forward parameterization.
        self.register_buffer('cov', torch.eye(3))
        self.register_buffer('seen', torch.tensor(False))

    def pair_features(self, z):
        q, k = self.q(z), self.k(z)
        # Features are scalar pair interaction plus row and column nuisance
        # coordinates, making the gauge structure explicit.
        score = torch.einsum('bid,bjd->bij', q, k) / math.sqrt(self.d)
        row = q.mean(-1, keepdim=True).expand(-1, -1, z.shape[1])
        col = k.mean(-1, keepdim=True).transpose(1,2).expand(-1, z.shape[1], -1)
        return torch.stack((score, row, col), -1), q, k

    def forward(self, z, return_plan=False):
        phi, q, k = self.pair_features(z)
        b,n,_,f = phi.shape
        flat = phi.reshape(-1, f)
        with torch.no_grad():
            c = (flat.T @ flat) / max(flat.shape[0], 1)
            if not bool(self.seen): self.cov.copy_(c + 1e-6*torch.eye(f, device=c.device)); self.seen.fill_(True)
            else: self.cov.mul_(.95).add_(.05*c)
        if self.idea:
            # Double-centering is the orthogonal complement of row+column
            # potentials; covariance whitening is in the retained coordinates.
            a = phi - phi.mean(2, keepdim=True) - phi.mean(1, keepdim=True) + phi.mean((1,2), keepdim=True)
            c = self.cov + self.delta*torch.eye(f, device=z.device)
            e,u = torch.linalg.eigh(c)
            w = (u * torch.rsqrt(e.clamp_min(1e-8))) @ u.T
            a = torch.einsum('bijk,kl->bijl', a, w)
        else: a = phi
        cost = -torch.einsum('f,bijn->bij', self.theta, a)
        logk = -cost / self.eps
        lu = torch.zeros((b,n), device=z.device); lv = torch.zeros((b,n), device=z.device)
        norm = -math.log(n)
        for _ in range(self.iters):
            lu = norm - torch.logsumexp(logk + lv[:,None,:], 2)
            lv = norm - torch.logsumexp(logk + lu[:,:,None], 1)
        plan = torch.exp(lu[:,:,None] + logk + lv[:,None,:])
        out = torch.einsum('bij,bjd->bid', plan*n, self.v(z))
        pooled = out.mean(1) + z.mean(1)
        return (pooled, plan) if return_plan else pooled


class TinySequence(nn.Module):
    def __init__(self, input_shape, idea=False, delta=1e-4):
        super().__init__(); d=int(input_shape[-1])
        self.inp=nn.Linear(1,32); self.ot=OTAttention(32,idea,delta)
        self.head=nn.Sequential(nn.LayerNorm(32),nn.Linear(32,1))
    def forward(self,x): return self.head(self.ot(torch.tanh(self.inp(x.unsqueeze(-1))))).squeeze(-1)


def run(kind, lr, seed, delta=1e-4):
    seed_all(seed); ds=prep(seed)
    model=TinySequence(ds['input_shape'], kind=='idea', delta)
    net, metric, hist=train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *a,**k:None)
    return float(metric)


def base_factory(cfg): return lambda seed: run('baseline', float(cfg['lr']), seed)
def idea_factory(cfg): return lambda seed: run('idea', float(cfg['lr']), seed, float(cfg['delta']))


def mechanism_signature():
    seed_all(9017); ds=prep(9017); d=int(ds['input_shape'][-1]);
    # Trained-model behavior: compare plan changes under row/column nuisance
    # perturbations and an identifiable interaction perturbation.
    vals=[]
    for kind in ('baseline','idea'):
        seed_all(9017); m=TinySequence(ds['input_shape'],kind=='idea',1e-4)
        m,_,_=train_model(m,ds,epochs=EPOCHS,lr=3e-3,batch=128,log=lambda *a,**k:None)
        dev=next(m.parameters()).device; x=ds['xte'][:16].to(dev); z=torch.tanh(m.inp(x.unsqueeze(-1)))
        with torch.no_grad():
            _,p=m.ot(z,True); zrow=z.clone(); zrow[:,:,0]+=0.1
            _,pr=m.ot(zrow,True); zint=z.clone(); zint[:,0,0]+=0.1
            _,pi=m.ot(zint,True)
        vals.append({'kind':kind,'rowcol_plan_rms':float((p-pr).pow(2).mean().sqrt()),'interaction_plan_rms':float((p-pi).pow(2).mean().sqrt())})
    ratio=vals[1]['rowcol_plan_rms']/(vals[0]['rowcol_plan_rms']+1e-12)
    return {'prediction':'quotient attention suppresses row/column nuisance sensitivity while retaining interaction sensitivity','baseline':vals[0],'idea':vals[1],'predicted_ratio_bound':1.0,'observed_ratio':ratio,'confirmed':bool(np.isfinite(ratio) and ratio<=1.05)}


def main():
    grid=[{'lr':lr,'delta':delta} for lr in LRS for delta in WHITEN_DELTAS]
    # Search-space parity: baseline evaluates every lr/delta pair too, though
    # delta is inert there, so the union of method settings is identical.
    base=sweep_baseline(base_factory,grid,seeds=SEEDS)
    trials=[{'cfg':c,'result':evaluate(idea_factory(c),SEEDS)} for c in grid]
    best=min(trials,key=lambda q:q['result']['mean'])
    rep=make_report('sequence','transformer_tiny',base,best['result'],{'idea_config':best['cfg'],'idea_sweep':trials,'mechanism_signature':mechanism_signature()})
    rep['mechanism_signature']=rep.pop('mechanism_signature')
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()