import json
import math
import sys
from pathlib import Path

import numpy as np
import torch
from scipy.special import roots_genlaguerre

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

SEED = 2977
N_FEATURES = 64
EPOCHS = 12
NTR, NTE = 1200, 300
LR_GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]


def quadrature_frequencies(d=10, n_features=N_FEATURES):
    """Product quadrature: generalized-Laguerre radial rule x fixed sphere rule.
    If r^2/2 ~ Gamma(d/2,1), Laguerre integrates its radial density; directions
    are a fixed antipodal quasi-uniform set, giving a deterministic Gaussian
    spectral rule with nonnegative equal weights.
    """
    ndir, nq = 16, n_features // 16
    t, wt = roots_genlaguerre(nq, d / 2.0 - 1.0)
    wt = wt / math.gamma(d / 2.0)
    rng = np.random.RandomState(18431)
    u = rng.normal(size=(ndir // 2, d))
    u /= np.linalg.norm(u, axis=1, keepdims=True)
    dirs = np.concatenate([u, -u], axis=0)
    # Pairing removes odd angular bias and is deterministic once constructed.
    W = np.concatenate([np.sqrt(2.0 * r) * dirs[j:j+1]
                        for r in t for j in range(ndir)], axis=0)
    a = np.repeat(wt / ndir, ndir)
    return W.astype(np.float32), a.astype(np.float32)


def random_frequencies(d=10, n_features=N_FEATURES, seed=0):
    return np.random.RandomState(seed).normal(size=(n_features, d)).astype(np.float32), \
           np.full(n_features, 1.0 / n_features, dtype=np.float32)


class FourierMLP(torch.nn.Module):
    def __init__(self, W, weights):
        super().__init__()
        self.register_buffer("W", torch.as_tensor(W))
        self.register_buffer("sqrt_a", torch.sqrt(torch.as_tensor(weights)))
        self.net = torch.nn.Sequential(
            torch.nn.Linear(2 * len(W), 64), torch.nn.ReLU(),
            torch.nn.Linear(64, 64), torch.nn.ReLU(), torch.nn.Linear(64, 1))

    def embedding(self, x):
        z = x @ self.W.T
        f = torch.cat((torch.cos(z), torch.sin(z)), dim=1)
        return f * self.sqrt_a.repeat(2).unsqueeze(0)

    def forward(self, x):
        return self.net(self.embedding(x))


def prep(seed):
    d = get_dataset("tabular", seed, n_train=NTR, n_test=NTE)
    # Standardization is shared preprocessing and prevents frequency scale from
    # being dominated by a coordinate's units.
    mu, sd = d["xtr"].mean(0, keepdim=True), d["xtr"].std(0, keepdim=True).clamp_min(1e-5)
    for k in ("xtr", "xte"):
        d[k] = (d[k] - mu) / sd
    return d


def train_value(kind, cfg, seed, return_model=False):
    torch.manual_seed(SEED + 100 * seed + (0 if kind == "baseline" else 10000))
    d = prep(seed)
    if kind == "baseline":
        W, a = random_frequencies(seed=SEED + seed)
    else:
        W, a = quadrature_frequencies()
    model, metric, _ = train_model(FourierMLP(W, a), d, epochs=EPOCHS,
                                   lr=cfg["lr"], batch=128, weight_decay=0.0,
                                   log=lambda *_: None)
    if return_model:
        return metric, model, d, W, a
    return metric


def make_train(kind):
    return lambda cfg: (lambda seed: train_value(kind, cfg, seed))


def spectral_error(model, d, W, a):
    with torch.no_grad():
        x = d["xte"][:120]
        z = x @ torch.as_tensor(W).T
        F = torch.cat((torch.cos(z), torch.sin(z)), 1) * torch.sqrt(torch.as_tensor(a)).repeat(2).unsqueeze(0)
        Kh = F @ F.T
        dist = ((x[:, None, :] - x[None, :, :]) ** 2).sum(-1)
        K = torch.exp(-0.5 * dist)
        le = torch.linalg.eigvalsh(K).flip(0)[:10]
        lh = torch.linalg.eigvalsh(Kh).flip(0)[:10]
        return float(torch.mean(torch.abs(le-lh) / le.clamp_min(1e-7)))


def main():
    # Baseline sweep and full evaluation are performed by the canonical protocol.
    base = sweep_baseline(make_train("baseline"), LR_GRID)
    # Same union of learning rates on idea side; select its best on the same
    # four tuning seeds, then evaluate that setting on all eight paired seeds.
    idea_trials = []
    for cfg in LR_GRID:
        vals = [train_value("idea", cfg, s) for s in range(4)]
        idea_trials.append({"cfg": cfg, "mean": float(np.mean(vals))})
    best_cfg = min(idea_trials, key=lambda x: x["mean"])["cfg"]
    idea = {"best_cfg": best_cfg, "sweep": idea_trials,
            "full": __import__("bench").protocol.evaluate(make_train("idea")(best_cfg))}

    # Re-test the claimed spectral effect on the actual trained systems.
    _, bmodel, bd, bW, ba = train_value("baseline", best_cfg, 0, True)
    _, imodel, idd, iW, ia = train_value("idea", best_cfg, 0, True)
    bspec = spectral_error(bmodel, bd, bW, ba)
    ispec = spectral_error(imodel, idd, iW, ia)
    sig = {"prediction": "deterministic quadrature has lower top-10 RBF Gram eigenvalue error",
           "predicted_baseline_top10_relative_error": bspec,
           "predicted_idea_top10_relative_error": ispec,
           "observed_delta_idea_minus_baseline": ispec - bspec,
           "confirmed": bool(ispec < bspec)}
    report = make_report("tabular", "fourier_mlp_shared", base, idea["full"],
                         {"mechanism_signature": sig,
                          "notes": "Tabular is the matched built-in track for an input embedding/kernel feature intervention."})
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()