import json
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

import minkowski_track as track

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
EPOCHS = 18
BATCH = 128
# The union is shared by baseline and idea, as required by the protocol.
GRID = [{"lr": 1e-3, "weight_decay": 0.0},
        {"lr": 3e-3, "weight_decay": 0.0},
        {"lr": 1e-2, "weight_decay": 0.0}]


class AdditiveLatentNet(nn.Module):
    """Same two-layer 64-wide MLP trunk, followed by nonnegative tuple weights."""
    def __init__(self, input_dim):
        super().__init__()
        self.trunk = nn.Sequential(nn.Linear(input_dim, 64), nn.ReLU(),
                                   nn.Linear(64, 64), nn.ReLU())
        self.weights = nn.Linear(64, 2)

    def forward(self, x):
        h = self.trunk(x)
        ab = torch.nn.functional.softplus(self.weights(h))
        a = x[:, :track.N_DIR * track.D]
        b = x[:, track.N_DIR * track.D:]
        # Componentwise Minkowski addition, with learned nonnegative scaling.
        return ab[:, :1] * a + ab[:, 1:2] * b


def dataset(seed):
    d = track.get_dataset(seed, 400, 400)
    for k in ("xtr", "ytr", "xte", "yte"):
        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
    d["input_shape"] = tuple(d["xtr"].shape[1:])
    d["out_dim"] = track.N_DIR * track.D
    return d


def train_baseline(cfg, seed, return_model=False):
    torch.manual_seed(10000 + seed)
    np.random.seed(10000 + seed)
    d = dataset(seed)
    net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
    net, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"],
                                 batch=BATCH, weight_decay=cfg["weight_decay"],
                                 log=lambda _: None)
    return (float(metric), net, d) if return_model else float(metric)


def train_idea(cfg, seed, return_model=False):
    torch.manual_seed(10000 + seed)
    np.random.seed(10000 + seed)
    d = dataset(seed)
    net = AdditiveLatentNet(int(np.prod(d["input_shape"])))
    net, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"],
                                 batch=BATCH, weight_decay=cfg["weight_decay"],
                                 log=lambda _: None)
    return (float(metric), net, d) if return_model else float(metric)


def baseline_factory(cfg):
    return lambda seed: train_baseline(cfg, seed)


def idea_sweep():
    rows = []
    for cfg in GRID:
        r = evaluate(lambda s, c=cfg: train_idea(c, s), seeds=SWEEP_SEEDS)
        rows.append({"cfg": cfg, "mean": r["mean"]})
    best = min(rows, key=lambda z: z["mean"])["cfg"]
    return best, rows


def mechanism_signature(cfg):
    add_err, ampl = [], []
    for seed in SEEDS:
        metric, net, d = train_idea(cfg, seed, return_model=True)
        net.eval()
        device = next(net.parameters()).device
        with torch.no_grad():
            x = d["xte"][:64].to(device)
            exact = d["yte"][:64].to(device)
            pred = net(x)
            # Trained-model output error relative to the observed task target.
            add_err.append(torch.linalg.vector_norm(pred - exact, dim=1).mean().item())
            noise = torch.randn_like(x) * 1e-3
            p0, p1 = net(x), net(x + noise)
            out_delta = torch.linalg.vector_norm(p1 - p0, dim=1)
            # This is the NN-scale observed amplification of the composed tuple.
            inp_delta = torch.linalg.vector_norm(noise[:, :track.N_DIR * track.D], dim=1)
            inp_delta = torch.maximum(inp_delta,
                                      torch.linalg.vector_norm(noise[:, track.N_DIR * track.D:], dim=1))
            ampl.extend((out_delta / (inp_delta + 1e-12)).cpu().numpy().tolist())
    observed = float(np.median(ampl))
    return {
        "quantity": "trained additive output error and local perturbation amplification",
        "predicted": {"componentwise_addition": "nonnegative weighted sum", "amplification_bound": "approximately alpha+beta"},
        "observed_mean_additivity_error": float(np.mean(add_err)),
        "observed_median_local_amplification": observed,
        "observed_p90_local_amplification": float(np.percentile(ampl, 90)),
        "confirmed": bool(np.mean(add_err) < 0.35 and observed < 3.0),
    }


def main():
    base = sweep_baseline(baseline_factory, GRID, seeds=SWEEP_SEEDS)
    best_idea, idea_sweep_rows = idea_sweep()
    idea = evaluate(lambda s: train_idea(best_idea, s), seeds=SEEDS)
    report = make_report(
        "minkowski_composition", "mlp_tiny", base, idea,
        {"custom_track": {"name": track.META["name"], "file": "minkowski_track.py", "domain": track.META["domain"]},
         "idea_sweep": idea_sweep_rows,
         "mechanism_signature": mechanism_signature(best_idea)})
    report["baseline"]["idea_union_grid"] = GRID
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()