from __future__ import annotations
import json, math, random, time
from pathlib import Path
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, evaluate, sweep_baseline, make_report

SEED = 2820
NTR, NTE, EPOCHS, BATCH = 400, 200, 6, 128
D, DEPTH, HEADS = 32, 1, 2
A, B, SIGMA_POS, SIGMA_NEG = 1.0, 0.8, 0.5, 2.0


def set_seed(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


def sanity_check():
    rng = np.random.default_rng(SEED)
    u = rng.uniform(-3, 3, 300)
    truth = A*np.exp(-0.5*(SIGMA_POS*u)**2) - B*np.exp(-0.5*(SIGMA_NEG*u)**2)
    rows = []
    for m in (16, 32, 64, 128, 256):
        es, eu = [], []
        for t in range(30):
            rr = np.random.default_rng(SEED + 1000*m + t)
            neg = rr.random(m) < B/(A+B)
            om = np.where(neg, rr.normal(0, SIGMA_NEG, m), rr.normal(0, SIGMA_POS, m))
            c = np.cos(u[:, None]*om)
            signed = (A+B)*np.mean(c*np.where(neg, -1., 1.), axis=1)
            unsigned = (A+B)*np.mean(c, axis=1)
            es.append(np.mean((signed-truth)**2)); eu.append(np.mean((unsigned-truth)**2))
        rows.append((m, float(np.sqrt(np.mean(es))), float(np.sqrt(np.mean(eu)))))
    slope = float(np.polyfit(np.log([r[0] for r in rows]), np.log([r[1] for r in rows]), 1)[0])
    positive_floor = float(np.sqrt(np.mean((A*np.exp(-.5*(SIGMA_POS*u)**2) + B*np.exp(-.5*(SIGMA_NEG*u)**2)-truth)**2)))
    return {"rows": [{"M":m,"signed_rmse":s,"positive_rmse":p} for m,s,p in rows],
            "observed_slope": slope, "predicted_slope": -.5,
            "positive_predicted_floor": positive_floor,
            "passed": bool(slope < -.25 and rows[-1][1] < rows[0][1] and rows[-1][2] > positive_floor*.95)}


class SoftmaxBlock(nn.Module):
    def __init__(self, d, heads, temperature=1.0):
        super().__init__(); self.temperature = temperature
        self.attn = nn.MultiheadAttention(d, heads, dropout=0., batch_first=True)
        self.n1 = nn.LayerNorm(d); self.ff = nn.Sequential(nn.Linear(d, 64), nn.GELU(), nn.Linear(64,d)); self.n2=nn.LayerNorm(d)
    def forward(self, x):
        h = self.n1(x)
        # Apply the swept temperature to the standard attention logits.
        w = self.attn.in_proj_weight
        b = self.attn.in_proj_bias
        q = torch.nn.functional.linear(h, w[:self.attn.embed_dim], b[:self.attn.embed_dim])
        k = torch.nn.functional.linear(h, w[self.attn.embed_dim:2*self.attn.embed_dim], b[self.attn.embed_dim:2*self.attn.embed_dim])
        v = torch.nn.functional.linear(h, w[2*self.attn.embed_dim:], b[2*self.attn.embed_dim:])
        q = q.view(x.shape[0], x.shape[1], self.attn.num_heads, -1).transpose(1,2)
        k = k.view(x.shape[0], x.shape[1], self.attn.num_heads, -1).transpose(1,2)
        v = v.view(x.shape[0], x.shape[1], self.attn.num_heads, -1).transpose(1,2)
        logits = torch.matmul(q, k.transpose(-2,-1)) / (math.sqrt(q.shape[-1]) * self.temperature)
        y = torch.matmul(torch.softmax(logits, dim=-1), v).transpose(1,2).reshape_as(h)
        y = self.attn.out_proj(y)
        return self.n2(x+y + self.ff(self.n2(x+y)))


class SignedBlock(nn.Module):
    def __init__(self, d, m, seed):
        super().__init__(); self.d=d; self.m=m; self.C=A+B
        rr=np.random.default_rng(seed)
        neg=rr.random(m)<B/(A+B)
        om=np.where(neg, rr.normal(0,SIGMA_NEG,(m,d)), rr.normal(0,SIGMA_POS,(m,d))).astype(np.float32)
        self.register_buffer('omega', torch.from_numpy(om)); self.register_buffer('signs', torch.from_numpy(np.where(neg,-1.,1.).astype(np.float32)))
        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.o=nn.Linear(d,d)
        self.n1=nn.LayerNorm(d); self.ff=nn.Sequential(nn.Linear(d,64),nn.GELU(),nn.Linear(64,d)); self.n2=nn.LayerNorm(d)
    def features(self, x, m=None):
        om=self.omega[:m] if m else self.omega
        a=x @ om.T
        return torch.cat((torch.cos(a), torch.sin(a)), dim=-1)
    def forward(self,x):
        h=self.n1(x); q=self.features(self.q(h)); k=self.features(self.k(h)); v=self.v(h)
        s=torch.repeat_interleave(self.signs,2); scale=self.C/self.m
        # Associative factorization: Zq D (Zk^T V), with positive feature normalization.
        kv=torch.einsum('bnd,bnr->bdr', k*s, v)
        out=scale*torch.einsum('bnd,bdr->bnr',q,kv)
        den=scale*torch.einsum('bnd,bd->bn', q, k.sum(dim=1)).unsqueeze(-1)
        # Positive denominator uses |spectral signs|, avoiding signed cancellation.
        den=den.abs().clamp_min(1e-3)
        y=self.o(out/den)
        z=x+y; return self.n2(z+self.ff(z))


class Net(nn.Module):
    def __init__(self, m=None, temp=1., seed=0):
        super().__init__(); self.inp=nn.Linear(1,D); self.pos=nn.Parameter(torch.zeros(1,32,D)); nn.init.normal_(self.pos,std=.02)
        self.block=SignedBlock(D,m,seed+91) if m else SoftmaxBlock(D,HEADS,temp)
        self.head=nn.Linear(32*D,1)
    def forward(self,x):
        h=self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]
        return self.head(self.block(h).reshape(x.shape[0],-1))


