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

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
EPOCHS, BATCH, GAMMA = 25, 64, 0.5

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def stein_h(x, gamma=GAMMA):
    # Target is N(0,1), score s(x)=-x, and RBF exp(-gamma*(x-y)^2).
    dif=x[:,None,:]-x[None,:,:]
    r2=(dif*dif).sum(-1)
    dot=x@x.T
    d=1
    return torch.exp(-gamma*r2)*(dot + 2*gamma*d - (2*gamma+4*gamma*gamma)*r2)

def objective(out, kind):
    h=stein_h(out)
    n=out.shape[0]
    if kind == 'v':
        return torch.sqrt(torch.relu(h.mean()) + 1e-8)
    u=(h.sum()-torch.diagonal(h).sum())/(n*(n-1))
    return torch.sqrt(torch.relu(u) + 1e-8)

def train_system(seed, cfg, kind, return_model=False):
    seed_all(seed)
    ds=get_dataset('gaussian_score_matching_1d', seed, 400, 400)
    net=make_model('mlp_tiny', (1,), 1)
    # train_model cannot express the intervention, so this local loop changes
    # only the loss while retaining Adam, batch size, epochs, and architecture.
    use_cuda=torch.cuda.is_available()
    device=torch.device('cuda' if use_cuda else 'cpu')
    try:
        net.to(device)
        x= torch.as_tensor(ds['xtr'], dtype=torch.float32, device=device)
        opt=torch.optim.Adam(net.parameters(), lr=cfg['lr'])
        for _ in range(EPOCHS):
            perm=torch.randperm(len(x), device=device)
            for i in range(0,len(x),BATCH):
                o=net(x[perm[i:i+BATCH]])
                loss=objective(o,kind)
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            pred=net(torch.as_tensor(ds['xte'],dtype=torch.float32,device=device))
            metric=float(((pred-torch.as_tensor(ds['yte'],dtype=torch.float32,device=device))**2).mean().cpu())
        return (metric, net.cpu(), ds) if return_model else metric
    except Exception:
        # Explicit CPU fallback for CUDA/runtime failures.
        seed_all(seed); device=torch.device('cpu'); net=make_model('mlp_tiny',(1,),1).to(device)
        x=torch.as_tensor(ds['xtr'],dtype=torch.float32)
        opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
        for _ in range(EPOCHS):
            perm=torch.randperm(len(x))
            for i in range(0,len(x),BATCH):
                loss=objective(net(x[perm[i:i+BATCH]]),kind)
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            p=net(torch.as_tensor(ds['xte'],dtype=torch.float32))
            metric=float(((p-torch.as_tensor(ds['yte'],dtype=torch.float32))**2).mean())
        return (metric,net,ds) if return_model else metric

def train_fn(kind, cfg):
    return lambda seed: train_system(seed,cfg,kind)

def signature(cfg, kind):
    metric, net, ds=train_system(0,cfg,kind,True)
    with torch.no_grad(): out=net(torch.as_tensor(ds['xte'],dtype=torch.float32))
    h=stein_h(out).numpy(); n=len(out)
    diag=np.diag(h); off=(h.sum()-diag.sum())/(n*(n-1)); v=float(np.sqrt(max(h.mean(),0))); u=float(np.sqrt(max(off,0)))
    # Bootstrap the trained model's actual outputs, not an analytic toy graph.
    rng=np.random.RandomState(395); vs=[]; us=[]
    for _ in range(80):
        ix=rng.choice(n,64,replace=False); hh=h[np.ix_(ix,ix)]; dd=np.diag(hh)
        vs.append(np.sqrt(max(float(hh.mean()),0))); us.append(np.sqrt(max(float((hh.sum()-dd.sum())/(63*64)),0)))
    # Empirical covariance spectrum gives the predicted HS/trace fluctuation ratio.
    arr=out.numpy().reshape(-1); c=np.atleast_2d(np.cov(arr, rowvar=False)); ev=np.linalg.eigvalsh(c); ev=np.maximum(ev,0)
    pred=float(np.sqrt(np.sqrt((ev*ev).sum())/max(ev.sum(),1e-12)))
    obs=float(np.std(us,ddof=1)/max(np.std(vs,ddof=1),1e-12))
    return {"predicted_hs_over_trace_scale":pred,"observed_bootstrap_sd_U_over_V":obs,
            "trained_output_v":v,"trained_output_u":u,"test_mse":metric,
            "confirmed": bool(np.isfinite(obs) and abs(np.log(max(obs,1e-12)/max(pred,1e-12)))<1.0)}

def main():
    base=sweep_baseline(lambda c: train_fn('v',c), GRID, seeds=SWEEP_SEEDS)
    # Equal-sized idea sweep on exactly the union of baseline learning rates.
    idea_trials=[]
    for c in GRID:
        r=evaluate(train_fn('u',c), seeds=SWEEP_SEEDS)
        idea_trials.append({"cfg":c,"mean":r['mean']})
    best=min(idea_trials,key=lambda q:q['mean'])['cfg']
    idea=evaluate(train_fn('u',best), seeds=SEEDS)
    sig=signature(best,'u')
    rep=make_report('gaussian_score_matching_1d','mlp_tiny',base,idea,
                    {"mechanism_signature":sig,"custom_track":{"name":"gaussian_score_matching_1d","file":"ksd_track_1d.py","domain":"loss"},
                     "idea_sweep":idea_trials,"training":{"epochs":EPOCHS,"batch":BATCH,"gamma":GAMMA}})
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()