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, sweep_baseline, evaluate, make_report, get_dataset
import orientation_track as ot

SEEDS = tuple(range(8))
LR_GRID = [1e-3, 3e-3, 1e-2]
EPOCHS = 18
BATCH = 128


def seed_all(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def ds_torch(d):
    out = dict(d)
    out["xtr"] = torch.tensor(d["xtr"], dtype=torch.float32)
    out["ytr"] = torch.tensor(d["ytr"], dtype=torch.long)
    out["xte"] = torch.tensor(d["xte"], dtype=torch.float32)
    out["yte"] = torch.tensor(d["yte"], dtype=torch.long)
    return out


def train_one(seed, lr, invariant):
    seed_all(seed)
    raw = get_dataset("orientation_phase_quotient", seed, 400, 160)
    raw_np = dict(raw)
    for k in ("xtr", "ytr", "xte", "yte"):
        if hasattr(raw_np[k], "detach"):
            raw_np[k] = raw_np[k].detach().cpu().numpy()
    d = ot.invariant_dataset(raw_np) if invariant else raw_np
    d = ds_torch(d)
    net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
    _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    return float(metric)


def baseline_factory(cfg):
    return lambda seed: train_one(seed, float(cfg["lr"]), False)


def idea_factory(cfg):
    return lambda seed: train_one(seed, float(cfg["lr"]), True)


def trained_signature(seed, lr):
    seed_all(seed)
    raw = get_dataset("orientation_phase_quotient", seed, 400, 160)
    raw_np = dict(raw)
    for k in ("xtr", "ytr", "xte", "yte"):
        if hasattr(raw_np[k], "detach"):
            raw_np[k] = raw_np[k].detach().cpu().numpy()
    inv = ot.invariant_dataset(raw_np)
    dr, di = ds_torch(raw_np), ds_torch(inv)
    br = make_model("mlp_tiny", dr["input_shape"], dr["out_dim"])
    bi = make_model("mlp_tiny", di["input_shape"], di["out_dim"])
    br, _, _ = train_model(br, dr, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    bi, _, _ = train_model(bi, di, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    rng = np.random.default_rng(seed + 991)
    u = raw_np["xte"].reshape(-1, 4 ** 3, 3)
    q, r = np.linalg.qr(rng.normal(size=(3, 3)))
    q = q @ np.diag(np.where(np.diag(r) >= 0, 1.0, -1.0))
    if np.linalg.det(q) < 0:
        q[:, 0] *= -1
    transformed = (u @ q.T) * rng.choice([-1.0, 1.0], size=(len(u), 4 ** 3, 1))
    raw2 = dict(raw_np)
    raw2["xte"] = transformed.reshape(len(u), -1).astype(np.float32)
    inv2 = ot.invariant_dataset(raw2)
    with torch.no_grad():
        devr = next(br.parameters()).device
        devi = next(bi.parameters()).device
        pbr = torch.softmax(br(torch.tensor(raw["xte"], dtype=torch.float32, device=devr)), 1)
        pbr2 = torch.softmax(br(torch.tensor(raw2["xte"], dtype=torch.float32, device=devr)), 1)
        pbi = torch.softmax(bi(torch.tensor(inv["xte"], dtype=torch.float32, device=devi)), 1)
        pbi2 = torch.softmax(bi(torch.tensor(inv2["xte"], dtype=torch.float32, device=devi)), 1)
    return {
        "baseline_mean_output_change": float(torch.abs(pbr - pbr2).mean().cpu()),
        "idea_mean_output_change": float(torch.abs(pbi - pbi2).mean().cpu()),
        "baseline_max_output_change": float(torch.abs(pbr - pbr2).max().cpu()),
        "idea_max_output_change": float(torch.abs(pbi - pbi2).max().cpu()),
    }


def main():
    # Core numerical check is performed before any neural training.
    math_err = ot.math_check()
    grid = [{"lr": x} for x in LR_GRID]
    base = sweep_baseline(baseline_factory, grid, seeds=(0, 1, 2, 3))
    idea_sweep = []
    for cfg in grid:
        short = evaluate(idea_factory(cfg), seeds=(0, 1, 2, 3))
        idea_sweep.append({"cfg": cfg, "mean": short["mean"]})
    best_idea_cfg = min(grid, key=lambda c: next(x["mean"] for x in idea_sweep if x["cfg"] == c))
    idea_full = evaluate(idea_factory(best_idea_cfg), seeds=SEEDS)
    sig = trained_signature(0, float(best_idea_cfg["lr"]))
    sig["predicted_invariant_output_change"] = 0.0
    sig["observed_feature_change"] = math_err
    sig["confirmed"] = bool(math_err < 1e-6 and sig["idea_mean_output_change"] < sig["baseline_mean_output_change"])
    sig["prediction"] = "P2 local-correlation inputs and their trained predictions should be unchanged by global rotation and independent apolar flips."
    report = make_report(
        "orientation_phase_quotient", "mlp_tiny",
        base,
        dict(idea_full, best_cfg=best_idea_cfg, sweep=idea_sweep),
        extra=sig,
    )
    result = {
        "bench_report": report,
        "custom_track": {"name": "orientation_phase_quotient", "file": "orientation_track.py", "domain": "molecular_orientation_symmetry"},
        "math_check": {"max_transformed_feature_error": math_err, "predicted": 0.0},
        "protocol": {"epochs": EPOCHS, "batch": BATCH, "lr_union": LR_GRID, "paired_seeds": list(SEEDS)},
    }
    Path("bench_results.json").write_text(json.dumps(result, indent=2, sort_keys=True))
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()