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

SEEDS = tuple(range(8))
# This union is used by both baseline and idea, satisfying step-size parity.
GRID = [
    {"lr": 1e-3, "epochs": 12},
    {"lr": 3e-3, "epochs": 12},
    {"lr": 1e-2, "epochs": 12},
]
EPSILON = 0.30
CONSTANT = 1.0


def jl_width(d, n, epsilon=EPSILON, constant=CONSTANT):
    raw = min(d, n - 1, math.log(2.0 + epsilon * epsilon * n) / (epsilon * epsilon))
    r = max(1, min(d, int(math.ceil(constant * raw))))
    return r, raw


def seed_all(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        try:
            torch.cuda.manual_seed_all(seed)
        except Exception:
            pass


class JLTransformer(nn.Module):
    """The bench transformer with a frozen Gaussian projection after embedding."""
    def __init__(self, win=32, epsilon=EPSILON, constant=CONSTANT):
        super().__init__()
        self.win = win
        self.d = 64
        self.r, self.raw_width = jl_width(self.d, win, epsilon, constant)
        self.inp = nn.Linear(1, self.d)
        gen = torch.Generator().manual_seed(1907 + self.r)
        p = torch.randn(self.r, self.d, generator=gen) / math.sqrt(self.r)
        self.register_buffer("P", p)
        self.pos = nn.Parameter(torch.zeros(1, win, self.r))
        nn.init.normal_(self.pos, std=0.02)
        # The shared transformer block is unchanged; only its width is reduced.
        layer = nn.TransformerEncoderLayer(
            self.r, nhead=2, dim_feedforward=128, batch_first=True, dropout=0.0
        )
        self.enc = nn.TransformerEncoder(layer, 2)
        self.head = nn.Linear(win * self.r, 1)

    def forward(self, x):
        h = self.inp(x.unsqueeze(-1))
        h = torch.matmul(h, self.P.t()) + self.pos[:, :x.shape[1]]
        return self.head(self.enc(h).reshape(x.shape[0], -1))


def run_one(kind, cfg, seed, keep=False):
    seed_all(seed)
    ds = get_dataset("sequence", int(seed), n_train=400, n_test=200)
    if kind == "baseline":
        net = make_model("transformer_tiny", ds["input_shape"], ds["out_dim"])
    else:
        net = JLTransformer(ds["input_shape"][0], EPSILON, CONSTANT)
    net, metric, history = train_model(
        net, ds, epochs=cfg["epochs"], lr=cfg["lr"], batch=128
    )
    if net is None:
        raise RuntimeError("bench training returned no model")
    if keep:
        return float(metric), net, ds, history
    return float(metric)


def train_fn(kind, cfg):
    return lambda seed: run_one(kind, cfg, seed)


def model_signature():
    """Measure the JL prediction on representations produced by a trained model."""
    metric, net, ds, _ = run_one("idea", GRID[1], 0, keep=True)
    net.eval()
    device = next(net.parameters()).device
    x = torch.as_tensor(ds["xte"][:64], dtype=torch.float32, device=device)
    with torch.no_grad():
        h = net.inp(x.unsqueeze(-1))
        z = torch.matmul(h, net.P.t())
        a = h.reshape(-1, net.d)
        b = z.reshape(-1, net.r)
        # Fixed paired token distances, avoiding zero denominators.
        dif_a = a[1:] - a[:-1]
        dif_b = b[1:] - b[:-1]
        da = torch.linalg.vector_norm(dif_a, dim=1)
        db = torch.linalg.vector_norm(dif_b, dim=1)
        keep = da > 1e-7
        ratios = (db[keep] / da[keep]).cpu().numpy()
    observed_mean = float(np.mean(ratios))
    observed_sd = float(np.std(ratios, ddof=1))
    predicted_sd = 1.0 / math.sqrt(2.0 * net.r)
    # Honest tolerance accounts for finite samples and non-isotropic trained states.
    confirmed = bool(abs(observed_sd / predicted_sd - 1.0) < 0.30)
    return {
        "prediction": "Gaussian JL fixed-pair distance ratio has mean near one and SD near 1/sqrt(2r)",
        "epsilon": EPSILON,
        "constant": CONSTANT,
        "original_width": int(net.d),
        "projected_width": int(net.r),
        "raw_width_rule": float(net.raw_width),
        "trained_test_metric": float(metric),
        "observed_ratio_mean": observed_mean,
        "observed_ratio_sd": observed_sd,
        "predicted_ratio_mean": 1.0,
        "predicted_ratio_sd": predicted_sd,
        "ratio_sd_observed_over_predicted": float(observed_sd / predicted_sd),
        "p95_absolute_distortion": float(np.percentile(np.abs(ratios - 1.0), 95)),
        "confirmed": confirmed,
    }


def main():
    # Canonical baseline sweep on the four tuning seeds.
    tuned = sweep_baseline(
        lambda cfg: train_fn("baseline", cfg), GRID, seeds=(0, 1, 2, 3)
    )
    # Full baseline results for every shared configuration, useful for parity audit.
    baseline_sweep = []
    for cfg in GRID:
        baseline_sweep.append({"cfg": cfg, **evaluate(train_fn("baseline", cfg), SEEDS)})
    best_cfg = tuned["best_cfg"]
    baseline = {
        "best_cfg": best_cfg,
        "sweep": baseline_sweep,
        "harness_tuning": tuned,
        "full": evaluate(train_fn("baseline", best_cfg), SEEDS),
    }
    idea_sweep = []
    for cfg in GRID:
        idea_sweep.append({"cfg": cfg, **evaluate(train_fn("idea", cfg), SEEDS)})
    idea_best = min(idea_sweep, key=lambda row: row["mean"])
    idea = {k: idea_best[k] for k in ("per_seed", "mean", "std", "n")}
    report = make_report(
        "sequence",
        "transformer_tiny",
        baseline,
        idea,
        {
            "idea_sweep": idea_sweep,
            "mechanism_signature": model_signature(),
            "matched_structure": "sequence-level jointly processed windows/tokens",
        },
    )
    report["custom_track"] = None
    report["idea_hyperparameters"] = {
        "epsilon": EPSILON,
        "constant": CONSTANT,
        "width_rule": "ceil(C*min(d,n-1,log(2+epsilon^2*n)/epsilon^2))",
    }
    with open("bench_report.json", "w") as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()