import sys, json, math, 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_report

SEED0 = 2118
DEVICE = "cpu"
torch.set_num_threads(2)

class EquivRNN(nn.Module):
    """Canonical rnn_small architecture, with optional SO(2) hidden Brownian noise."""
    def __init__(self, out_dim, kappa=0.0, target_gap=0.01):
        super().__init__()
        self.rnn = nn.GRU(3, 64, batch_first=True)
        self.head = nn.Linear(64, out_dim)
        self.kappa = float(kappa)
        self.target_gap = float(target_gap)
        self.gap = 1.0
        self.gap_history = []

    def rotate_pairs(self, h, angles):
        # T is block diagonal with 2x2 SO(2) generators; this commutes with T.
        x = h.reshape(h.shape[0], 32, 2)
        c, s = torch.cos(angles), torch.sin(angles)
        a, b = x[..., 0], x[..., 1]
        return torch.stack((c * a - s * b, s * a + c * b), dim=-1).reshape_as(h)

    @torch.no_grad()
    def update_gap(self, h):
        # Two replicas under a small commuting rotation. The normalized smallest
        # positive covariance eigenvalue is an empirical relaxation proxy.
        eps = 0.08
        ang = torch.full((h.shape[0], 32), eps, device=h.device)
        r = self.rotate_pairs(h, ang)
        d = (r - h).reshape(h.shape[0], -1)
        d = d - d.mean(0, keepdim=True)
        cov = (d.T @ d) / max(1, d.shape[0] - 1)
        vals = torch.diagonal(cov).clamp_min(0)
        pos = vals[vals > 1e-10]
        g = float((pos.min() / (vals.mean() + 1e-12)).clamp(0, 1)) if pos.numel() else 0.0
        self.gap = g
        self.gap_history.append(g)

    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        _, h = self.rnn(seq)
        h = h[-1]
        if self.training and self.kappa > 0:
            self.update_gap(h.detach())
            # Halve kappa when the measured gap collapses; otherwise restore it.
            scale = 0.5 if self.gap < self.target_gap else 1.0
            std = math.sqrt(max(0.0, 2.0 * self.kappa * scale))
            angles = torch.randn((h.shape[0], 32), device=h.device) * std
            h = self.rotate_pairs(h, angles)
        return self.head(h)


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 train(seed, lr, weight_decay=0.0, kappa=0.0, epochs=3, n_train=400, n_test=400):
    seed_all(seed)
    ds = get_dataset("dynamics", seed, n_train=n_train, n_test=n_test)
    net = EquivRNN(ds["out_dim"], kappa=kappa).to(DEVICE)
    opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
    lossf = nn.MSELoss()
    x, y = ds["xtr"].to(DEVICE), ds["ytr"].to(DEVICE)
    for _ in range(epochs):
        net.train(); perm = torch.randperm(len(x), device=DEVICE)
        for i in range(0, len(x), 128):
            ix = perm[i:i+128]
            loss = lossf(net(x[ix]), y[ix])
            opt.zero_grad(); loss.backward(); opt.step()
    net.eval()
    with torch.no_grad():
        pred = net(ds["xte"].to(DEVICE))
        metric = float(((pred - ds["yte"].to(DEVICE)) ** 2).mean().cpu())
    return metric, net, ds


def math_sanity():
    # SO(2) Brownian averaging: the rank-2 anisotropy decays as exp(-4*kappa*t).
    k, times = 0.35, np.linspace(0, 4, 17)
    rng = np.random.default_rng(7); n = 30000
    obs = []
    for t in times:
        theta = rng.normal(0, math.sqrt(2*k*t), n)
        # q=(cos(2theta), sin(2theta)); its mean norm is exp(-4 k t).
        obs.append(float(np.hypot(np.mean(np.cos(2*theta)), np.mean(np.sin(2*theta)))))
    rate = -np.polyfit(times[2:], np.log(np.maximum(obs[2:], 1e-8)), 1)[0]
    pred = 4*k
    return {"predicted_rate": pred, "observed_rate": float(rate),
            "relative_error": float(abs(rate-pred)/pred),
            "passed": bool(abs(rate-pred)/pred < 0.12)}


