import sys, json, random, math
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, evaluate, sweep_baseline, make_report

SEEDS = tuple(range(8))
LR_GRID = [0.001, 0.003, 0.006]
EPOCHS = 8
IDEA_GRID = [(0.85, 0.5), (0.90, 0.5), (0.90, 1.0)]


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)


class HabituationTransformer(nn.Module):
    def __init__(self, win, out_dim, rho=0.9, beta=0.5, enabled=True):
        super().__init__()
        d = 64
        self.win, self.rho, self.beta, self.enabled = win, float(rho), float(beta), enabled
        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.0)
        self.enc = nn.TransformerEncoder(layer, 2)
        self.head = nn.Linear(win*d, out_dim)
        self.last_states = None
        self.last_gains = None

    def forward(self, x, capture=False):
        v = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]]
        if not self.enabled:
            h = self.enc(v)
            return self.head(h.reshape(x.shape[0], -1))
        # Scalar per-token stimulation, with stop-gradient to isolate habituation dynamics.
        a = torch.zeros(v.shape[0], device=v.device, dtype=v.dtype)
        hs, states, gains = [], [], []
        for t in range(v.shape[1]):
            vt = v[:, t]
            s = torch.sqrt((vt.detach() ** 2).mean(dim=-1) + 1e-8)
            a = self.rho * a + (1.0 - self.rho) * s
            g = 1.0 / (1.0 + self.beta * a)
            hs.append(vt * g[:, None])
            if capture:
                states.append(a.detach()); gains.append(g.detach())
        h = self.enc(torch.stack(hs, dim=1))
        if capture:
            self.last_states = torch.stack(states, dim=1)
            self.last_gains = torch.stack(gains, dim=1)
        return self.head(h.reshape(x.shape[0], -1))


def train_one(seed, cfg, idea, capture=False):
    seed_all(seed)
    ds = get_dataset('sequence', seed, n_train=400, n_test=200)
    win = ds['input_shape'][0]
    net = HabituationTransformer(win, ds['out_dim'], cfg.get('rho', .9),
                                 cfg.get('beta', .5), enabled=idea)
    net, metric, history = train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128,
                                       log=lambda *_: None)
    if not capture:
        return float(metric)
    net.eval()
    dev = next(net.parameters()).device
    xte = ds['xte'].to(dev)
    with torch.no_grad():
        pred = net(xte, capture=True)
    gains = net.last_gains.cpu().numpy() if net.last_gains is not None else np.ones((len(ds['xte']), win))
    states = net.last_states.cpu().numpy() if net.last_states is not None else np.zeros_like(gains)
    # Signature is measured on the trained network's actual test-window representations.
    observed_final = float(gains[:, -1].mean())
    observed_initial = float(gains[:, 0].mean())
    stim = float(states[:, -1].mean())
    predicted_final = 1.0 / (1.0 + cfg.get('beta', 0.0) * stim) if idea else 1.0
    # Fit temporal state pole using the first test sample's measured state trajectory.
    if idea and states.shape[1] > 2:
        y = states[0, 1:]; x = states[0, :-1]
        slope = float(np.dot(x, y) / max(np.dot(x, x), 1e-12))
    else:
        slope = 0.0
    return {'metric': float(metric), 'initial_gain': observed_initial,
            'final_gain': observed_final, 'mean_state': stim,
            'predicted_steady_gain': predicted_final, 'fitted_state_pole': slope,
            'rho': cfg.get('rho', 0.0), 'beta': cfg.get('beta', 0.0)}


def math_check():
    rows=[]
    for rho,beta in [(0.85,.5),(.9,.5),(.9,1.0)]:
        a=0.0
        for _ in range(300): a=rho*a+(1-rho)*1.7
        gain=1/(1+beta*a); pred=1/(1+beta*1.7)
        rec=[]; z=a
        for _ in range(100): z=rho*z; rec.append(z)
        half=next((i+1 for i,q in enumerate(rec) if q<=a/2),None)
        ph=math.ceil(math.log(.5)/math.log(rho))
        rows.append({'rho':rho,'steady_gain_abs_error':abs(gain-pred),
                     'half_observed':half,'half_predicted':ph})
    return {'rows':rows,'confirmed':all(r['steady_gain_abs_error']<1e-10 and r['half_observed']==r['half_predicted'] for r in rows)}


def main():
    # Baseline sweep includes every LR used by the idea side, satisfying union parity.
    baseline_grid=[{'lr':lr,'epochs':EPOCHS,'rho':.0,'beta':0.0} for lr in LR_GRID]
    base=sweep_baseline(lambda cfg: (lambda seed: train_one(seed,cfg,False)), baseline_grid, seeds=SEEDS)
    best=base['best_cfg']
    idea_cfgs=[]
    for lr in LR_GRID:
        rho,beta=IDEA_GRID[len(idea_cfgs)%len(IDEA_GRID)]
        idea_cfgs.append({'lr':lr,'epochs':EPOCHS,'rho':rho,'beta':beta})
    idea_runs=[]
    for cfg in idea_cfgs:
        result=evaluate(lambda seed,cfg=cfg: train_one(seed,cfg,True), seeds=SEEDS)
        idea_runs.append((cfg,result))
    idea_cfg, idea = min(idea_runs, key=lambda z:z[1]['mean'])
    sig=train_one(0, idea_cfg, True, capture=True)
    confirmed=(abs(sig['final_gain']-sig['predicted_steady_gain']) < .08 and
               abs(sig['fitted_state_pole']-idea_cfg['rho']) < .12)
    report=make_report('sequence','transformer_tiny',base,idea,extra={
        'track_rationale':'Sequence forecast contains multi-token temporal correlations and matches the proposed per-token fading-memory transformer gate.',
        'observed_best_cfg':idea_cfg,
        'idea_sweep':[{'cfg':c,'result':r} for c,r in idea_runs],
        'mechanism_signature':{'trained_model_seed':0, **sig, 'confirmed':bool(confirmed)},
        'math_check':math_check(),
        'protocol':{'seeds':list(SEEDS),'epochs':EPOCHS,'baseline_lr_grid':LR_GRID,'idea_lr_grid':LR_GRID}
    })
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()