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

SEEDS = tuple(range(8))
# Shared union: every idea learning rate is also evaluated by baseline sweep.
GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
EPOCHS = 12
NTRAIN, NTEST = 400, 400


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)


class QuotientFibreGRU(nn.Module):
    """64-state triangular recurrent model: autonomous quotient and fibre."""
    def __init__(self, out_dim=1, qdim=32, fdim=42):
        super().__init__()
        self.qdim, self.fdim = qdim, fdim
        self.qcell = nn.GRUCell(3, qdim)
        self.fcell = nn.GRUCell(3 + qdim, fdim)
        self.head = nn.Linear(qdim + fdim, out_dim)

    def forward(self, x, return_states=False):
        seq = x.view(x.shape[0], -1, 3)
        z = x.new_zeros((x.shape[0], self.qdim))
        y = x.new_zeros((x.shape[0], self.fdim))
        zs, ys = [], []
        for t in range(seq.shape[1]):
            u = seq[:, t]
            z = self.qcell(u, z)
            # Conditional fibre transition uses current quotient, never y to update z.
            y = self.fcell(torch.cat([u, z], dim=-1), y)
            zs.append(z); ys.append(y)
        out = self.head(torch.cat([z, y], dim=-1))
        if return_states:
            return out, torch.stack(zs, 1), torch.stack(ys, 1)
        return out


def train_one(kind, seed, lr):
    seed_all(seed)
    ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
    if kind == "baseline":
        model = make_model("rnn_small", ds["xtr"].shape[1:], ds["ytr"].shape[-1] if ds["ytr"].ndim > 1 else 1)
    else:
        model = QuotientFibreGRU(out_dim=1)
    _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
    return float(metric)


def model_for_signature(kind, seed, lr):
    seed_all(seed)
    ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
    model = (make_model("rnn_small", ds["xtr"].shape[1:], 1)
             if kind == "baseline" else QuotientFibreGRU(1))
    model, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
    return model, ds, metric


def fitted_rate(values, start=1):
    a = np.maximum(np.asarray(values, dtype=float), 1e-8)
    if len(a) <= start + 2: return float("nan")
    slope = np.polyfit(np.arange(start, len(a)), np.log(a[start:]), 1)[0]
    return float(np.exp(slope))


def mechanism_signature(best_lr):
    """Re-test max-rate and slow-branch invariance on trained benchmark models."""
    model, ds, _ = model_for_signature("idea", 0, best_lr)
    model.eval()
    dev = next(model.parameters()).device
    x = ds["xte"][:32].clone().to(dev)
    # Empirical perturbation of initial hidden states, with identical observed inputs.
    with torch.no_grad():
        _, z1, y1 = model(x, True)
        # Since the canonical model starts at zero, rerun with perturbed branch initial
        # states through the same trained cells to measure each branch's decay.
        seq = x.view(x.shape[0], -1, 3)
        z = torch.zeros((len(x), model.qdim), device=dev); y = torch.zeros((len(x), model.fdim), device=dev)
        z2 = torch.zeros_like(z); y2 = torch.zeros_like(y)
        z2[:, 0] = 1.0; y2[:, 0] = 1.0
        qnorm, fnorm, fullnorm = [], [], []
        for t in range(seq.shape[1]):
            u = seq[:, t]
            z = model.qcell(u, z); y = model.fcell(torch.cat([u, z], -1), y)
            z2 = model.qcell(u, z2); y2 = model.fcell(torch.cat([u, z2], -1), y2)
            qnorm.append((z2-z).norm(dim=1).mean().item())
            fnorm.append((y2-y).norm(dim=1).mean().item())
            fullnorm.append(torch.sqrt((z2-z).pow(2).sum(1)+(y2-y).pow(2).sum(1)).mean().item())
    aq, af, observed = fitted_rate(qnorm), fitted_rate(fnorm), fitted_rate(fullnorm)
    pred = max(aq, af)
    rel = abs(observed-pred) / max(abs(pred), 1e-8)
    return {"predicted_max_rate": pred, "observed_full_rate": observed,
            "observed_quotient_rate": aq, "observed_fibre_rate": af,
            "relative_error": rel, "confirmed": bool(np.isfinite(rel) and rel <= .20),
            "measurement": "trained quotient-fibre GRU hidden-state perturbation decay"}


def main():
    # Baseline sweep on four seeds, then full eight-seed reevaluation by harness.
    base = sweep_baseline(lambda cfg: (lambda s: train_one("baseline", s, cfg["lr"])), GRID)
    best_lr = float(base["best_cfg"]["lr"])
    idea_cfgs = [best_lr] + [float(g["lr"]) for g in GRID if float(g["lr"]) != best_lr]
    idea_runs = []
    for lr in idea_cfgs:
        r = evaluate(lambda s, lr=lr: train_one("idea", s, lr), seeds=SEEDS)
        idea_runs.append({"lr": lr, "result": r})
    idea_best = min(idea_runs, key=lambda q: q["result"]["mean"])
    report = make_report("dynamics", "rnn_small", base, idea_best["result"], {
        "predicted_vs_observed": mechanism_signature(idea_best["lr"]),
        "idea_lr_sweep": idea_runs,
        "parameter_counts": {"baseline": count_params(make_model("rnn_small", (24,), 1)), "idea": count_params(QuotientFibreGRU(1))},
        "track_choice": "dynamics structurally matches stability/control and multi-step dynamical memory"
    })
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__":
    main()