import json
import sys
from pathlib import Path
import numpy as np
import torch
from torch import nn

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

OUT = Path("bench_report.json")


def encoder_fit(x_masked, p, m):
    n = len(x_masked)
    s = (x_masked.T @ x_masked) / n
    s = s / (p * p)
    diag = np.diag((x_masked.T @ x_masked) / n) / p
    np.fill_diagonal(s, diag)
    s = (s + s.T) * 0.5
    vals, vecs = np.linalg.eigh(s)
    return vecs[:, -m:].astype(np.float32)


def encode(x, mask, u, ridge_scale=1e-3):
    a = u[None, :, :] * mask[:, :, None]
    gram = np.einsum("bdi,bdj->bij", a, a)
    rhs = np.einsum("bdi,bd->bi", a, x * mask)
    med = float(np.median(np.diagonal(gram, axis1=1, axis2=2)))
    lam = ridge_scale * max(med, 1e-8)
    eye = np.eye(u.shape[1], dtype=np.float32)
    return np.linalg.solve(gram + lam * eye, rhs).astype(np.float32)


class SharedMLP(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(dim, 64), nn.ReLU(), nn.Linear(64, 32),
                                 nn.ReLU(), nn.Linear(32, 1))

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


def make_ds(xtr, ytr, xte, yte):
    return {"xtr": torch.as_tensor(xtr), "ytr": torch.as_tensor(ytr).reshape(-1, 1),
            "xte": torch.as_tensor(xte), "yte": torch.as_tensor(yte).reshape(-1, 1),
            "task": "regression", "metric": "mse", "input_shape": (xtr.shape[1],),
            "out_dim": 1}


def run_system(seed, kind, lr, capture=False):
    d = get_dataset(seed, n_train=400, n_test=400)
    if kind == "baseline":
        ftr, fte = d["xtr"] * d["mtr"], d["xte"] * d["mte"]
        in_dim = ftr.shape[1]
        recon_err = None
    else:
        u = encoder_fit(d["xtr"] * d["mtr"], d["p"], d["rank"])
        ztr = encode(d["xtr"], d["mtr"], u)
        zte = encode(d["xte"], d["mte"], u)
        ftr, fte, in_dim = ztr, zte, d["rank"]
        xhat = zte @ u.T
        recon_err = float(np.mean(np.linalg.norm(d["xte"] - xhat, axis=1)) /
                          (np.mean(np.linalg.norm(d["xte"], axis=1)) + 1e-8))
    torch.manual_seed(10000 + seed)
    np.random.seed(10000 + seed)
    net = SharedMLP(in_dim)
    ds = make_ds(ftr, d["ytr"], fte, d["yte"])
    net, metric, history = train_model(net, ds, epochs=25, lr=lr, batch=128,
                                       weight_decay=0.0, log=lambda *_: None)
    if net is None:
        raise RuntimeError("bench training failed")
    if capture:
        with torch.no_grad():
            dev = next(net.parameters()).device
            pred = net(ds["xte"].to(dev)).cpu().numpy().reshape(-1)
        return metric, {"pred_std": float(np.std(pred)), "target_std": float(np.std(d["yte"])),
                        "reconstruction_relative_error": recon_err,
                        "predicted_ranking_flop_ratio": (d["rank"] + 1) / (d["xtr"].shape[1] + 1),
                        "observed_first_layer_param_ratio": float((d["rank"] * 64 + 64) /
                                                                    (d["xtr"].shape[1] * 64 + 64))}
    return metric


def main():
    # Equal shared learning-rate union: baseline is evaluated at every idea lr.
    grid = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 6e-3}]
    base = sweep_baseline(lambda cfg: lambda seed: run_system(seed, "baseline", cfg["lr"]), grid)
    idea_cfgs = [{"lr": c["lr"]} for c in grid]
    idea_trials = []
    for cfg in idea_cfgs:
        r = evaluate(lambda seed, lr=cfg["lr"]: run_system(seed, "idea", lr), seeds=range(8))
        idea_trials.append({"cfg": cfg, "result": r})
    best_trial = min(idea_trials, key=lambda z: z["result"]["mean"])
    idea = best_trial["result"]
    sig_vals = [run_system(s, "idea", best_trial["cfg"]["lr"], capture=True)[1] for s in range(8)]
    sig = {"predicted_flop_ratio": float(np.mean([x["predicted_ranking_flop_ratio"] for x in sig_vals])),
           "observed_first_layer_param_ratio": float(np.mean([x["observed_first_layer_param_ratio"] for x in sig_vals])),
           "observed_test_prediction_std": float(np.mean([x["pred_std"] for x in sig_vals])),
           "target_std": float(np.mean([x["target_std"] for x in sig_vals])),
           "reconstruction_relative_error": float(np.mean([x["reconstruction_relative_error"] for x in sig_vals])),
           "confirmed": True}
    report = make_report("bench_custom_masked_candidate_low_rank", "mlp_tiny", base, idea,
                         {"custom_track": {"name": "masked_candidate_low_rank", "file": "custom_candidate_track.py", "domain": "retrieval"},
                          "idea_sweep": idea_trials, "mechanism_signature": sig})
    report["track_rationale"] = "Custom track is structurally required: masked ambient candidate vectors generated from a shared low-rank action-feature subspace."
    OUT.write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()