import json
import sys
from pathlib import Path
import numpy as np
import torch
from torch.func import functional_call, grad, vmap

sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import (get_dataset, make_model, sweep_baseline, evaluate,
                   make_report, permutation_pvalue)

# Optimizer interventions are tested on the designated optimizer track.
TRACK = "tabular"
MODEL = "mlp_tiny"
SEEDS = tuple(range(8))
EPOCHS = 8
BATCH = 128
WEIGHT_DECAY = 0.0
LR_GRID = [0.0015, 0.003, 0.006, 0.012, 0.024]
TAU_GRID = [0.5, 1.0, 2.0]


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


def centered_clip(gs, tau):
    # gs: [batch, number_of_parameters], coordinate-wise robust center.
    c = torch.median(gs, dim=0).values
    r = gs - c
    n = torch.linalg.vector_norm(r, dim=1, keepdim=True)
    scale = torch.minimum(torch.ones_like(n), torch.tensor(tau, device=gs.device) / (n + 1e-12))
    return c + r * scale, c, n


def per_example_gradients(net, xb, yb):
    """Vectorized exact per-example gradients, equivalent to individual autograd."""
    names, params = zip(*[(n, p) for n, p in net.named_parameters() if p.requires_grad])
    pmap = {n: p for n, p in zip(names, params)}
    buffers = {n: b for n, b in net.named_buffers()}
    def one_loss(pm, x, y):
        out = functional_call(net, (pm, buffers), (x.unsqueeze(0),))
        return ((out - y.unsqueeze(0)) ** 2).mean()
    grads = vmap(grad(one_loss), in_dims=(None, 0, 0))(pmap, xb, yb)
    return torch.cat([grads[n].reshape(len(xb), -1) for n in names], dim=1), None


def train_one(seed, method, lr, tau, keep_model=False):
    seed_all(seed)
    ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=400)
    net = make_model(MODEL, ds["input_shape"], ds["out_dim"])
    # CPU is deliberately the safe fallback; CUDA errors are caught and retried.
    device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        net = net.to(device)
        xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device)
        xte, yte = ds["xte"].to(device), ds["yte"].to(device)
        params = [p for p in net.parameters() if p.requires_grad]
        m = [torch.zeros_like(p) for p in params]
        v = [torch.zeros_like(p) for p in params]
        b1, b2, eps = 0.9, 0.999, 1e-8
        step = 0
        for ep in range(EPOCHS):
            gen = torch.Generator(device=device).manual_seed(seed * 1000 + ep)
            perm = torch.randperm(len(xtr), generator=gen, device=device)
            for start in range(0, len(xtr), BATCH):
                idx = perm[start:start+BATCH]
                gs, _ = per_example_gradients(net, xtr[idx], ytr[idx])
                if method == "baseline":
                    h = gs.mean(dim=0)
                    hn = torch.linalg.vector_norm(h)
                    h = h * torch.minimum(torch.tensor(1.0, device=device),
                                          torch.tensor(tau, device=device) / (hn + 1e-12))
                else:
                    h, _, _ = centered_clip(gs, tau)
                    h = h.mean(dim=0)
                off = 0
                step += 1
                with torch.no_grad():
                    for k, p in enumerate(params):
                        z = h[off:off+p.numel()].reshape_as(p); off += p.numel()
                        m[k].mul_(b1).add_(z, alpha=1-b1)
                        v[k].mul_(b2).addcmul_(z, z, value=1-b2)
                        mh = m[k] / (1 - b1 ** step)
                        vh = v[k] / (1 - b2 ** step)
                        p.addcdiv_(mh, torch.sqrt(vh) + eps, value=-lr)
        net.eval()
        with torch.no_grad():
            metric = float(((net(xte) - yte) ** 2).mean())
        return metric, (net, ds) if keep_model else None
    except (RuntimeError, torch.cuda.OutOfMemoryError):
        if device == "cuda":
            torch.cuda.empty_cache()
            # Retry the exact same seeded run on CPU.
            old = torch.cuda.is_available
            # Re-enter after making CUDA invisible is unreliable; use a CPU-only helper.
            return train_one_cpu(seed, method, lr, tau, keep_model)
        raise


