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

SEEDS = tuple(range(8))
# Union-parity: every lr tried by the idea is also evaluated by baseline.
GRID = [
    {'lr': 0.0003, 'beta2': 0.99},
    {'lr': 0.0010, 'beta2': 0.99},
    {'lr': 0.0030, 'beta2': 0.999},
]
IDEA_GRID = [
    {'lr': 0.0003, 'beta2': 0.99, 'metric_beta': -0.5},
    {'lr': 0.0010, 'beta2': 0.99, 'metric_beta': -0.5},
    {'lr': 0.0030, 'beta2': 0.999, 'metric_beta': -0.5},
]

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 train_one(seed, cfg, idea=False, collect=False):
    seed_all(seed)
    ds = get_dataset('tabular', seed, n_train=400, n_test=200)
    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
    use_cuda = torch.cuda.is_available()
    device = 'cuda' if use_cuda else 'cpu'
    try:
        net.to(device)
    except Exception:
        device = 'cpu'; net.to(device)
    x, y = ds['xtr'].to(device), ds['ytr'].to(device)
    params = list(net.parameters())
    mtrace=[]
    vtrace=[]
    if idea:
        v = [torch.zeros_like(p) for p in params]
    else:
        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], betas=(0.9, cfg['beta2']))
    loss_fn = nn.MSELoss(); violations = 0.0; hist=[]
    for ep in range(12):
        perm = torch.randperm(len(x), device=device)
        total = 0.0
        for start in range(0, len(x), 128):
            idx = perm[start:start+128]
            loss = loss_fn(net(x[idx]), y[idx])
            net.zero_grad(set_to_none=True); loss.backward()
            if not idea:
                opt.step()
            else:
                # State-dependent diagonal metric M=diag(exp(beta log(vhat))).
                with torch.no_grad():
                    for j,p in enumerate(params):
                        g = p.grad
                        v[j].mul_(cfg['beta2']).addcmul_(g, g, value=1-cfg['beta2'])
                    step = ep * ((len(x)+127)//128) + start//128 + 1
                    for j,p in enumerate(params):
                        g=p.grad
                        vhat=v[j] / (1.0-cfg['beta2']**step)
                        logm=torch.clamp(cfg['metric_beta']*torch.log(vhat+1e-8), np.log(1e-3), np.log(10.0))
                        m=torch.exp(logm)
                        z=p - cfg['lr']*m*g
                        proposal=torch.clamp(z, -2.0, 2.0)
                        p.add_(proposal-p)
                        mtrace.append(float(m.mean().cpu()))
                        vtrace.append(float(vhat.mean().cpu()))
                        violations=max(violations, float(torch.relu(-2.0-p).max().cpu()), float(torch.relu(p-2.0).max().cpu()))
            total += float(loss.detach().cpu()) * len(idx)
        hist.append(total/len(x))
    net.eval()
    with torch.no_grad():
        pred=net(ds['xte'].to(device)); metric=float(((pred-ds['yte'].to(device))**2).mean().cpu())
    result={'metric':metric, 'history':hist, 'max_constraint_violation':violations}
    if collect:
        result['model']=net; result['signature_metric_means']=mtrace; result['signature_variances']=vtrace
    return result

def baseline_factory(cfg):
    return lambda seed: train_one(seed, cfg, idea=False)['metric']

def idea_eval(cfg):
    vals=[]
    for s in SEEDS:
        vals.append(train_one(s,cfg,idea=True)['metric'])
    return {'per_seed':vals, 'mean':float(np.mean(vals))}

def signature(seed, base_cfg, idea_cfg):
    b=train_one(seed,base_cfg,False,True); q=train_one(seed,idea_cfg,True,True)
    # Prediction tested on trained NN: inverse-variance metric gives larger m
    # to coordinates with smaller observed second moments.
    a=np.asarray(q['signature_metric_means']); vv=np.asarray(q['signature_variances'])
    ok=len(a)>2 and np.isfinite(a).all() and np.isfinite(vv).all()
    corr=float(np.corrcoef(np.log(a+1e-12), -np.log(vv+1e-12))[0,1]) if ok else float('nan')
    confirmed=bool(ok and corr > 0.8)
    return {'predicted':'m decreases as observed v increases (inverse-variance metric)',
            'observed_metric_mean':float(np.mean(a)) if len(a) else None,
            'observed_log_metric_vs_negative_log_variance_corr':corr,
            'baseline_final_train_loss':b['history'][-1],
            'idea_final_train_loss':q['history'][-1],
            'max_constraint_violation':q['max_constraint_violation'],
            'confirmed':confirmed}

def main():
    base=sweep_baseline(baseline_factory, GRID, seeds=(0,1,2,3))
    best=base['best_cfg']
    idea_cfgs=[c for c in IDEA_GRID if c['lr'] in {x['lr'] for x in GRID}]
    idea_runs=[(c,idea_eval(c)) for c in idea_cfgs]
    best_idea_cfg,best_idea=max(idea_runs,key=lambda z:z[1]['mean']) if False else min(idea_runs,key=lambda z:z[1]['mean'])
    extra=signature(0,best,best_idea_cfg)
    rep=make_report('tabular','mlp_tiny',base,best_idea,extra)
    rep['idea']['best_cfg']=best_idea_cfg
    rep['idea']['nearby_settings']=[{'cfg':c,'mean':r['mean']} for c,r in idea_runs]
    rep['protocol_note']='Baseline and idea use identical tabular task, mlp_tiny architecture, 12 epochs, batch 128, and the same lr union; baseline Adam beta2 was swept.'
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()