import json
import sys
from pathlib import Path
import numpy as np
import torch

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

SEEDS = tuple(range(8))
NTRAIN, NTEST = 400, 100
EPOCHS, BATCH = 18, 128
# Shared union: both systems are evaluated at every learning rate.
GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]


def ds_torch(seed, kind):
    raw = track.get_dataset(seed, NTRAIN, NTEST)
    d = track.encode(raw, kind)
    return {**d,
            "xtr": torch.tensor(d["xtr"], dtype=torch.float32),
            "ytr": torch.tensor(d["ytr"], dtype=torch.float32),
            "xte": torch.tensor(d["xte"], dtype=torch.float32),
            "yte": torch.tensor(d["yte"], dtype=torch.float32)}


def train_one(seed, kind, cfg, return_model=False):
    torch.manual_seed(10000 + int(seed))
    np.random.seed(10000 + int(seed))
    d = ds_torch(seed, kind)
    net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
    model, metric, history = train_model(net, d, epochs=EPOCHS,
                                         lr=float(cfg["lr"]), batch=BATCH,
                                         log=lambda *_: None)
    if return_model:
        return model, metric, d
    return float(metric)


def baseline_factory(cfg):
    return lambda seed: train_one(seed, "baseline", cfg)


def idea_factory(cfg):
    return lambda seed: train_one(seed, "idea", cfg)


def signature(cfg):
    """Measure the trained models, not an analytic proxy.
    A simultaneous cyclic shift of both endpoints must preserve g_uv."""
    raw = track.get_dataset(0, NTRAIN, NTEST)
    n = raw["n_nodes"]
    shift = 7
    shifted = dict(raw)
    shifted["xte"] = (np.asarray(raw["xte"]) + shift) % n
    vals = {}
    for kind in ("baseline", "idea"):
        model, _, d = train_one(0, kind, cfg, return_model=True)
        if model is None:
            vals[kind] = float("nan")
            continue
        model.eval()
        a = track.encode(raw, kind)
        b = track.encode(shifted, kind)
        device = next(model.parameters()).device
        with torch.no_grad():
            pa = model(torch.tensor(a["xte"], dtype=torch.float32, device=device)).cpu()
            pb = model(torch.tensor(b["xte"], dtype=torch.float32, device=device)).cpu()
        vals[kind] = float(torch.abs(pa - pb).mean())
    confirmed = bool(np.isfinite(vals["idea"]) and vals["idea"] < 0.25 * max(vals["baseline"], 1e-12))
    return {"prediction": "relative group differences are invariant to simultaneous vertex translation",
            "baseline_mean_output_change": vals["baseline"],
            "idea_mean_output_change": vals["idea"],
            "predicted_idea_change": 0.0, "confirmed": confirmed}


def main():
    math = track.math_check()
    assert math["path_independence"] and math["cycle_sum_mod_n"] == 0
    base = sweep_baseline(baseline_factory, GRID, seeds=SEEDS)
    idea_trials = []
    for cfg in GRID:
        r = evaluate(idea_factory(cfg), seeds=SEEDS)
        idea_trials.append({"cfg": cfg, "mean": r["mean"], "per_seed": r["per_seed"]})
    best_cfg = min(idea_trials, key=lambda x: x["mean"])["cfg"]
    idea_full = evaluate(idea_factory(best_cfg), seeds=SEEDS)
    idea_full["best_cfg"] = best_cfg
    idea_full["sweep"] = [{"cfg": x["cfg"], "mean": x["mean"]} for x in idea_trials]
    rep = make_report("cayley_cycle_distance", "mlp_tiny", base, idea_full,
                      {"mechanism_signature": signature(best_cfg),
                       "custom_track": {"name": track.META["name"], "file": "cayley_graph_track.py", "domain": track.META["domain"]},
                       "math_check": math,
                       "protocol_notes": {"epochs": EPOCHS, "batch": BATCH, "paired_seeds": list(SEEDS), "grid_union": GRID,
                                          "structure_justification": "The custom task is a cyclic Cayley graph distance prediction problem, directly containing generator increments, cycle constraints, and relative group displacements."}})
    Path("bench_report.json").write_text(json.dumps(rep, indent=2, sort_keys=True))
    print(json.dumps(rep, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()