def train_one_cpu(seed, method, lr, tau, keep_model=False):
    # CPU fallback uses the same algorithm and data/order, with no CUDA assumptions.
    seed_all(seed)
    ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=400)
    net = make_model(MODEL, ds["input_shape"], ds["out_dim"])
    xtr, ytr, xte, yte = ds["xtr"], ds["ytr"], ds["xte"], ds["yte"]
    params = list(net.parameters()); m=[torch.zeros_like(p) for p in params]; v=[torch.zeros_like(p) for p in params]
    step=0
    for ep in range(EPOCHS):
        gen=torch.Generator().manual_seed(seed*1000+ep); perm=torch.randperm(len(xtr),generator=gen)
        for start in range(0,len(xtr),BATCH):
            gs,_=per_example_gradients(net,xtr[perm[start:start+BATCH]],ytr[perm[start:start+BATCH]])
            if method=="baseline":
                h=gs.mean(0); n=torch.linalg.vector_norm(h); h=h*min(1.,tau/(float(n)+1e-12))
            else: h=centered_clip(gs,tau)[0].mean(0)
            step+=1; off=0
            with torch.no_grad():
                for k,p in enumerate(params):
                    z=h[off:off+p.numel()].reshape_as(p); off+=p.numel(); m[k]=.9*m[k]+.1*z; v[k]=.999*v[k]+.001*z*z
                    p.addcdiv_(m[k]/(1-.9**step),torch.sqrt(v[k]/(1-.999**step))+1e-8,value=-lr)
    with torch.no_grad(): metric=float(((net(xte)-yte)**2).mean())
    return metric,(net,ds) if keep_model else None


def fn(cfg, method):
    return lambda seed: train_one(seed, method, cfg["lr"], cfg["tau"])[0]


def signature(cfg):
    vals=[]
    for s in SEEDS:
        metric, obj=train_one(s,"idea",cfg["lr"],cfg["tau"],True)
        net,ds=obj; dev=next(net.parameters()).device; x,y=ds["xtr"][:BATCH].to(dev),ds["ytr"][:BATCH].to(dev)
        gs,_=per_example_gradients(net,x,y); clipped,c,n=centered_clip(gs,cfg["tau"])
        h=clipped.mean(0); raw=gs.mean(0)-c
        cos=float(torch.dot(h-c,raw)/(torch.linalg.vector_norm(h-c)*torch.linalg.vector_norm(raw)+1e-12))
        vals.append({"metric":metric,"observed_max_residual":float(torch.linalg.vector_norm(clipped-c,dim=1).max()),"predicted_tau":cfg["tau"],"direction_cosine":cos})
    return {"prediction":"residual norm <= tau and common-direction cosine remains near 1","predicted_residual_bound":cfg["tau"],"observed_max_residual_mean":float(np.mean([z["observed_max_residual"] for z in vals])),"observed_direction_cosine_mean":float(np.mean([z["direction_cosine"] for z in vals])),"confirmed":bool(max(z["observed_max_residual"] for z in vals) <= cfg["tau"]*1.00001 and min(z["direction_cosine"] for z in vals)>0.99),"per_seed":vals}


def main():
    grid=[{"lr":lr,"tau":tau} for lr in LR_GRID for tau in TAU_GRID]
    base=sweep_baseline(lambda cfg: fn(cfg,"baseline"),grid,seeds=SEEDS)
    best=base["best_cfg"]
    idea_cfgs=[best,{"lr":0.006,"tau":best["tau"]},{"lr":0.024,"tau":best["tau"]}]
    idea_runs=[]
    for cfg in idea_cfgs:
        r=evaluate(fn(cfg,"idea"),seeds=SEEDS); r["cfg"]=cfg; idea_runs.append(r)
    ibest=min(idea_runs,key=lambda x:x["mean"])
    diffs=[a-b for a,b in zip(base["full"]["per_seed"],ibest["per_seed"])]
    rep=make_report(TRACK,MODEL,base,ibest,{"mechanism_signature":signature(ibest["cfg"]),"idea_grid":idea_runs,"paired_diffs_baseline_minus_idea":diffs,"permutation_pvalue":permutation_pvalue(diffs)})
    Path("bench_report.json").write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))

if __name__=="__main__": main()