import json
import random
import sys
from pathlib import Path
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
import cochain_track as ct

SEEDS = tuple(range(8))
SWEEP_SEEDS = (0, 1, 2, 3)
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS = 24
NTR, NTE = 400, 200
WIDTH = 16
M0 = np.ones(3)
M1 = np.ones(3)


def weighted_qr(a, masses):
    q, _ = np.linalg.qr(np.sqrt(masses)[:, None] * a)
    return q / np.sqrt(masses)[:, None]


def mass_norm(a, mi, mo):
    s = np.sqrt(mo)[:, None] * a * (1.0 / np.sqrt(mi))[None, :]
    return float(np.linalg.svd(s, compute_uv=False)[0])


def math_check():
    assert np.max(np.abs(ct.D1 @ ct.D0)) < 1e-12
    q = weighted_qr(ct.D0[:, :2], M1)
    p0 = np.eye(3)
    p1 = q @ q.T @ np.diag(M1)
    # A deliberately unrelated rank-2 edge subspace is the control.
    r = weighted_qr(np.array([[1., 0.], [0., 1.], [1., 1.]]), M1)
    bad = r @ r.T @ np.diag(M1)
    return {
        "incidence_d1_d0_max": float(np.max(np.abs(ct.D1 @ ct.D0))),
        "compatible_idempotence": mass_norm(p1 @ p1 - p1, M1, M1),
        "compatible_bound": mass_norm(p1, M1, M1),
        "compatible_commutation_d0": mass_norm(ct.D0 @ p0 - p1 @ ct.D0, M0, M1),
        "compatible_commutation_d1": mass_norm(ct.D1 @ p1, M1, np.ones(1)),
        "independent_commutation_d0": mass_norm(ct.D0 @ p0 - bad @ ct.D0, M0, M1),
        "independent_commutation_d1": mass_norm(ct.D1 @ bad, M1, np.ones(1)),
    }

Q1 = torch.tensor(weighted_qr(ct.D0[:, :2], M1), dtype=torch.float32)
P1 = Q1 @ Q1.T @ torch.diag(torch.tensor(M1, dtype=torch.float32))


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


def tensors(seed):
    d = ct.get_dataset(seed, NTR, NTE)
    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:])
    return d


class CochainNet(nn.Module):
    def __init__(self, projected):
        super().__init__()
        self.projected = projected
        self.node = nn.Linear(2, WIDTH)
        self.edge = nn.Linear(2, WIDTH)
        self.head = nn.Sequential(nn.Linear(6 * WIDTH, 32), nn.ReLU(), nn.Linear(32, 1))

    def forward(self, x, return_hidden=False):
        hn = torch.relu(self.node(x[:, :3, :]))
        he = torch.relu(self.edge(x[:, 3:, :]))
        if self.projected:
            he = torch.einsum("ij,bjc->bic", P1.to(he.device), he)
        h = torch.cat((hn, he), dim=1).flatten(1)
        out = self.head(h)
        return (out, hn, he) if return_hidden else out


def run(kind, lr, seed, return_model=False):
    seed_all(seed)
    d = tensors(seed)
    net, metric, hist = train_model(CochainNet(kind == "idea"), d,
                                    epochs=EPOCHS, lr=float(lr), batch=128,
                                    log=lambda *a, **k: None)
    if return_model:
        return float(metric), net, d
    return float(metric)


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


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


def mechanism_signature():
    # Retrain/evaluate actual benchmark systems, then measure the spatial
    # derivative commutation defect on their learned hidden edge representations.
    vals = {"baseline": [], "idea": []}
    for seed in (0, 1, 2, 3):
        for kind in vals:
            _, net, d = run(kind, 0.003, seed, return_model=True)
            net = net.to("cpu")
            with torch.no_grad():
                _, hn, he = net(d["xte"], return_hidden=True)
                boundary = torch.einsum("ij,bjc->bic", torch.tensor(ct.D1, dtype=torch.float32), he)
                vals[kind].append(float(torch.sqrt(torch.mean(boundary ** 2))))
    b, i = np.mean(vals["baseline"]), np.mean(vals["idea"])
    return {"prediction": "trained projected edge representations have lower face-boundary residual",
            "predicted_direction": "idea < baseline",
            "baseline_boundary_rms": float(b), "idea_boundary_rms": float(i),
            "reduction_fraction": float((b - i) / b) if b else 0.0,
            "confirmed": bool(i < b and np.isfinite(i) and np.isfinite(b))}


def main():
    checks = math_check()
    grid = [{"lr": x} for x in LRS]
    base = sweep_baseline(base_factory, grid, seeds=SWEEP_SEEDS)
    trials = [{"cfg": c, "result": evaluate(idea_factory(c), SEEDS)} for c in grid]
    best = min(trials, key=lambda z: z["result"]["mean"])
    # Full baseline at every idea-tested learning rate enforces search-space parity.
    full_base = [{"cfg": c, "result": evaluate(base_factory(c), SEEDS)} for c in grid]
    chosen_base = min(full_base, key=lambda z: z["result"]["mean"])
    base["full"] = chosen_base["result"]
    report = make_report("simplicial_cochain_regression", "cochain_mlp_tiny",
                         base, best["result"], {
                             "math_check": checks,
                             "baseline_full_grid": full_base,
                             "idea_sweep": trials,
                             "idea_config": best["cfg"],
                             "custom_track": {"name": "simplicial_cochain_regression",
                                              "file": "cochain_track.py",
                                              "domain": "simplicial_geometry_pde_like"},
                             "mechanism_signature": mechanism_signature()})
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()