import sys, json, math, 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, sweep_baseline, make_report
from bench.protocol import evaluate

ALPHA = 0.10
GAMMA = 0.25
K = 8
EPOCHS = 12
CAL_FRAC = 0.75
# This is the complete shared union: every idea lr is also a baseline lr.
GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 6e-3}]
SEEDS = tuple(range(8))


def cvar_lower(z, gamma):
    z = np.sort(np.asarray(z, dtype=float))
    h = gamma * len(z)
    m = int(math.floor(h))
    total = z[:m].sum()
    if h - m > 1e-12 and m < len(z):
        total += (h - m) * z[m]
    return float(total / h)


def cvar_eta(z, gamma):
    z = np.asarray(z, dtype=float)
    cand = np.r_[z, z.min() - 1., z.max() + 1.]
    vals = cand - np.maximum(cand[:, None] - z[None, :], 0).sum(1) / (gamma * len(z))
    return float(vals.max())


def conformal_q(r, alpha):
    r = np.sort(np.asarray(r, dtype=float))
    k = int(math.ceil((len(r) + 1) * (1 - alpha)))
    return float("inf") if k > len(r) else float(r[k - 1])


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 fit(seed, lr, idea):
    seed_all(seed)
    # The built-in dynamics task is the structural match: controlled pendulum
    # rollout and stability/safety margins. Use a held-out calibration slice.
    ds = get_dataset("dynamics", seed, n_train=400, n_test=400)
    ncal = int(len(ds["xtr"]) * (1 - CAL_FRAC))
    train_ds = dict(ds)
    train_ds["xtr"] = ds["xtr"][:len(ds["xtr"]) - ncal]
    train_ds["ytr"] = ds["ytr"][:len(ds["ytr"]) - ncal]
    net = make_model("rnn_small", ds["input_shape"], K)
    net, _, _ = train_model(net, train_ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
    if net is None: return float("nan"), {}
    device = next(net.parameters()).device
    net.eval()
    with torch.no_grad():
        pc = net(ds["xtr"][-ncal:].to(device)).detach().cpu().numpy()
        pt = net(ds["xte"].to(device)).detach().cpu().numpy()
    yc = ds["ytr"][-ncal:].numpy().reshape(-1)
    yt = ds["yte"].numpy().reshape(-1)
    # Clearance is signed distance to the pendulum angle limit |theta|=1.5.
    # It is computed on the selected rollout output, not analytically.
    mc = 1.5 - np.abs(pc)
    mt = 1.5 - np.abs(pt)
    clearance_cvar = np.array([cvar_lower(row, GAMMA) for row in mc])
    clearance_test_cvar = np.array([cvar_lower(row, GAMMA) for row in mt])
    q = conformal_q(clearance_cvar - (1.5 - np.abs(yc)), ALPHA)
    cert = clearance_test_cvar - q
    # Preserve predicted direction while enforcing the lower clearance bound.
    sign = np.sign(pt.mean(1)); sign[sign == 0] = 1
    pred = sign * (1.5 - cert)
    pred = np.clip(pred, -3., 3.)
    if not idea:
        pred = pt.mean(1)
    mse = float(np.mean((pred - yt) ** 2))
    # Signature is measured from this trained model's predictions and observations.
    realized_clear = 1.5 - np.abs(yt)
    coverage = float(np.mean(realized_clear >= cert)) if idea else float(np.mean(realized_clear >= clearance_test_cvar))
    return mse, {"q": q, "coverage": coverage, "certificate_nonnegative": float(np.mean(cert >= 0)),
                 "predicted_clearance_mean": float(np.mean(clearance_test_cvar)),
                 "observed_clearance_mean": float(np.mean(realized_clear))}


def baseline_factory(cfg):
    return lambda seed: fit(seed, cfg["lr"], False)[0]


def idea_factory(cfg):
    return lambda seed: fit(seed, cfg["lr"], True)[0]


def main():
    # Core math sanity check comes before any neural training.
    rng = np.random.default_rng(2747)
    errs = [abs(cvar_eta((z := rng.normal(size=K)), GAMMA) - cvar_lower(z, GAMMA)) for _ in range(20)]
    rr = np.sort(rng.normal(size=31)); q = conformal_q(rr, ALPHA)
    math_check = {"cvar_max_abs_error": float(max(errs)), "q_is_order_statistic": bool(q in rr),
                  "q_rank": int(np.where(rr == q)[0][0] + 1)}

    base = sweep_baseline(baseline_factory, GRID, seeds=(0, 1, 2, 3))
    idea_cfg = base["best_cfg"]
    # Full idea sweep at exactly the same three learning rates; best is selected
    # on the four sweep seeds, then rerun on all eight paired seeds.
    idea_sweep = []
    for cfg in GRID:
        r = evaluate(idea_factory(cfg), seeds=(0, 1, 2, 3))
        idea_sweep.append({"cfg": cfg, "mean": r["mean"]})
    idea_cfg = min(idea_sweep, key=lambda x: x["mean"])["cfg"]
    idea_full = evaluate(idea_factory(idea_cfg), seeds=SEEDS)
    # Mechanism signature is recomputed using trained models on all paired seeds.
    sig = [fit(s, idea_cfg["lr"], True)[1] for s in SEEDS]
    signature = {"definition": "trained RNN predicted-vs-observed clearance coverage",
                 "alpha": ALPHA, "gamma": GAMMA, "K": K,
                 "coverage_mean": float(np.mean([x["coverage"] for x in sig])),
                 "coverage_per_seed": [x["coverage"] for x in sig],
                 "q_mean": float(np.mean([x["q"] for x in sig])),
                 "certificate_nonnegative_mean": float(np.mean([x["certificate_nonnegative"] for x in sig])),
                 "target_coverage": 1 - ALPHA,
                 "confirmed": bool(abs(float(np.mean([x["coverage"] for x in sig])) - (1-ALPHA)) <= 0.06),
                 "math_check": math_check,
                 "idea_sweep": idea_sweep}
    report = make_report("dynamics", "rnn_small", {**base, "sweep_union": GRID}, idea_full, signature)
    report["idea"]["selected_cfg"] = idea_cfg
    report["math_check"] = math_check
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()