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

ROOT = Path(__file__).parent
SEEDS = tuple(range(8))
# Union is shared: baseline is evaluated at every idea candidate too.
GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
EPOCHS, BATCH = 18, 128

class IMMMonitor:
    """Small diagonal IMM over delays 0..D for observed training dynamics."""
    def __init__(self, dim=5, D=2, q=0.03, r=0.08, persistence=.96, threshold=.08):
        self.D, self.m, self.dim = D, D+1, dim
        self.q, self.r, self.threshold = q, r, threshold
        self.T = np.full((self.m,self.m), (1-persistence)/D)
        np.fill_diagonal(self.T, persistence)
        self.pi = np.ones(self.m)/self.m
        self.hist = []
        self.alarm_run = 0
        self.last = None

    def step(self, z):
        z = np.asarray(z, dtype=float)
        self.hist.append(z.copy())
        # Calibration adapts observation noise from clean one-step residuals.
        if len(self.hist) > 1:
            e = self.hist[-1] - self.hist[-2]
            self.r = max(.015, .90*self.r + .10*float(np.mean(e*e)))
        c = self.T.T @ self.pi
        preds = []
        for i in range(self.m):
            idx = len(self.hist)-1-i
            preds.append(self.hist[max(0, idx)])
        preds = np.asarray(preds)
        ll = -.5*np.sum((z[None,:]-preds)**2, axis=1)/self.r - .5*self.dim*math.log(2*math.pi*self.r)
        u = c*np.exp(ll-np.max(ll)); self.pi = u/max(u.sum(), 1e-300)
        if self.pi[0] <= self.threshold: self.alarm_run += 1
        else: self.alarm_run = 0
        alarm = self.alarm_run >= 3
        self.last = {'posterior': self.pi.copy(), 'alarm': alarm, 'innovation': float(np.linalg.norm(z-preds[0]))}
        return alarm

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 vec_norm(model):
    return math.sqrt(sum(float((p.detach()**2).sum()) for p in model.parameters()))

def train_imm(ds, *, epochs=EPOCHS, lr=3e-3, batch=BATCH, weight_decay=0.0):
    seed_all(int(ds.get('_seed', 0)))
    # Same architecture and optimizer as bench.train_model; only loop intervention differs.
    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
    lossf = nn.MSELoss()
    opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net.to(device); x, y = ds['xtr'].to(device), ds['ytr'].to(device)
        mon = IMMMonitor(); hist=[]; reduced=False; alarms=0; post=[]; innovations=[]
        for ep in range(epochs):
            net.train(); perm=torch.randperm(len(x), device=device); total=0.
            for j in range(0,len(x),batch):
                idx=perm[j:j+batch]; loss=lossf(net(x[idx]),y[idx])
                opt.zero_grad(); loss.backward()
                gn=math.sqrt(sum(float((p.grad.detach()**2).sum()) for p in net.parameters() if p.grad is not None))
                pn=vec_norm(net); un=lr*gn
                z=np.log1p(np.array([float(loss), max(float(loss),0.), gn, pn, un]))
                alarm=mon.step(z); alarms += int(alarm)
                if mon.last: post.append(mon.last['posterior'].copy()); innovations.append(mon.last['innovation'])
                if alarm and not reduced:
                    for group in opt.param_groups: group['lr']=lr/4.
                    reduced=True
                opt.step(); total += float(loss)*len(idx)
            hist.append(total/len(x))
        net.eval()
        with torch.no_grad(): metric=float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
        ds['_signature']={'alarms':alarms,'reduced':reduced,'posterior_final':post[-1].tolist() if post else [],'innovation_mean':float(np.mean(innovations)) if innovations else 0.}
        return net, metric, hist
    except RuntimeError:
        # CPU fallback, matching bench's robustness.
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        return train_imm_cpu(ds, epochs=epochs, lr=lr, batch=batch, weight_decay=weight_decay)

def train_imm_cpu(ds, *, epochs, lr, batch, weight_decay):
    old=torch.cuda.is_available; torch.cuda.is_available=lambda:False
    try: return train_imm(ds, epochs=epochs, lr=lr, batch=batch, weight_decay=weight_decay)
    finally: torch.cuda.is_available=old

def baseline_fn(cfg):
    def run(seed):
        seed_all(seed); ds=get_dataset('tabular', seed, n_train=400, n_test=400); ds['_seed']=seed
        net, metric, hist=train_model(make_model('mlp_tiny',ds['input_shape'],ds['out_dim']),ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH)
        return metric
    return run

def idea_fn(cfg):
    def run(seed):
        ds=get_dataset('tabular', seed, n_train=400, n_test=400); ds['_seed']=seed
        return train_imm(ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH)[1]
    return run

def main():
    base=sweep_baseline(baseline_fn, GRID)
    idea_cfgs=[base['best_cfg']]+[c for c in GRID if c != base['best_cfg']]
    # evaluate all three shared-grid configurations, select best on full paired seeds.
    idea_runs=[]
    for cfg in idea_cfgs: idea_runs.append((cfg,evaluate(idea_fn(cfg), seeds=SEEDS)))
    best_cfg, idea=min(idea_runs,key=lambda x:x[1]['mean'])
    sig={'prediction':'persistent stale-feedback should lower no-delay posterior and trigger intervention','observed':'trained tabular MLP monitor posterior/alarm telemetry across paired runs','confirmed':False,'note':'No injected delay is available in the canonical bench; signature is behavioral telemetry, not toy arithmetic.'}
    rep=make_report('tabular','mlp_tiny',base,idea,extra=sig)
    rep['idea_sweep']=[{'cfg':c,'mean':r['mean']} for c,r in idea_runs]; rep['selected_idea_cfg']=best_cfg
    (ROOT/'bench_report.json').write_text(json.dumps(rep,indent=2,allow_nan=False))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()