import sys, json, random
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 get_dataset, make_model, evaluate, sweep_baseline, make_report

TRACK = "dynamics"
MODEL = "rnn_small"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

class Potential(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(dim, 32), nn.Tanh(), nn.Linear(32, 1))
    def forward(self, x):
        z = self.net(x)
        return z - z.mean()                 # minibatch gauge fixing

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 make_next(x, pred):
    # A learned one-coordinate state transition: retain the observed window and
    # replace its final angle by the model's next-angle prediction.
    z = x.clone()
    z[:, -3] = pred[:, 0]
    return z

def train_one(seed, cfg, mode, capture=False):
    global DEVICE
    seed_all(seed)
    ds = get_dataset(TRACK, seed, n_train=400, n_test=200)
    net = make_model(MODEL, ds["input_shape"], ds["out_dim"])
    pot = Potential(ds["xtr"].shape[1]) if mode == "coh" else None
    c = nn.Parameter(torch.tensor(0.0)) if mode == "coh" else None
    params = list(net.parameters()) + ([] if pot is None else list(pot.parameters()) + [c])
    opt = torch.optim.Adam(params, lr=cfg["lr"])
    dev = DEVICE
    try:
        net.to(dev); dsx, dsy = ds["xtr"].to(dev), ds["ytr"].to(dev)
        if pot is not None: pot.to(dev); c.data = c.data.to(dev)
        for ep in range(cfg["epochs"]):
            net.train()
            perm = torch.randperm(len(dsx), device=dev)
            for ii in range(0, len(dsx), 64):
                x = dsx[perm[ii:ii+64]].detach().requires_grad_(True)
                pred = net(x)
                task = ((pred - dsy[perm[ii:ii+64]]) ** 2).mean()
                # selected 1D unstable-coordinate log Jacobian proxy
                g = torch.autograd.grad(pred[:, 0].sum(), x, create_graph=True)[0]
                ell = torch.log(g[:, -3].abs() + 1e-3)
                if mode == "none":
                    reg = 0.0
                elif mode == "point":
                    reg = ((ell - c0(ell)) ** 2).mean()
                else:
                    nxt = make_next(x, pred)
                    reg = (ell - pot(nxt)[:, 0] + pot(x)[:, 0] - c) .pow(2).mean()
                loss = task + cfg["lam"] * reg
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(params, 5.0); opt.step()
        net.eval()
        with torch.no_grad():
            metric = float(((net(ds["xte"].to(dev)) - ds["yte"].to(dev))**2).mean())
        if capture:
            return metric, net, pot, c, ds
        return metric
    except RuntimeError:
        # Explicit CPU fallback for constrained CUDA/cuDNN environments.
        if dev != "cpu":
            DEVICE = "cpu"
            return train_one(seed, cfg, mode, capture)
        raise

def c0(ell):
    return ell.mean().detach()  # pointwise constant baseline, no learned potential

def signature(net, pot, c, ds, seed=0):
    if pot is None: return {"confirmed": False, "reason": "no potential"}
    net.eval(); pot.eval(); dev = next(net.parameters()).device
    x = ds["xte"][:48].to(dev)
    rows = []
    for k in (1, 4, 8):
        z = x.clone(); s = torch.zeros(len(z), device=dev)
        for _ in range(k):
            z.requires_grad_(True); p = net(z)
            gg = torch.autograd.grad(p[:,0].sum(), z, create_graph=False)[0]
            ell = torch.log(gg[:, -3].abs() + 1e-3)
            s += ell
            z = make_next(z.detach(), p.detach())
        with torch.no_grad():
            r = s - (pot(z)[:,0] - pot(x)[:,0] + k*c)
            rows.append({"k": k, "observed_std_R_over_k": float((r/k).std()),
                         "observed_std_ell_minus_c": float((s/k-c).std())})
    vals = np.array([q["observed_std_R_over_k"] for q in rows])
    # Stage-1 prediction: coboundary residual should not grow linearly with horizon.
    confirmed = bool(vals[-1] <= 1.25 * vals[0] + 1e-8)
    return {"prediction": "finite-horizon residual per step remains bounded rather than accumulating linearly",
            "observed": rows, "confirmed": confirmed}

def main():
    # Same union is used by baseline sweep and idea settings (parity).
    grid = [{"lr": lr, "lam": lam, "epochs": 10} for lr in (1e-3, 3e-3, 1e-2) for lam in (1e-3, 1e-2)]
    cache = {}
    def maker(mode):
        def f(cfg):
            return lambda seed: train_one(seed, cfg, mode)
        return f
    base = sweep_baseline(maker("point"), grid)
    best = base["best_cfg"]
    # Best baseline config plus two nearby settings are all represented in grid.
    idea_grid = [best, {"lr": 1e-3, "lam": best["lam"], "epochs": 10},
                 {"lr": 1e-2, "lam": best["lam"], "epochs": 10}]
    idea_runs = []
    chosen = idea_grid[0]
    for cfg in idea_grid:
        r = evaluate(maker("coh")(cfg))
        idea_runs.append({"cfg": cfg, "result": r})
    chosen_run = min(idea_runs, key=lambda q: q["result"]["mean"])
    idea = chosen_run["result"]
    # Refit one paired seed for a mechanism signature from trained networks.
    m, n, p, cc, ds = train_one(0, chosen_run["cfg"], "coh", True)
    sig = signature(n, p, cc, ds)
    report = make_report(TRACK, MODEL, base, idea, {"mechanism_signature": sig,
        "idea_config_sweep": idea_runs, "device": DEVICE,
        "primary_metric": "test MSE"})
    report["baseline"]["method"] = "pointwise constant selected-coordinate log-Jacobian"
    report["idea"]["method"] = "cohomological potential residual"
    report["protocol_notes"] = "8 paired seeds; baseline sweep uses 4 seeds then full reevaluation; shared lr/lambda union"
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()