import json, sys, random
from pathlib import Path
import numpy as np
import torch
sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report

N_TRAIN, N_TEST, EPOCHS, M = 400, 400, 3, 2
LR_GRID = [1e-3, 3e-3, 1e-2]
SEEDS = tuple(range(8))
ALPHA = 0.10

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 train_system(seed, lr, mode, kappa=1.0):
    seed_all(seed)
    d = get_dataset("vision", seed, N_TRAIN, N_TEST)
    n = len(d["xtr"]); cut = int(.8*n)
    tr = dict(d); tr["xtr"] = d["xtr"][:cut]; tr["ytr"] = d["ytr"][:cut]
    calx, caly = d["xtr"][cut:], d["ytr"][cut:]
    testx, testy = d["xte"], d["yte"]
    probs_cal, probs_test, unc = [], [], []
    for j in range(M):
        seed_all(seed * 100 + j)
        net = make_model("cnn_small", d["input_shape"], d["out_dim"])
        net, _, _ = train_model(net, tr, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
        if net is None: return float("nan"), {}
        dev = next(net.parameters()).device
        with torch.no_grad():
            pc = torch.softmax(net(calx.to(dev)), 1).cpu().numpy()
            pt = torch.softmax(net(testx.to(dev)), 1).cpu().numpy()
        probs_cal.append(pc); probs_test.append(pt)
        unc.append(float(-(pc * np.log(np.maximum(pc, 1e-8))).sum(1).mean()))
    pc = np.stack(probs_cal); pt = np.stack(probs_test)
    if mode == "baseline":
        fused = pt.mean(0)
        pred = fused.argmax(1)
        return float((pred != testy.numpy()).mean()), {"accuracy": float((pred == testy.numpy()).mean())}
    idx = np.arange(len(caly.numpy()))
    raw = np.stack([1.0 / (10.0 * np.maximum(pc[j], 1e-8)[idx, caly.numpy()]) for j in range(M)])
    normalizers = raw.mean(1) * 1.01
    ev = 1.0 / (10.0 * np.maximum(pt, 1e-8)) / normalizers[:, None, None]
    weights = np.exp(-kappa * np.asarray(unc)); weights /= weights.sum()
    fused_e = np.sum(ev * weights[:, None, None], axis=0)
    pred = np.argmin(fused_e, 1)
    sets = fused_e < 1.0 / ALPHA
    true_e = fused_e[np.arange(len(testy)), testy.numpy()]
    return float((pred != testy.numpy()).mean()), {"accuracy": float((pred == testy.numpy()).mean()), "coverage": float(sets[np.arange(len(testy)), testy.numpy()].mean()), "avg_set_size": float(sets.sum(1).mean()), "mean_true_e": float(true_e.mean()), "predicted_mean_true_e": float(weights @ (ev[:, np.arange(len(testy)), testy.numpy()].mean(1))), "convexity_abs_error": float(abs(true_e.mean() - weights @ (ev[:, np.arange(len(testy)), testy.numpy()].mean(1))))}

def main():
    # Equal-budget sweep: every idea learning rate is also evaluated for baseline.
    grid = [{"lr": x} for x in LR_GRID]
    base = sweep_baseline(lambda cfg: lambda s: train_system(s, cfg["lr"], "baseline")[0], grid, seeds=(0,1,2,3))
    best_lr = base["best_cfg"]["lr"]
    idea_grid = [best_lr] + [x for x in LR_GRID if x != best_lr]
    # Record the same three learning rates on the idea side before choosing one.
    idea_sweep = [{"lr": lr, "result": evaluate(lambda s, lr=lr: train_system(s, lr, "idea", 1.0)[0], seeds=(0,1,2,3))} for lr in idea_grid]
    idea_best_lr = min(idea_sweep, key=lambda z: z["result"]["mean"])["lr"]
    idea = evaluate(lambda s: train_system(s, idea_best_lr, "idea", 1.0)[0], seeds=SEEDS)
    extra = {"prediction": "fused mean e-value equals uncertainty-weighted mean local e-values", "observed": [], "confirmed": False}
    for s in SEEDS:
        _, z = train_system(s, best_lr, "idea", 1.0); extra["observed"].append(z)
    errs = [z["convexity_abs_error"] for z in extra["observed"]]
    extra["max_abs_error"] = float(max(errs)); extra["confirmed"] = bool(extra["max_abs_error"] < 1e-6)
    base["idea_grid"] = idea_sweep
    report = make_report("vision", "cnn_small", base, idea, extra)
    report["idea_sweep"] = idea_sweep
    report["notes"] = "Vision selected because labels/candidate-label e-values are classification-specific; N_TRAIN/N_TEST=400 and 3 epochs were used for budget."
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()