import json
import math
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import make_model, train_model, evaluate, sweep_baseline, make_report
from sphere_track import get_dataset, META

OUT = Path("bench_report.json")


def exp_sphere(x, v):
    r = torch.linalg.vector_norm(v, dim=-1, keepdim=True)
    return torch.cos(r) * x + torch.sinc(r / math.pi) * v


def tangent(x, u):
    return u - (u * x).sum(dim=-1, keepdim=True) * x


class ResidualSystem(nn.Module):
    """Same MLP backbone; only the residual integration rule differs."""
    def __init__(self, idea, h):
        super().__init__()
        self.backbone = make_model("mlp_tiny", (3,), 3)
        self.idea = bool(idea)
        self.h = float(h)

    def forward(self, x):
        raw = self.backbone(x)
        if self.idea:
            return exp_sphere(x, self.h * tangent(x, raw))
        return x + self.h * raw


def dataset(seed):
    d = get_dataset(seed, 400, 400)
    out = {k: torch.as_tensor(v, dtype=torch.float32) for k, v in d.items()
           if k in ("xtr", "ytr", "xte", "yte")}
    out.update({"task": "regression", "metric": "mse", "input_shape": (3,), "out_dim": 3})
    return out


def run_one(idea, cfg, seed, return_model=False):
    torch.manual_seed(10000 + int(seed))
    np.random.seed(10000 + int(seed))
    net, metric, history = train_model(
        ResidualSystem(idea, cfg["h"]), dataset(seed),
        epochs=20, lr=cfg["lr"], batch=128, weight_decay=0.0,
        log=lambda *_args, **_kwargs: None)
    if metric is None:
        raise RuntimeError("benchmark training failed")
    return (float(metric), net) if return_model else float(metric)


def main():
    # Cheap numerical verification is performed before any training.
    torch.manual_seed(3)
    x = torch.randn(1000, 3)
    x = x / torch.linalg.vector_norm(x, dim=-1, keepdim=True)
    u = torch.randn_like(x)
    v = 0.12 * tangent(x, u)
    exp_norm_error = float((torch.linalg.vector_norm(exp_sphere(x, v), dim=-1) - 1).abs().max())
    add_norm_error = float((torch.linalg.vector_norm(x + v, dim=-1) - 1).abs().mean())
    if not (exp_norm_error < 1e-6 and add_norm_error > 1e-5):
        raise RuntimeError("core exponential-map norm sanity check failed")

    lrs = [1e-3, 3e-3, 1e-2]
    hs = [0.06, 0.12, 0.24]
    grid = [{"lr": lr, "h": h} for lr in lrs for h in hs]
    # Baseline sees the complete union of all idea hyperparameters.
    base = sweep_baseline(
        lambda cfg: lambda seed: run_one(False, cfg, seed), grid)
    best = base["best_cfg"]
    idea_grid = [best]
    for cfg in grid:
        if cfg != best and len(idea_grid) < 3:
            idea_grid.append(cfg)
    idea_trials = []
    for cfg in idea_grid:
        r = evaluate(lambda seed, c=cfg: run_one(True, c, seed))
        idea_trials.append({"cfg": cfg, "result": r})
    best_idea_trial = min(idea_trials, key=lambda z: z["result"]["mean"])
    idea = best_idea_trial["result"]
    rep = make_report("sphere_one_step_dynamics", "mlp_tiny", base, idea, extra={
        "sanity_check": {
            "prediction": "Exp preserves unit norm while additive Euler drifts",
            "exp_max_norm_error": exp_norm_error,
            "additive_mean_norm_error": add_norm_error,
            "confirmed": bool(exp_norm_error < 1e-6 and add_norm_error > 1e-5),
        },
        "idea_trials": idea_trials,
    })
    rep["custom_track"] = {
        "name": META["name"], "file": "sphere_track.py", "domain": META["domain"]}
    # Signature is measured from independently trained systems at their selected configs.
    cfgb = best
    mb_metric, mb = run_one(False, cfgb, 0, return_model=True)
    mi_metric, mi = run_one(True, best_idea_trial["cfg"], 0, return_model=True)
    with torch.no_grad():
        xx = dataset(0)["xte"]
        xx_b = xx.to(next(mb.parameters()).device)
        xx_i = xx.to(next(mi.parameters()).device)
        pb = mb(xx_b); pi = mi(xx_i)
        rep["mechanism_signature"] = {
            "trained_model_seed": 0,
            "baseline_mean_abs_norm_error": float((pb.norm(dim=1) - 1).abs().mean()),
            "idea_mean_abs_norm_error": float((pi.norm(dim=1) - 1).abs().mean()),
            "predicted_idea_lower_constraint_error": True,
            "confirmed": bool((pi.norm(dim=1) - 1).abs().mean() < (pb.norm(dim=1) - 1).abs().mean()),
        }
    OUT.write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))


if __name__ == "__main__":
    main()