import sys, json, random
import numpy as np
import torch
import torch.nn.functional as F

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

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
LRS = [1e-3, 3e-3, 1e-2]
ALPHAS = [0.3, 0.5, 0.7]
EPOCHS = 4
NTR, NTE = 300, 150
NOISE = 0.20


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        try: torch.cuda.manual_seed_all(seed)
        except Exception: pass


def dpd_loss(logits, y, alpha):
    p = F.softmax(logits.float(), dim=-1)
    A = p.pow(1.0 + alpha).sum(dim=-1)
    q = p.gather(1, y[:, None]).squeeze(1).clamp_min(1e-12)
    return (A - (1.0 + 1.0 / alpha) * q.pow(alpha)).mean()


def corrupted_ds(seed):
    d = get_dataset("vision", seed=seed, n_train=NTR, n_test=NTE)
    rng = np.random.RandomState(seed + 99173)
    y = d["ytr"].clone()
    mask = torch.as_tensor(rng.rand(len(y)) < NOISE)
    new = torch.as_tensor(rng.randint(0, 10, len(y)), dtype=torch.long)
    new = torch.where(new == y, (new + 1) % 10, new)
    y[mask] = new[mask]
    d["ytr"] = y
    d["corrupt_mask"] = mask
    return d


def run(cfg, seed, return_model=False):
    seed_all(seed)
    d = corrupted_ds(seed)
    net = make_model("cnn_small", d["input_shape"], d["out_dim"])
    # This is a custom loop only because the intervention changes the loss.
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        net = net.to(device)
        x, y = d["xtr"].to(device), d["ytr"].to(device)
        opt = torch.optim.Adam(net.parameters(), lr=float(cfg["lr"]))
        gen = torch.Generator(device="cpu").manual_seed(seed + 12345)
        for _ in range(EPOCHS):
            order = torch.randperm(len(x), generator=gen)
            for start in range(0, len(x), 128):
                ix = order[start:start+128].to(device)
                logits = net(x[ix])
                loss = (F.cross_entropy(logits, y[ix]) if cfg["kind"] == "ce"
                        else dpd_loss(logits, y[ix], float(cfg["alpha"])))
                opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
        with torch.no_grad():
            pred = net(d["xte"].to(device)).argmax(1).cpu()
            metric = float((pred != d["yte"]).float().mean())
        if return_model:
            return metric, net, d, device
        del net
        if torch.cuda.is_available(): torch.cuda.empty_cache()
        return metric
    except RuntimeError:
        # Explicit CPU fallback for a shared/fragmented GPU slot.
        seed_all(seed)
        d = corrupted_ds(seed)
        net = make_model("cnn_small", d["input_shape"], d["out_dim"]).to("cpu")
        opt = torch.optim.Adam(net.parameters(), lr=float(cfg["lr"]))
        gen = torch.Generator().manual_seed(seed + 12345)
        for _ in range(EPOCHS):
            for start in range(0, NTR, 128):
                ix = torch.randperm(NTR, generator=gen)[start:start+128]
                z = net(d["xtr"][ix]); loss = (F.cross_entropy(z, d["ytr"][ix]) if cfg["kind"] == "ce" else dpd_loss(z, d["ytr"][ix], float(cfg["alpha"])))
                opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
        with torch.no_grad(): metric = float((net(d["xte"]).argmax(1) != d["yte"]).float().mean())
        if return_model: return metric, net, d, torch.device("cpu")
        return metric


def base_factory(cfg):
    return lambda seed: run({"kind": "ce", "lr": cfg["lr"]}, seed)


def idea_factory(cfg):
    return lambda seed: run({"kind": "dpd", "lr": cfg["lr"], "alpha": cfg["alpha"]}, seed)


def mechanism_signature():
    # Measure the proposed score suppression on outputs of a trained NN.
    cfg = {"kind": "dpd", "lr": 3e-3, "alpha": 0.5}
    _, net, d, device = run(cfg, 0, return_model=True)
    net.eval(); x = d["xtr"].to(device); y = d["ytr"].to(device)
    mask = d["corrupt_mask"].to(device)
    idx = torch.where(mask)[0][:min(48, int(mask.sum()))]
    z = net(x[idx]).detach().requires_grad_(True)
    yy = y[idx]
    p = F.softmax(z, -1); q = p.gather(1, yy[:, None]).squeeze(1).clamp_min(1e-12)
    g_ce = torch.autograd.grad((-q.log()).sum(), z, retain_graph=True)[0]
    a = cfg["alpha"]
    g_obs = torch.autograd.grad((-(1+1/a)*q.pow(a)).sum(), z)[0]
    observed = (g_obs.norm(dim=1) / g_ce.norm(dim=1).clamp_min(1e-12)).detach().cpu().numpy()
    predicted = ((1+a)*q.pow(a)).detach().cpu().numpy()
    rel = np.abs(observed-predicted) / np.maximum(predicted, 1e-12)
    return {"prediction": "observed-label score gradient ratio equals (1+alpha)*q^alpha", "alpha": a,
            "n_corrupted": int(len(idx)), "mean_predicted_ratio": float(predicted.mean()),
            "mean_observed_ratio": float(observed.mean()), "max_relative_error": float(rel.max()),
            "corrupted_mean_q": float(q.detach().mean()),
            "confirmed": bool(np.isfinite(rel).all() and float(rel.max()) < 1e-4)}


def main():
    base_grid = [{"lr": lr} for lr in LRS]
    base = sweep_baseline(base_factory, base_grid, seeds=SWEEP_SEEDS)
    trials = []
    idea_lr = float(base["best_cfg"]["lr"])
    for alpha in ALPHAS:
        cfg = {"lr": idea_lr, "alpha": alpha}
        trials.append({"cfg": cfg, "result": evaluate(idea_factory(cfg), SEEDS)})
    best = min(trials, key=lambda r: r["result"]["mean"])
    extra = {"track_choice": "vision: categorical class probability output matches exact DPD integral; label corruption tests robustness", "noise_rate": NOISE,
             "idea_sweep": trials, "mechanism_signature": mechanism_signature()}
    rep = make_report("vision", "cnn_small", base, best["result"], extra)
    with open("bench_report.json", "w") as f: json.dump(rep, f, indent=2)
    print(json.dumps(rep, indent=2))

if __name__ == "__main__": main()