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

import sys
sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report

SEED = 2841
EPOCHS = 12
NTRAIN, NTEST = 1200, 400
LR_GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]

class OrbitalRNN(nn.Module):
    """Phase-only recurrent replacement for bench rnn_small.

    The directed ring plus two-step skip graph is fixed and positive. Input
    drive, common frequency, and readout are learned end-to-end; amplitudes
    are fixed at one, as in the MVP implementation plan.
    """
    def __init__(self, hidden=64, dt=0.15, kappa=0.8):
        super().__init__()
        self.hidden, self.dt, self.kappa = hidden, dt, kappa
        self.inp = nn.Linear(3, hidden, bias=False)
        self.omega = nn.Parameter(torch.randn(hidden) * 0.05 + 0.12)
        A = torch.zeros(hidden, hidden)
        for i in range(hidden):
            A[i, (i + 1) % hidden] = 0.35
            A[i, (i + 2) % hidden] = 0.15
        self.register_buffer("A", A)
        self.head = nn.Linear(2 * hidden, 1)

    def phase_step(self, theta, drive):
        # d_ij = theta_j - theta_i; A[i,j] is directed i <- j.
        diff = theta[:, None, :] - theta[:, :, None]
        coupling = torch.einsum("bij,ij->bi", torch.sin(diff), self.A)
        return theta + self.dt * (self.omega + 0.08 * drive + self.kappa * coupling)

    def forward_with_theta(self, x, theta0=None, return_states=False):
        b = x.shape[0]
        seq = x.view(b, -1, 3)
        theta = torch.zeros(b, self.hidden, device=x.device) if theta0 is None else theta0
        states = []
        for k in range(seq.shape[1]):
            theta = self.phase_step(theta, self.inp(seq[:, k]))
            states.append(torch.cat((torch.cos(theta), torch.sin(theta)), dim=1))
        out = self.head(states[-1])
        return (out, torch.stack(states, dim=1), theta) if return_states else out

    def forward(self, x):
        return self.forward_with_theta(x)


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_one(kind, lr, seed, keep_model=False):
    seed_all(seed)
    ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
    model = make_model("rnn_small", ds["input_shape"], ds["out_dim"]) if kind == "baseline" else OrbitalRNN()
    try:
        _, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
    except Exception:
        # train_model already has a CUDA fallback; this protects unusual driver errors.
        model = model.cpu(); ds = {k: (v.cpu() if torch.is_tensor(v) else v) for k, v in ds.items()}
        _, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
    if metric is None: raise RuntimeError("benchmark training failed")
    return (float(metric), model, ds) if keep_model else float(metric)


def baseline_factory(cfg):
    return lambda seed: train_one("baseline", float(cfg["lr"]), seed)

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


def mechanism_signature(model, ds):
    """Retest orbital prediction on trajectories produced by the trained model.

    At each observed phase state, compute the transverse Jacobian spectral
    abscissa and compare it with finite-perturbation decay of the trained cell.
    """
    device = next(model.parameters()).device
    x = ds["xte"][:64].to(device)
    model.eval()
    with torch.no_grad():
        _, _, theta = model.forward_with_theta(x, return_states=True)
    th = theta[0].detach().cpu().numpy()
    A = model.A.detach().cpu().numpy(); k = model.kappa
    J = np.zeros((model.hidden, model.hidden))
    for i in range(model.hidden):
        for j in range(model.hidden):
            if i != j:
                J[i, j] = k * A[i, j] * math.cos(float(th[j] - th[i]))
        J[i, i] = -np.sum([k * A[i, j] * math.cos(float(th[j] - th[i])) for j in range(model.hidden) if j != i])
    eig = np.linalg.eigvals(J)
    transverse = [z.real for z in eig if abs(z) > 1e-6]
    predicted = float(max(transverse)) if transverse else 0.0
    # Empirical cell perturbation on the same trained model and observed input.
    xone = x[:1]
    theta0 = torch.zeros(1, model.hidden, device=device)
    d = torch.linspace(-1, 1, model.hidden, device=device).unsqueeze(0)
    d = d - d.mean(); d = 1e-5 * d / d.norm()
    norms = []
    with torch.no_grad():
        for t in range(8):
            theta0 = model.phase_step(theta0, model.inp(xone[:, t*3:(t+1)*3]))
            # Compare a perturbed parallel rollout at this same input step.
            # Recompute from initial perturbation for a direct finite response.
            if t == 0: pert = d.clone()
            pert = model.phase_step(pert, model.inp(xone[:, t*3:(t+1)*3]))
            norms.append(float((pert - theta0).norm().cpu()))
    # fit log decay after first point; positive slope means instability.
    observed = float(np.polyfit(np.arange(len(norms))[1:], np.log(np.maximum(norms, 1e-30))[1:], 1)[0])
    # Euler predicts log multiplier log(1 + dt*alpha)/dt.
    predicted_euler = math.log(max(1e-8, 1.0 + model.dt * predicted)) / model.dt
    confirmed = bool(abs(observed - predicted_euler) < 0.35 and predicted < 0)
    return {"predicted_transverse_alpha": predicted, "predicted_euler_slope": predicted_euler,
            "observed_perturbation_slope": observed, "max_edge_angle_rad": float(np.max(np.abs(((th[:,None]-th[None,:]+np.pi)%(2*np.pi))-np.pi))),
            "confirmed": confirmed}


def main():
    torch.set_num_threads(4)
    # Same union of method step sizes on both sides; baseline is selected on 4 seeds.
    base = sweep_baseline(baseline_factory, LR_GRID)
    best_lr = float(base["best_cfg"]["lr"])
    idea_grid = LR_GRID
    # Explicitly evaluate all three idea settings, then report the best setting.
    idea_runs = []
    for cfg in idea_grid:
        r = evaluate(idea_factory(cfg))
        idea_runs.append({"cfg": cfg, "result": r})
    best_idea = min(idea_runs, key=lambda z: z["result"]["mean"])
    # Retain the trained best model on seed 0 for behavior signature.
    _, trained, ds = train_one("idea", float(best_idea["cfg"]["lr"]), 0, keep_model=True)
    report = make_report("dynamics", "rnn_small", base, best_idea["result"], {
        "prediction": "stable transverse perturbations should contract with slope near log(1+dt*alpha)/dt",
        "trained_model": "orbital phase RNN, seed 0, best benchmark lr",
        **mechanism_signature(trained, ds)
    })
    report["idea_sweep"] = idea_runs
    report["protocol"] = {"epochs": EPOCHS, "n_train": NTRAIN, "n_test": NTEST,
                           "paired_seeds": list(range(8)), "lr_union": [c["lr"] for c in LR_GRID],
                           "baseline_best_lr": best_lr, "idea_best_lr": best_idea["cfg"]["lr"]}
    with open("bench_report.json", "w") as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()