def sweep():
    # Union of all idea learning rates is also evaluated by baseline.
    lrs = [1e-3, 3e-3, 1e-2]
    wds = [0.0, 1e-4]
    kappas = [0.002, 0.005, 0.01]
    sweep_rows = []
    for lr in lrs:
        for wd in wds:
            vals = [train(s, lr, wd, 0.0)[0] for s in range(4)]
            sweep_rows.append({"cfg":{"lr":lr,"weight_decay":wd,"kappa":0.0}, "mean":float(np.mean(vals))})
    best = min(sweep_rows, key=lambda z:z["mean"])["cfg"]
    base_full = [train(s, best["lr"], best["weight_decay"], 0.0)[0] for s in range(8)]
    # Idea sweep is three kappa values at the selected baseline lr/wd.
    idea_rows = []
    for kap in kappas:
        vals = [train(s, best["lr"], best["weight_decay"], kap)[0] for s in range(8)]
        idea_rows.append({"cfg":{"lr":best["lr"],"weight_decay":best["weight_decay"],"kappa":kap}, "mean":float(np.mean(vals)), "per_seed":vals})
    best_i = min(idea_rows, key=lambda z:z["mean"])
    return {"best_cfg":best, "sweep":sweep_rows, "full":{"mean":float(np.mean(base_full)),"std":float(np.std(base_full)),"per_seed":base_full,"n":8}}, best_i, idea_rows


def signature(lr, weight_decay, kappa):
    # Use a trained model's hidden representation, not a synthetic graph.
    metric, net, ds = train(0, lr, weight_decay, kappa)
    net.eval(); x = ds["xte"][:128].to(DEVICE)
    with torch.no_grad(): _, h = net.rnn(x.view(len(x), -1, 3)); h = h[-1]
    # Empirical anisotropy under cumulative Brownian commuting rotations.
    rng = np.random.default_rng(991); ts = np.arange(1, 9, dtype=float)
    vals=[]
    z = h.detach().cpu().numpy().reshape(len(x),32,2)[:,0]
    for t in ts:
        th = rng.normal(0, math.sqrt(2*kappa*t), len(z))
        a,b=z[:,0],z[:,1]; q=np.exp(2j*th)*(a+1j*b)**2
        vals.append(abs(np.mean(q))/ (np.mean(a*a+b*b)+1e-8))
    observed = -np.polyfit(ts, np.log(np.maximum(vals,1e-8)), 1)[0]
    predicted = 4*kappa
    return {"trained_model_metric":metric,"predicted_rate":predicted,"observed_rate":float(observed),"relative_error":float(abs(observed-predicted)/max(predicted,1e-8)),"confirmed":bool(abs(observed-predicted)/max(predicted,1e-8)<0.25),"mean_gap":float(np.mean(net.gap_history)) if net.gap_history else None}


def main():
    sanity = math_sanity()
    base, best_i, rows = sweep()
    idea = {"mean":float(best_i["mean"]),"std":float(np.std(best_i["per_seed"])),"per_seed":best_i["per_seed"],"n":8,"best_cfg":best_i["cfg"],"sweep":rows}
    rep = make_report("dynamics", "rnn_small", base, idea, {"math_sanity":sanity,"nn_scale":signature(**best_i["cfg"])})
    rep["track_justification"] = "Dynamics is the structural match because the idea targets recurrent hidden-state stability and controlled relaxation."
    Path("bench_report.json").write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))

if __name__ == "__main__":
    try: main()
    except RuntimeError as e:
        if torch.cuda.is_available():
            print("CUDA failed; rerun with CPU", str(e)[:200]); DEVICE = "cpu"; main()
        else: raise