import json, sys, math
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, evaluate, sweep_baseline, make_report

SEED0 = 1270
N_TRAIN, N_TEST = 1200, 400
EPOCHS, BATCH = 18, 128
LR_GRID = [1e-3, 3e-3, 1e-2]
LAMBDA_PATCH = 0.01
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

class GaugeRNN(nn.Module):
    """rnn_small-compatible GRU with two local hidden patches and SO(2)^32 gauges."""
    def __init__(self, hidden=64):
        super().__init__()
        self.rnn = nn.GRU(3, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)
        # One angle per 2D coordinate plane; exp(skew(angle)) is exactly orthogonal.
        self.angles = nn.Parameter(torch.zeros(hidden // 2))

    def forward(self, x, return_aux=False):
        seq = x.view(x.shape[0], -1, 3)
        _, hfull = self.rnn(seq)
        pred = self.head(hfull[-1])
        if not return_aux:
            return pred
        mid = seq.shape[1] // 2
        _, h0 = self.rnn(seq[:, :mid])
        _, h1 = self.rnn(seq[:, mid:])
        a = self.angles
        c, s = torch.cos(a), torch.sin(a)
        z0, z1 = h0[-1].view(h0[-1].shape[0], self.angles.numel(), 2), h1[-1].view(h1[-1].shape[0], self.angles.numel(), 2)
        # R(a) @ [x,y] = [c*x-s*y, s*x+c*y]
        zg = torch.stack((c[None] * z1[..., 0] - s[None] * z1[..., 1],
                          s[None] * z1[..., 0] + c[None] * z1[..., 1]), dim=-1)
        aligned = zg.reshape_as(h1[-1])
        return pred, h0[-1], h1[-1], aligned

def patch_loss(model, x):
    _, h0, h1, aligned = model(x, return_aux=True)
    return ((h0 - aligned) ** 2).mean()

def train_one(seed, lr, lam):
    global DEVICE
    torch.manual_seed(seed); np.random.seed(seed)
    try:
        ds = get_dataset("dynamics", seed, n_train=N_TRAIN, n_test=N_TEST)
        model = GaugeRNN().to(DEVICE)
        opt = torch.optim.Adam(model.parameters(), lr=lr)
        xtr, ytr = ds["xtr"].to(DEVICE), ds["ytr"].to(DEVICE)
        model.train()
        gen = torch.Generator(device="cpu").manual_seed(seed + 91)
        for ep in range(EPOCHS):
            order = torch.randperm(len(xtr), generator=gen)
            for ix in order.split(BATCH):
                xb, yb = xtr[ix], ytr[ix]
                pred = model(xb)
                task = ((pred - yb) ** 2).mean()
                loss = task if lam == 0.0 else task + lam * patch_loss(model, xb)
                opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
        model.eval()
        with torch.no_grad():
            pred, h0, h1, aligned = model(ds["xte"].to(DEVICE), return_aux=True)
            metric = ((pred - ds["yte"].to(DEVICE)) ** 2).mean().item()
            disagreement = ((h0 - aligned) ** 2).mean().item()
            norms = torch.linalg.vector_norm(aligned, dim=1) - torch.linalg.vector_norm(h1, dim=1)
            norm_err = norms.abs().max().item()
            angle_mag = model.angles.abs().mean().item()
        return metric, {"disagreement": disagreement, "norm_error": norm_err, "angle_mean_abs": angle_mag}
    except RuntimeError as exc:
        # CUDA failures get one deterministic CPU retry; programming errors propagate.
        msg = str(exc).lower()
        if DEVICE != "cuda" or not any(k in msg for k in ("cuda", "cudnn", "out of memory")):
            raise
        old = DEVICE; DEVICE = "cpu"
        try: return train_one(seed, lr, lam)
        finally: DEVICE = old

def metric_fn(lam, lr):
    def f(seed): return train_one(seed, lr, lam)[0]
    return f

def main():
    # Cheap numerical verification before training: block rotations preserve norms.
    torch.manual_seed(SEED0)
    a = torch.randn(17, 32); theta = torch.linspace(-2, 2, 16)
    c, s = torch.cos(theta), torch.sin(theta)
    b = torch.stack((c[None] * a[:, 0::2] - s[None] * a[:, 1::2],
                     s[None] * a[:, 0::2] + c[None] * a[:, 1::2]), -1).reshape_as(a)
    math_check = {"max_norm_error": float((a.norm(dim=1)-b.norm(dim=1)).abs().max()),
                  "predicted": "orthogonal gauge preserves hidden norm"}

    # Baseline sweep uses every learning rate also tried by the idea.
    base = sweep_baseline(lambda cfg: metric_fn(0.0, cfg["lr"]),
                          [{"lr": x} for x in LR_GRID])
    idea_runs = []
    for lr in LR_GRID:
        r = evaluate(metric_fn(LAMBDA_PATCH, lr))
        idea_runs.append({"cfg": {"lr": lr, "lambda_patch": LAMBDA_PATCH}, "result": r})
    best = min(idea_runs, key=lambda z: z["result"]["mean"])
    idea = best["result"]

    # Re-test trained models on all paired seeds for mechanism signature.
    base_beh, idea_beh = [], []
    for s in range(8):
        bm, bx = train_one(s, base["best_cfg"]["lr"], 0.0)
        im, ix = train_one(s, best["cfg"]["lr"], LAMBDA_PATCH)
        base_beh.append({"metric": bm, **bx}); idea_beh.append({"metric": im, **ix})
    bdisc = float(np.mean([z["disagreement"] for z in base_beh]))
    idisc = float(np.mean([z["disagreement"] for z in idea_beh]))
    inorm = float(max(z["norm_error"] for z in idea_beh))
    signature = {"prediction": "orthogonal transitions preserve hidden norms and patch loss reduces disagreement",
                 "predicted_norm_error": 0.0, "observed_max_norm_error": inorm,
                 "predicted_disagreement_change": "negative", "observed_baseline_disagreement": bdisc,
                 "observed_idea_disagreement": idisc,
                 "confirmed": bool(inorm < 1e-5 and idisc < bdisc),
                 "math_sanity": math_check, "trained_model_behavior": {"baseline": base_beh, "idea": idea_beh}}
    rep = make_report("dynamics", "rnn_small", base, idea,
                      {**signature, "idea_sweep": idea_runs,
                       "track_rationale": "Dynamics is the built-in stability/control track; local temporal patches are neighboring state windows."})
    Path("bench_report.json").write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))

if __name__ == "__main__": main()