"""Tensile2d training with the equilibrium regularizer — the SUPERVISED-tensor transfer experiment.

Unlike Geo-FNO Elasticity (scalar von Mises target -> latent tensor), Tensile2d directly supervises the
full Cauchy tensor (sig11, sig22, sig12) per node. So here:
  - the data loss is relative-L2 on the 3 SUPERVISED stress components (not a derived von Mises);
  - the model conditions on the 6 input scalars (P, p1..p5) via fun_dim=6 (broadcast per node);
  - the physics loss is the same meshfree MLS divergence ||div(sigma)||^2 on interior nodes.
Meshes are variable-size (~6k-12k nodes), so we keep per-sample lists (batch_size=1).

This run lets us test the key question raised by the Elasticity cross-operator finding: does supervising
the true tensor REDUCE the operator-specific 'gaming' (MLS-vs-FE residual gap)? (Checked post-hoc.)
"""
from __future__ import annotations

import os
import time
from typing import Any, Dict, Optional

import torch

from .data.tensile2d import build_tensile_splits
from .losses.equilibrium import build_mls_gradient_operators_sparse, interior_mask_knn
from .losses.relative_l2 import relative_l2
from .models.transolver import build_model, count_parameters
from .seeds import set_seed
from .utils.logging import MODAL_RATES_PER_SEC, write_run_log


def _precompute_ops(coords_list, k, device):
    ops = []
    for c in coords_list:
        Gx, Gy = build_mls_gradient_operators_sparse(c, k=k)   # sparse-direct (scales to ~10k nodes)
        ops.append((Gx.to(device), Gy.to(device), interior_mask_knn(c, k=k).to(device)))
    return ops


def _div_residual(stress3, Gx_s, Gy_s, mask):
    sxx, syy, sxy = stress3[:, 0:1], stress3[:, 1:2], stress3[:, 2:3]
    dx = torch.sparse.mm(Gx_s, sxx) + torch.sparse.mm(Gy_s, sxy)
    dy = torch.sparse.mm(Gx_s, sxy) + torch.sparse.mm(Gy_s, syy)
    return (dx.squeeze(-1) ** 2 + dy.squeeze(-1) ** 2)[mask].mean()


