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

SEEDS = tuple(range(8))
SWEEP_SEEDS = (0, 1, 2, 3)
EPOCHS = 12
BATCH = 128
# The idea's a priori schedules at m=30, C_N=1.
KAPPAS = (0.5, 1.0, 1.5)
M_FINE = 30
IDEA_N = {k: int(math.ceil(M_FINE ** k)) for k in KAPPAS}
# Include every idea lr and every idea data budget in the baseline grid.
LRS = (0.0015, 0.003, 0.006)
BASE_GRID = [{"lr": lr, "n_train": n} for lr in LRS for n in sorted(set(IDEA_N.values()))]


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 dataset(seed, n_train, m=M_FINE, n_test=400):
    # Track has exact 30-point PDE observations. Coarser resolutions are
    # represented by nested spatial subsampling, preserving the same operator.
    d = poisson_data(seed, n_train, n_test)
    if m == 30:
        return d
    idx = np.linspace(0, d["out_dim"] - 1, m).round().astype(int)
    d = dict(d)
    d["ytr"] = d["ytr"][:, idx].copy(); d["yte"] = d["yte"][:, idx].copy()
    d["out_dim"] = m; d["n_grid"] = m
    return d


def run_one(seed, cfg, m=M_FINE):
    seed_all(seed)
    d = dataset(seed, int(cfg["n_train"]), m=m)
    for key in ("xtr", "ytr", "xte", "yte"):
        d[key] = torch.from_numpy(np.asarray(d[key], dtype=np.float32))
    net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
    t0 = time.perf_counter()
    net, metric, hist = train_model(net, d, epochs=EPOCHS, lr=float(cfg["lr"]), batch=BATCH, log=lambda *_: None)
    wall = time.perf_counter() - t0
    if net is None: raise RuntimeError("canonical training failed")
    # Return trained behavior for the mechanism signature.
    with torch.no_grad():
        try:
            device = next(net.parameters()).device
            pred = net(d["xte"].to(device)).detach().cpu().numpy()
        except Exception:
            pred = np.zeros_like(d["yte"])
    return float(metric), {"model": net, "pred": pred, "truth": d["yte"].cpu().numpy(), "wall": wall,
                           "history": hist, "n_train": int(cfg["n_train"]), "m": m}


def factory(cfg):
    return lambda seed: run_one(seed, cfg)[0]


def main():
    t0 = time.perf_counter()
    base = sweep_baseline(factory, BASE_GRID, seeds=SWEEP_SEEDS)
    best_base_cfg = base["best_cfg"]
    # Explicitly run idea at baseline's lr and two nearby learning-rate settings;
    # each setting uses one of the predeclared resolution-aware budgets.
    idea_cfgs = [
        {"lr": best_base_cfg["lr"], "n_train": IDEA_N[1.0], "kappa": 1.0},
        {"lr": best_base_cfg["lr"], "n_train": IDEA_N[1.5], "kappa": 1.5},
        {"lr": 0.006 if best_base_cfg["lr"] != 0.006 else 0.0015,
         "n_train": IDEA_N[1.5], "kappa": 1.5},
    ]
    idea_runs = []
    for cfg in idea_cfgs:
        r = evaluate(factory(cfg), seeds=SEEDS)
        idea_runs.append({"cfg": cfg, "result": r})
    best_idea = min(idea_runs, key=lambda z: z["result"]["mean"])
    report = make_report("poisson_jfb_short_trace", "mlp_tiny", base,
                         best_idea["result"], extra={
        "mechanism_signature": mechanism_signature(best_base_cfg, best_idea["cfg"]),
        "audit": {"idea_sweep": idea_runs, "epochs": EPOCHS, "batch": BATCH,
                  "resolution_scaling": resolution_scaling(best_base_cfg)},
        "custom_track": {"name": "poisson_jfb_short_trace",
                         "file": "/home/maxwelhelp/all/math2nn/bench/custom_tracks/poisson_jfb_short_trace.py",
                         "domain": "pde"}
    })
    report["elapsed_seconds"] = time.perf_counter() - t0
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


def mechanism_signature(base_cfg, idea_cfg):
    # Re-test the claimed mechanism on trained models: increasing N should
    # improve prediction error and input coverage, not merely an analytic proxy.
    vals = []
    for n in (IDEA_N[0.5], IDEA_N[1.0], IDEA_N[1.5]):
        metric, aux = run_one(0, {"lr": idea_cfg["lr"], "n_train": n})
        err = float(np.mean((aux["pred"] - aux["truth"]) ** 2))
        vals.append({"n_train": n, "observed_test_mse": metric,
                     "observed_prediction_mse_recomputed": err})
    confirmed = vals[-1]["observed_test_mse"] < vals[0]["observed_test_mse"]
    return {"claim": "more operator pairs at fixed output resolution reduce learned error",
            "trained_model_observations": vals, "confirmed": bool(confirmed)}


def resolution_scaling(base_cfg):
    out = []
    for m in (10, 20, 30):
        n = int(math.ceil(m ** 1.0))
        metric, aux = run_one(0, {"lr": base_cfg["lr"], "n_train": n}, m=m)
        out.append({"m": m, "n_train": n, "test_mse": metric, "wall_seconds": aux["wall"]})
    return out

if __name__ == "__main__": main()