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

OUT = Path('bench_report.json')
SEEDS = tuple(range(8))
# Shared union: baseline is evaluated at every lr used by the idea.
GRID = [
    {'lr': 0.001, 'weight_decay': 0.0},
    {'lr': 0.003, 'weight_decay': 0.0},
    {'lr': 0.006, 'weight_decay': 0.0},
]
EPOCHS = 18
BATCH = 128


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 c_transition(k):
    k = max(float(k), 1.0)
    return math.sqrt(((k+1)+math.sqrt(max(0., (k+1)**2-4)))/2)


def calibrated_t(k, delta):
    k = max(float(k), 1.0)
    raw = math.sqrt(1 + math.sqrt(max(0., (k-1)*(1/delta-1))))
    return max(raw, c_transition(k))


def math_check():
    max_inv = 0.; max_poly = 0.; min_margin = 1e9
    for k in [1.01, 1.1, 2., 3., 10., 100.]:
        c = c_transition(k)
        max_poly = max(max_poly, abs(c**4-(k+1)*c**2+1))
        for d in [.001, .01, .05, .2]:
            raw = math.sqrt(1 + math.sqrt((k-1)*(1/d-1)))
            if raw >= c:
                tail = (k-1)/((raw*raw-1)**2+k-1)
                max_inv = max(max_inv, abs(tail-d))
            min_margin = min(min_margin, calibrated_t(k,d)-c)
    return {'max_inverse_error': max_inv, 'max_transition_polynomial_error': max_poly,
            'min_projected_tail_margin': min_margin, 'passed': max_inv < 1e-12 and max_poly < 1e-10}


def run(seed, cfg, idea=False, collect=False):
    seed_all(seed)
    ds = get_dataset('tabular', seed, n_train=400, n_test=400)
    # Explicitly use the matched canonical architecture.
    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net = net.to(device)
    except Exception:
        device = 'cpu'; net = net.to(device)
    opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg.get('weight_decay',0.))
    x, y = ds['xtr'].to(device), ds['ytr'].to(device)
    lossf = nn.MSELoss(); n = len(x)
    beta=.95; delta=.02; safety=1.5; cap=1000.
    stats={'raw_spike':0, 'clipped_entries':0, 'entries':0, 'tail_sum':0., 'k_sum':0., 'tau_sum':0., 'steps':0}
    for ep in range(EPOCHS):
        perm = torch.randperm(n, device=device)
        for start in range(0,n,BATCH):
            ix=perm[start:start+BATCH]; opt.zero_grad(set_to_none=True)
            loss=lossf(net(x[ix]), y[ix]); loss.backward()
            for p in net.parameters():
                if p.grad is None: continue
                g=p.grad
                if not idea:
                    stats['raw_spike'] += int(torch.isfinite(g).all() and g.detach().abs().max().item() > 5.0)
                    continue
                z=g.detach(); m=z.mean(); centered=z-m
                v_batch=(centered*centered).mean(); q_batch=(centered**4).mean()
                # Per-tensor EMA moments; current-batch central moments avoid lag bias.
                if not hasattr(p, '_kc_state'):
                    p._kc_state=[m.item(), max(v_batch.item(),1e-12), max(q_batch.item(),1e-12)]
                st=p._kc_state
                st[0]=beta*st[0]+(1-beta)*m.item(); st[1]=beta*st[1]+(1-beta)*v_batch.item(); st[2]=beta*st[2]+(1-beta)*q_batch.item()
                k=min(cap,max(1., safety*st[2]/(st[1]**2+1e-12))); t=calibrated_t(k,delta); tau=t*math.sqrt(st[1]+1e-12)
                stats['raw_spike'] += int(torch.isfinite(g).all() and z.abs().max().item()>5.0)
                stats['clipped_entries'] += int((centered.abs()>tau).sum().item()); stats['entries'] += z.numel()
                stats['tail_sum'] += float((centered.abs()>tau).float().mean().item()); stats['k_sum'] += k; stats['tau_sum'] += tau
                g.copy_(m + centered.clamp(-tau,tau))
            opt.step(); stats['steps'] += 1
    with torch.no_grad(): metric=float(lossf(net(ds['xte'].to(device)), ds['yte'].to(device)).item())
    if collect:
        stats['mean_tail_above_tau']=stats['tail_sum']/max(stats['steps']*len(list(net.parameters())),1)
        stats['mean_kurtosis_bound']=stats['k_sum']/max(stats['steps']*len(list(net.parameters())),1)
        stats['mean_tau']=stats['tau_sum']/max(stats['steps']*len(list(net.parameters())),1)
        stats['clip_fraction']=stats['clipped_entries']/max(stats['entries'],1)
        return metric, stats
    return metric


def fn(idea, cfg): return lambda s: run(s,cfg,idea=idea)

# Baseline sweep uses exactly the idea's three learning rates and four sweep seeds.
base = sweep_baseline(lambda cfg: fn(False,cfg), GRID, seeds=(0,1,2,3))
best = base['best_cfg']
# Idea at best baseline lr and two nearby settings, same grid; choose by the same 4-seed tuning protocol.
idea_trials=[]
for cfg in GRID:
    r=[run(s,cfg,idea=True) for s in (0,1,2,3)]
    idea_trials.append({'cfg':cfg,'mean':float(np.mean(r))})
idea_cfg=min(GRID,key=lambda c: next(q['mean'] for q in idea_trials if q['cfg']==c))
idea_res={'cfg':idea_cfg, **__import__('bench').protocol.evaluate(fn(True,idea_cfg), seeds=SEEDS)}
# Collect trained-model behavior on all paired seeds (not an analytical toy signature).
beh=[]
for s in SEEDS:
    val, st=run(s,idea_cfg,idea=True,collect=True); beh.append(st)
mean_tail=float(np.mean([z['mean_tail_above_tau'] for z in beh])); mean_clip=float(np.mean([z['clip_fraction'] for z in beh]))
mean_k=float(np.mean([z['mean_kurtosis_bound'] for z in beh]))
# Prediction: calibrated threshold should make observed two-sided exceedance no larger than target order.
signature={'prediction':'trained gradient-coordinate tail above calibrated tau is controlled near target delta=.02',
           'target_delta':0.02, 'observed_mean_tail_above_tau':mean_tail,
           'mean_clip_fraction':mean_clip, 'mean_kurtosis_bound':mean_k,
           'confirmed': bool(mean_tail <= 0.02*2.0)}
report=make_report('tabular','mlp_tiny',base,idea_res,{'mechanism_signature':signature,
    'math_check':math_check(),'idea_sweep':idea_trials,
    'protocol_note':'8 paired seeds; baseline sweep and idea sweep share the full learning-rate union; Adam and architecture are otherwise identical.'})
report['idea']['behavior_all_seeds']=beh
OUT.write_text(json.dumps(report,indent=2))
print(json.dumps(report,indent=2))