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

TRACK, MODEL = "tabular", "mlp_tiny"
EPOCHS, BATCH = 8, 128
LRS = [0.001, 0.003, 0.006]  # identical union on baseline and idea sides
LAMBDAS = [0.02, 0.05, 0.10]
torch.set_num_threads(1)


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


def mgs_numpy(Z, Y, d, eps=1e-6):
    qs, ys, res = [], [], []
    for z, y in zip(Z, Y):
        nz = np.linalg.norm(z)
        if nz < eps: continue
        u = z / nz; v = u.copy()
        for q in qs: v -= q * np.dot(q, v)
        nv = np.linalg.norm(v); res.append(nv)
        if nv > eps:
            qs.append(v / nv); ys.append(y / nz)
            if len(qs) == d: break
    if not qs: return None, None, res
    return np.asarray(qs, dtype=np.float32), np.asarray(ys, dtype=np.float32), res


def run(seed, lr, lam=0.0, signature=False):
    seed_all(seed)
    ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=200)
    net = make_model(MODEL, ds["input_shape"], ds["out_dim"])
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"
    try:
        net.to(device)
        x, y = ds["xtr"].to(device), ds["ytr"].to(device)
        xt, yt = ds["xte"].to(device), ds["yte"].to(device)
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        lossf = nn.MSELoss()
        head = [m for m in net.modules() if isinstance(m, nn.Linear)][-1]
        memz, memy, Q = [], [], None
        ranks, orth, eigs = [], [], []
        for ep in range(EPOCHS):
            net.train(); perm = torch.randperm(len(x), device=device)
            for st in range(0, len(x), BATCH):
                ix = perm[st:st+BATCH]; loss = lossf(net(x[ix]), y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
                if lam:
                    with torch.no_grad():
                        h = x[ix]
                        for layer in list(net.children())[:-1]: h = layer(h)
                        memz.append(h.detach().cpu().numpy()); memy.append(y[ix].detach().cpu().numpy())
            if lam:
                # Refresh once per epoch, using at most 256 recent examples.
                Z = np.concatenate(memz)[-256:]; Y = np.concatenate(memy)[-256:]
                q, qy, _ = mgs_numpy(Z, Y, head.in_features)
                if q is not None:
                    Q = torch.from_numpy(q).to(device); Qy = torch.from_numpy(qy).to(device)
                    with torch.no_grad():
                        old = head.weight.detach().clone()
                        residual = Qy - Q @ head.weight.T
                        head.weight.sub_(lr * lam * (residual.T @ Q))
                        G = Q.T @ Q
                        ranks.append(int(Q.shape[1])); orth.append(float(torch.max(torch.abs(G-torch.eye(G.shape[0],device=device))).cpu()))
                        eigs.append(float(torch.linalg.eigvalsh(G).max().cpu()))
        with torch.no_grad(): metric = float(torch.mean((net(xt)-yt)**2).cpu())
        sig = None
        if signature and eigs:
            sig = {"feature_dim": head.in_features, "observed_rank_last": ranks[-1],
                   "predicted_identity_max_eigen_mean": float(np.mean(eigs)),
                   "observed_QtQ_error_max": max(orth),
                   "prediction": "MGS memory predicts identity Gramian after d independent directions",
                   "confirmed": bool(ranks[-1] == head.in_features and max(orth) < 1e-4 and abs(np.mean(eigs)-1)<1e-4)}
        return metric, sig
    except RuntimeError:
        if use_cuda:
            torch.cuda.empty_cache()
            # Explicit CPU retry, avoiding recursive CUDA detection.
            torch.cuda.is_available = lambda: False
            try: return run(seed, lr, lam, signature)
            finally: torch.cuda.is_available = lambda: True
        raise


def fn(cfg):
    return lambda s: run(s, cfg["lr"], cfg.get("lambda", 0.0))[0]


def main():
    base = sweep_baseline(fn, [{"lr":lr, "weight_decay":0.0} for lr in LRS])
    ideas = []
    for lam in LAMBDAS:
        cfg = {"lr": base["best_cfg"]["lr"], "lambda": lam}
        ideas.append({"cfg":cfg, "result":evaluate(fn(cfg))})
    best = min(ideas, key=lambda z:z["result"]["mean"])
    _, sig = run(0, best["cfg"]["lr"], best["cfg"]["lambda"], True)
    report = make_report(TRACK, MODEL, base, best["result"], {
        "prediction": sig, "idea_configs": ideas,
        "confirmed": bool(sig and sig["confirmed"])
    })
    report["track_justification"] = "Tabular is the prescribed optimizer track; Friedman regression with mlp_tiny has a trainable final linear head and fixed feature dimension, matching the memory correction."
    Path("bench_report.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()