def run(cfg, idea, seed, capture=False):
    set_seed(seed); ds=get_dataset('sequence',seed,n_train=NTR,n_test=NTE)
    net=Net(m=cfg['m'], seed=seed) if idea else Net(temp=cfg['temperature'])
    net, metric, hist=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,log=lambda *_:None)
    if capture: captured[(idea,seed)] = net
    return float(metric)


def signature():
    net=captured.get((True,0));
    if net is None: return {"confirmed":False,"reason":"model capture failed"}
    net.eval()
    dev=next(net.parameters()).device
    x=get_dataset('sequence',0,n_train=NTR,n_test=32)['xte'].to(dev)
    with torch.no_grad():
        h=net.inp(x.unsqueeze(-1))+net.pos[:,:32]
        q=net.block.q(net.block.n1(h)); k=net.block.k(net.block.n1(h));
        q=q.reshape(-1,D); k=k.reshape(-1,D); u=q[:128]-k[:128]
        exact=A*torch.exp(-.5*(SIGMA_POS*u).pow(2).sum(1))-B*torch.exp(-.5*(SIGMA_NEG*u).pow(2).sum(1))
        M=net.block.m; vals=[]
        for mm in sorted(set((max(4,M//4),max(8,M//2),M))):
            zq=net.block.features(q[:128],mm); zk=net.block.features(k[:128],mm); s=torch.repeat_interleave(net.block.signs[:mm],2)
            est=(A+B)/mm*(zq*(zk*s)).sum(1); vals.append((mm,float(torch.sqrt(torch.mean((est-exact)**2)))))
        slope=float(np.polyfit(np.log([a for a,b in vals]),np.log([b for a,b in vals]),1)[0]) if len(vals)>1 else float('nan')
        unsigned=(A+B)/M*(net.block.features(q[:128],M)*(net.block.features(k[:128],M))).sum(1)
        ur=float(torch.sqrt(torch.mean((unsigned-exact)**2))); sr=vals[-1][1]
    return {"trained_model":True,"M":M,"observed_rmse_by_M":[{"M":a,"rmse":b} for a,b in vals],"observed_slope":slope,"predicted_slope":-.5,"signed_rmse":sr,"unsigned_rmse":ur,"confirmed":bool(sr<ur and slope < -.15)}


if __name__=='__main__':
    print(json.dumps({"sanity_check":sanity_check()},indent=2))
    captured={}
    # Union of all learning rates is evaluated by baseline; temperature is the baseline's method knob.
    lrs=[1e-3,3e-3,6e-3]; temps=[0.7,1.0]
    grid=[{"lr":lr,"temperature":t} for lr in lrs for t in temps]
    base=sweep_baseline(lambda cfg: lambda s: run(cfg,False,s,capture=False),grid)
    best_lr=base['best_cfg']['lr']
    # Three-point idea sweep: selected baseline lr plus two nearby rates; all are in baseline grid.
    idea_cfgs=[{"lr":lr,"m":32} for lr in (1e-3,3e-3,6e-3)]
    idea_all=[]
    for cfg in idea_cfgs:
        r=evaluate(lambda s,cfg=cfg: run(cfg,True,s,capture=False))
        idea_all.append({"cfg":cfg,"result":r})
    best=min(idea_all,key=lambda z:z['result']['mean']); bestcfg=best['cfg']
    idea=evaluate(lambda s: run(bestcfg,True,s,capture=(s==0)))
    extra=signature()
    report=make_report('sequence','transformer_tiny',base,idea,extra)
    report['idea_sweep']=idea_all; report['sanity_check']=sanity_check(); report['protocol']={"seeds":list(range(8)),"n_train":NTR,"n_test":NTE,"epochs":EPOCHS,"baseline_grid":grid,"idea_grid":idea_cfgs}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))