import os, sys, json
import numpy as np
import torch
import torch.nn as nn

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


def seed_all(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 Dense3D(nn.Module):
    def __init__(self):
        super().__init__()
        self.op = nn.Linear(64, 64, bias=False)
        self.head = nn.Sequential(nn.ReLU(), nn.Linear(64, 1))
        nn.init.xavier_uniform_(self.op.weight)

    def forward(self, x):
        return self.head(self.op(x.reshape(x.shape[0], 64)))


class TT3D(nn.Module):
    """Learned TT-matrix over the three 4-way spatial modes."""
    def __init__(self, rank=2):
        super().__init__()
        self.rank = rank
        self.cores = nn.ParameterList([
            nn.Parameter(torch.randn(1, 4, 4, rank) * 0.16),
            nn.Parameter(torch.randn(rank, 4, 4, rank) * 0.16),
            nn.Parameter(torch.randn(rank, 4, 4, 1) * 0.16),
        ])
        self.head = nn.Sequential(nn.ReLU(), nn.Linear(64, 1))

    def forward(self, x):
        g1, g2, g3 = self.cores
        # x[b,i,j,k], g1[a,l,i,c], g2[c,n,j,d], g3[d,o,k,f].
        y = torch.einsum("bijk,alic,cnjd,dokf->blno", x, g1, g2, g3)
        return self.head(y.reshape(x.shape[0], 64))


def train_one(kind, seed, lr, epochs, rank=2):
    seed_all(seed)
    d0 = get_dataset(seed, 400, 100)
    d = {**d0, **{k: torch.as_tensor(d0[k], dtype=torch.float32) for k in ("xtr", "ytr", "xte", "yte")}}
    model = Dense3D() if kind == "dense" else TT3D(rank)
    net, metric, _ = train_model(model, d, epochs=epochs, lr=lr, batch=128,
                                 weight_decay=0.0, log=lambda *_: None)
    if net is None:
        return float("nan"), None
    return float(metric), net


def main():
    # The union of step sizes is identical for baseline and idea.
    grid = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
    epochs = 18

    def baseline_factory(cfg):
        return lambda s: train_one("dense", int(s), cfg["lr"], epochs)[0]

    base = sweep_baseline(baseline_factory, grid)
    idea_trials = []
    for cfg in grid:
        r = evaluate(lambda s, c=cfg: train_one("tt", int(s), c["lr"], epochs, rank=2)[0])
        idea_trials.append({"cfg": cfg, **r})
    best_idea = min(idea_trials, key=lambda r: r["mean"])

    # Signature is measured on trained models and test inputs.
    best_lr = base["best_cfg"]["lr"]
    _, bnet = train_one("dense", 0, best_lr, epochs)
    _, inet = train_one("tt", 0, best_idea["cfg"]["lr"], epochs, rank=2)
    d = get_dataset(0, 400, 100)
    xb = d["xte"]
    with torch.no_grad():
        dense_out = bnet.op(xb.reshape(-1, 64)).numpy()
        g1, g2, g3 = [g.detach() for g in inet.cores]
        tt_out = torch.einsum("bijk,alic,cnjd,dokf->blno", xb, g1, g2, g3)
        tt_out = tt_out.reshape(-1, 64).numpy()
    sig = {
        "quantity": "trained operator output norm ratio on held-out PDE fields",
        "predicted": "TT has fewer parameters than the dense spatial map",
        "observed": {
            "dense_params": count_params(bnet),
            "tt_params": count_params(inet),
            "operator_output_norm_ratio": float(np.linalg.norm(tt_out) / (np.linalg.norm(dense_out) + 1e-12)),
        },
        "confirmed": bool(count_params(inet) < count_params(bnet)),
    }
    idea = {"best_cfg": best_idea["cfg"], "trials": idea_trials,
            "mean": best_idea["mean"], "std": best_idea["std"],
            "per_seed": best_idea["per_seed"], "n": best_idea["n"]}
    rep = make_report("separable_diffusion_field", "custom_dense_vs_tt3d", base, idea, sig)
    rep["custom_track"] = {"name": META["name"], "file": "custom_track.py", "domain": META["domain"]}
    rep["parameter_counts"] = {"baseline": count_params(bnet), "idea": count_params(inet)}
    with open("bench_report.json", "w") as f:
        json.dump(rep, f, indent=2)
    print(json.dumps(rep, indent=2))


if __name__ == "__main__":
    main()