def run_training_tensile(
    config: Dict[str, Any], seed: int, data_dir: str = "data/tensile2d/samples",
    device: Optional[str] = None, gpu_name: str = "CPU",
    results_path: Optional[str] = None, ckpt_path: Optional[str] = None,
    log_every: int = 25, max_epochs: Optional[int] = None,
    splits=None, lambda_override: Optional[float] = None,
) -> Dict[str, Any]:
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    set_seed(seed)
    tr_cfg, model_cfg, eq_cfg = config["train"], config["model"], config["equilibrium"]
    lam = float(lambda_override) if lambda_override is not None else float(eq_cfg.get("lambda", 0.01))
    k = int(eq_cfg.get("knn_k", 12))  # interior_mask_knn is geometry-agnostic; no interior_tol needed

    if splits is None:
        d = config["data"]
        splits = build_tensile_splits(data_dir, ntrain=d.get("ntrain", 400), ntest=d.get("ntest", 100), seed=0)
    S = splits.scale_S
    trc = [c.to(device) for c in splits.train_coords]
    trs = [s.to(device).float() for s in splits.train_sigma]
    trsc = [sc.to(device).float() for sc in splits.train_scalars]
    tec = [c.to(device) for c in splits.test_coords]
    tes = [s.to(device).float() for s in splits.test_sigma]
    tesc = [sc.to(device).float() for sc in splits.test_scalars]
    n_train, n_test = len(trc), len(tec)

    print(f"[tensile seed {seed}] lambda={lam} k={k} S={S:.2f} ntrain={n_train} ntest={n_test}", flush=True)
    t_build = time.time()
    train_ops = _precompute_ops(splits.train_coords, k, device)
    test_ops = _precompute_ops(splits.test_coords, k, device)
    print(f"[tensile seed {seed}] precomputed operators in {time.time()-t_build:.0f}s", flush=True)

    model = build_model(model_cfg).to(device)
    n_params = count_parameters(model)
    epochs = max_epochs or int(tr_cfg.get("epochs", 500))
    opt = torch.optim.AdamW(model.parameters(), lr=float(tr_cfg.get("lr", 1e-3)),
                            weight_decay=float(tr_cfg.get("weight_decay", 1e-5)),
                            betas=tuple(tr_cfg.get("betas", (0.9, 0.999))))
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
    mgn = tr_cfg.get("max_grad_norm", None)
    eval_every = int(tr_cfg.get("eval_every", 10))
    gen = torch.Generator().manual_seed(seed)

    def fx_of(coords, scalars):  # broadcast the 6 input scalars to every node
        return scalars.view(1, 1, -1).expand(1, coords.shape[0], -1)

    def losses(coords, sigma, scalars, ops_i):
        out = model(coords.unsqueeze(0), fx_of(coords, scalars))[0]   # (N,3) scaled stress
        data = relative_l2(out.unsqueeze(0), sigma.unsqueeze(0), reduction="mean")  # SUPERVISED tensor
        phys = _div_residual(out, *ops_i)
        return data, phys

    @torch.no_grad()
    def evaluate():
        model.eval()
        d, p = 0.0, 0.0
        for i in range(n_test):
            dl, ph = losses(tec[i], tes[i], tesc[i], test_ops[i])
            d += dl.item(); p += ph.item()
        return d / n_test, p / n_test

    t0 = time.time()
    best, test_rel, test_phys, hist = float("inf"), float("nan"), float("nan"), []
    for ep in range(epochs):
        model.train()
        run = 0.0
        for i in torch.randperm(n_train, generator=gen).tolist():
            opt.zero_grad()
            dl, ph = losses(trc[i], trs[i], trsc[i], train_ops[i])
            (dl + lam * ph).backward()
            if mgn is not None:
                torch.nn.utils.clip_grad_norm_(model.parameters(), mgn)
            opt.step(); run += dl.item()
        sched.step()
        train_rel = run / n_train
        if (ep % eval_every == 0) or (ep >= epochs - 5):
            test_rel, test_phys = evaluate(); best = min(best, test_rel)
        hist.append({"epoch": ep, "train_rel": train_rel, "test_rel": test_rel, "test_phys": test_phys})
        if ep % log_every == 0 or ep == epochs - 1:
            print(f"[tensile seed {seed}] epoch {ep:4d} train_rel={train_rel:.5f} "
                  f"test_rel={test_rel:.5f} test_resid={test_phys:.4e}", flush=True)

    wall = time.time() - t0
    rate = MODAL_RATES_PER_SEC.get(gpu_name, 0.0)
    metrics = {
        "test_rel_l2": round(test_rel, 6), "best_test_rel_l2": round(best, 6),
        "test_residual": test_phys * (S ** 2), "test_residual_scaled": test_phys,
        "scale_S": S, "train_rel_l2": round(train_rel, 6),
        "n_params": n_params, "epochs": epochs, "lambda": lam, "dataset": "tensile2d",
    }
    if ckpt_path is not None:
        os.makedirs(os.path.dirname(ckpt_path) or ".", exist_ok=True)
        torch.save({"state_dict": model.state_dict(), "scale_S": S, "config": config,
                    "scalar_mean": splits.scalar_mean, "scalar_std": splits.scalar_std,
                    "seed": seed, "metrics": metrics}, ckpt_path)
    if results_path is None:
        results_path = os.path.join("results", f"{config.get('name','tensile')}_seed{seed}.json")
    write_run_log(results_path, config, seed, metrics, wall, gpu_name, wall * rate,
                  extra={"history_tail": hist[-5:]})
    print(f"[tensile seed {seed}] DONE test_rel={test_rel:.6f} resid={test_phys:.4e} "
          f"wall={wall:.0f}s est_cost=${wall*rate:.4f}", flush=True)
    return metrics