import os, sys, json
import numpy as np
import torch
import torch.nn as nn

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

TRACK = "dynamics"
EPOCHS = 18
LRS = [1e-3, 3e-3, 1e-2]
D = 4


def fixed_ops():
    # Hermitian Pauli interactions on two latent qubits; X and Z terms do not commute.
    I = torch.eye(2, dtype=torch.float32)
    X = torch.tensor([[0., 1.], [1., 0.]])
    Y = torch.tensor([[0., -1.], [1., 0.]])  # real representation of iY
    Z = torch.tensor([[1., 0.], [0., -1.]])
    kron = torch.kron
    return torch.stack([kron(X, I), kron(Z, X), kron(Y, Y), kron(Z, Z)])


class MatchedRNN(nn.Module):
    def __init__(self, idea, seed):
        super().__init__()
        # This is the bench rnn_small backbone, kept identical between systems.
        self.rnn = nn.GRU(3, 64, batch_first=True)
        self.proj = nn.Linear(64, D, bias=False)
        g = torch.Generator().manual_seed(10000 + int(seed))
        with torch.no_grad():
            self.proj.weight.copy_(torch.randn(D, 64, generator=g) / np.sqrt(64))
        for p in self.proj.parameters():
            p.requires_grad_(False)
        self.idea = idea
        if idea:
            self.theta = nn.Parameter(torch.randn(D, generator=g) * 0.15)
            self.register_buffer("ops", fixed_ops())
        else:
            self.readout = nn.Linear(D, 1, bias=False)

    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        _, h = self.rnn(seq)
        z = self.proj(h[-1])
        if not self.idea:
            return self.readout(z)
        # rho is a pure-state density matrix made from the shared latent feature.
        q = z / (torch.linalg.vector_norm(z, dim=1, keepdim=True) + 1e-6)
        rho = q.unsqueeze(2) * q.unsqueeze(1)
        H = torch.einsum("j,jab->ab", self.theta, self.ops)
        lam, U = torch.linalg.eigh(H)
        A = (U * torch.tanh(lam).unsqueeze(0)) @ U.T
        pred = torch.einsum("bij,ji->b", rho, A)
        return pred.unsqueeze(1)


def run_one(idea, lr, seed, keep_model=False):
    torch.manual_seed(7000 + int(seed))
    np.random.seed(7000 + int(seed))
    ds = get_dataset(TRACK, int(seed), n_train=400, n_test=400)
    model = MatchedRNN(idea, seed)
    net, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
    if net is None:
        raise RuntimeError("bench training failed")
    return (float(metric), net, ds) if keep_model else float(metric)


def make_train(idea, cfg=None):
    cfg = cfg or {"lr": 3e-3, "epochs": EPOCHS}
    return lambda seed: run_one(idea, float(cfg["lr"]), seed)


def behavior_signature(idea_res):
    # Re-test the claimed nonlinear expressivity on predictions of trained systems.
    vals = []
    nonlinear = []
    comm = fixed_ops()
    comm_norm = float(torch.linalg.matrix_norm(comm[0] @ comm[1] - comm[1] @ comm[0]))
    for seed in range(8):
        metric, net, ds = run_one(True, idea_res["cfg"]["lr"], seed, True)
        net = net.to("cpu")
        with torch.no_grad():
            p = net(ds["xte"]).squeeze().numpy()
            # Shared latent coordinates are reconstructed from the trained backbone.
            seq = ds["xte"].view(len(ds["xte"]), -1, 3)
            _, h = net.rnn(seq)
            z = net.proj(h[-1]).numpy()
        X = np.column_stack([np.ones(len(z)), z])
        fit = X @ np.linalg.lstsq(X, p, rcond=None)[0]
        nl = float(np.sqrt(np.mean((p - fit) ** 2)))
        vals.append(float(metric)); nonlinear.append(nl)
    return {
        "prediction": "noncommuting Hamiltonian terms should create a measurable nonlinear readout",
        "commutator_frobenius": comm_norm,
        "trained_idea_test_mse_mean": float(np.mean(vals)),
        "trained_idea_nonlinear_residual_mean": float(np.mean(nonlinear)),
        "confirmed": bool(comm_norm > 1e-6 and np.mean(nonlinear) > 1e-4),
    }


def main():
    # Baseline and idea use the same union of step sizes, satisfying search parity.
    grid = [{"lr": lr, "epochs": EPOCHS} for lr in LRS]
    base = sweep_baseline(lambda cfg: make_train(False, cfg), grid)
    idea_trials = []
    best = None
    for cfg in grid:
        r = evaluate(make_train(True, cfg))
        idea_trials.append({"cfg": cfg, **r})
        if best is None or r["mean"] < best["mean"]:
            best = {"cfg": cfg, **r}
    sig = behavior_signature(best)
    report = make_report(TRACK, "rnn_small_matched_spectral_head", base, best,
                         {"mechanism_signature": sig,
                          "idea_sweep": idea_trials,
                          "structural_match": "dynamics: controlled pendulum rollout"})
    report["budget"] = {"epochs": EPOCHS, "n_train": 400, "n_test": 400, "seeds": 8}
    with open("bench_report.json", "w") as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()