import json, random, sys
import numpy as np
import torch
from torch import nn

sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import train_model, evaluate, sweep_baseline, make_report, get_dataset as bench_get_dataset
# registered track supplied by bench

SEEDS = tuple(range(8))
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS = 28
NTR, NTE = 400, 200

class PersistenceLandscape(nn.Module):
    def __init__(self, n_grid=48, levels=4):
        super().__init__()
        self.levels = levels
        self.register_buffer("grid", torch.linspace(0., 1., n_grid))

    def forward(self, x):
        b, d = x[..., 0:1], x[..., 1:2]
        t = self.grid.view(1, 1, -1)
        v = torch.relu(torch.minimum(t - b, d - t))
        vals, _ = torch.topk(v, k=min(self.levels, v.shape[1]), dim=1)
        if vals.shape[1] < self.levels:
            z = torch.zeros(vals.shape[0], self.levels - vals.shape[1], vals.shape[2], device=x.device)
            vals = torch.cat([vals, z], dim=1)
        return vals.clamp_min(0.)


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


def as_tensors(ds):
    return {**ds,
      "xtr": torch.as_tensor(ds["xtr"], dtype=torch.float32),
      "ytr": torch.as_tensor(ds["ytr"], dtype=torch.long),
      "xte": torch.as_tensor(ds["xte"], dtype=torch.float32),
      "yte": torch.as_tensor(ds["yte"], dtype=torch.long)}


def representation(ds, kind):
    q = as_tensors(ds)
    if kind == "baseline":
        q["xtr"] = q["xtr"].flatten(1)
        q["xte"] = q["xte"].flatten(1)
        q["input_shape"] = q["xtr"].shape[1:]
    else:
        layer = PersistenceLandscape()
        with torch.no_grad():
            q["xtr"] = layer(q["xtr"]).flatten(1)
            q["xte"] = layer(q["xte"]).flatten(1)
        q["input_shape"] = q["xtr"].shape[1:]
    return q


def run(kind, lr, seed, return_model=False):
    seed_all(seed)
    ds = representation(bench_get_dataset("persistence_diagrams", seed=seed, n_train=NTR, n_test=NTE), kind)
    model = nn.Sequential(nn.Linear(int(np.prod(ds["input_shape"])), 64), nn.ReLU(), nn.Linear(64, 2))
    net, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *a, **k: None)
    if return_model:
        return float(metric), net, ds
    return float(metric)


def base_factory(cfg):
    return lambda seed: run("baseline", float(cfg["lr"]), seed)


def idea_factory(cfg):
    return lambda seed: run("idea", float(cfg["lr"]), seed)


def mechanism_signature():
    seed_all(9001)
    ds = bench_get_dataset("persistence_diagrams", seed=9001, n_train=16, n_test=16)
    x = torch.as_tensor(ds["xte"], dtype=torch.float32)
    layer = PersistenceLandscape()
    a, b = x[:8], x[:8].clone()
    b[:, :, 0] += 0.006; b[:, :, 1] += 0.006
    with torch.no_grad():
        la, lb = layer(a), layer(b)
    observed = float(torch.sqrt(torch.mean((la-lb)**2)))
    tent_a = torch.relu(torch.minimum(layer.grid.view(1,1,-1)-a[...,0:1], a[...,1:2]-layer.grid.view(1,1,-1)))
    tent_b = torch.relu(torch.minimum(layer.grid.view(1,1,-1)-b[...,0:1], b[...,1:2]-layer.grid.view(1,1,-1)))
    matched = float(torch.sqrt(torch.mean((tent_a-tent_b)**2)))
    ratio = observed / matched if matched else 0.0
    return {"prediction":"landscape perturbation does not exceed matched tent perturbation",
            "predicted_ratio_bound":1.0, "observed_ratio":ratio,
            "observed_landscape_rms":observed, "matched_tent_rms":matched,
            "confirmed": bool(np.isfinite(ratio) and ratio <= 1.02)}


def main():
    grid = [{"lr": lr} for lr in LRS]
    base = sweep_baseline(base_factory, grid, seeds=(0,1,2,3))
    idea_trials = [{"cfg": c, "result": evaluate(idea_factory(c), SEEDS)} for c in grid]
    best = min(idea_trials, key=lambda z: z["result"]["mean"])
    rep = make_report("persistence_diagrams", "mlp_tiny", base, best["result"], {
        "custom_track": {"name":"persistence_diagrams", "file":"persistence_track.py", "domain":"topological_representation"},
        "idea_config": best["cfg"], "idea_sweep": idea_trials,
        "mechanism_signature": mechanism_signature()
    })
    rep["mechanism_signature"] = rep.pop("mechanism_signature")
    with open("bench_report.json", "w") as f: json.dump(rep, f, indent=2)
    print(json.dumps(rep, indent=2))

if __name__ == "__main__": main()