"""Training with the equilibrium-residual regularizer (Stage 3 / Gate C).

Two approaches (master plan §2.6), selected by ``config['equilibrium']['approach']``:

  A (principled): model outputs 3 channels (sigma_xx, sigma_yy, sigma_xy). The data loss compares
    the derived von Mises stress to the scalar target; the physics loss penalizes the discrete
    divergence residual ||div(sigma)||^2 on interior nodes (operators validated analytically).
    Target is scaled by a constant S so outputs stay O(1); relative-L2 is scale-invariant, so the
    reported number equals the physical relative-L2.

  B (fallback): model keeps the 1-channel scalar output (identical accuracy path to LinearNO);
    the physics loss is a graph-Laplacian smoothness prior — a plausibility prior *motivated by*
    (not equal to) equilibrium. Honest framing required.

The per-sample discrete operators depend only on the mesh, so they are precomputed once (sparse)
and reused every epoch. Total loss: L = data_loss + lambda * physics_loss.
"""
from __future__ import annotations

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

import torch

from .data.dataset import build_splits
from .losses.equilibrium import (
    build_graph_laplacian,
    build_mls_gradient_operators,
    interior_mask,
    von_mises,
)
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_grad_ops(
    coords: torch.Tensor, k: int, tol: float, device
) -> List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
    """Per-sample (Gx_sparse, Gy_sparse, interior_mask) for Approach A. Built once (CPU build)."""
    ops = []
    for s in range(coords.shape[0]):
        c = coords[s]
        Gx, Gy = build_mls_gradient_operators(c, k=k)
        mask = interior_mask(c, tol).to(device)
        ops.append((Gx.to_sparse().to(device), Gy.to_sparse().to(device), mask))
    return ops


def _precompute_laplacians(coords: torch.Tensor, k: int, device) -> List[torch.Tensor]:
    """Per-sample sparse graph Laplacian for Approach B."""
    return [build_graph_laplacian(coords[s], k=k).to_sparse().to(device) for s in range(coords.shape[0])]


def _sparse_div_residual(stress3: torch.Tensor, Gx_s, Gy_s, mask) -> torch.Tensor:
    """L_eq = mean_interior ||div(sigma)||^2 for one sample. stress3: (N, 3)."""
    sxx = stress3[:, 0:1]
    syy = stress3[:, 1:2]
    sxy = stress3[:, 2:3]
    div_x = torch.sparse.mm(Gx_s, sxx) + torch.sparse.mm(Gy_s, sxy)   # (N,1)
    div_y = torch.sparse.mm(Gx_s, sxy) + torch.sparse.mm(Gy_s, syy)
    sq = (div_x.squeeze(-1) ** 2 + div_y.squeeze(-1) ** 2)            # (N,)
    return sq[mask].mean()


def run_training_eqreg(
    config: Dict[str, Any],
    seed: int,
    data_dir: str,
    device: Optional[str] = None,
    gpu_name: str = "CPU",
    results_path: Optional[str] = None,
    ckpt_path: Optional[str] = None,
    log_every: int = 50,
    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)

    data_cfg = config["data"]
    train_cfg = config["train"]
    model_cfg = config["model"]
    eq_cfg = config["equilibrium"]
    approach = eq_cfg.get("approach", "A").upper()
    lam = float(lambda_override) if lambda_override is not None else float(eq_cfg.get("lambda", 0.05))
    knn_k = int(eq_cfg.get("knn_k", 12))
    tol = float(eq_cfg.get("interior_tol", 0.03))

    if splits is None:
        ntrain = data_cfg.get("ntrain", 1000)
        ntest = data_cfg.get("ntest", 200)
        splits = build_splits(data_dir, ntrain=ntrain, ntest=ntest)
    ntest = splits.test_coords.shape[0]
    normalizer = splits.normalizer.to(device)

    train_coords = splits.train_coords.to(device)
    test_coords = splits.test_coords.to(device)
    train_sigma = (splits.train_sigma if splits.train_sigma.dim() == 3 else splits.train_sigma.unsqueeze(-1)).to(device)
    test_sigma = (splits.test_sigma if splits.test_sigma.dim() == 3 else splits.test_sigma.unsqueeze(-1)).to(device)
    n_train = train_coords.shape[0]

    # Scale for Approach A: keep von Mises target O(1). relative-L2 is scale-invariant, so the
    # reported metric equals the physical relative-L2 regardless of S.
    S = float(splits.train_sigma.mean()) if approach == "A" else 1.0

    print(f"[eqreg seed {seed}] approach={approach} lambda={lam} k={knn_k} S={S:.2f}", flush=True)
    t_build = time.time()
    if approach == "A":
        train_ops = _precompute_grad_ops(splits.train_coords, knn_k, tol, device)
        test_ops = _precompute_grad_ops(splits.test_coords, knn_k, tol, device)
    else:
        train_ops = _precompute_laplacians(splits.train_coords, knn_k, device)
        test_ops = _precompute_laplacians(splits.test_coords, knn_k, device)
    print(f"[eqreg 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)

    lr = float(train_cfg.get("lr", 1e-3))
    wd = float(train_cfg.get("weight_decay", 1e-5))
    betas = tuple(train_cfg.get("betas", (0.9, 0.999)))
    epochs = max_epochs or int(train_cfg.get("epochs", 500))
    max_grad_norm = train_cfg.get("max_grad_norm", None)
    eval_every = int(train_cfg.get("eval_every", 10))

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd, betas=betas)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    gen = torch.Generator().manual_seed(seed)

    def data_and_phys(out, sigma_phys, ops_i):
        """Return (data_loss, phys_residual) for one sample. out: (N, C).

        Approach A scale convention (IMPORTANT — do NOT 'decode' ``out``):
        ``out`` is the RAW 3-channel model output (no normalizer applied). It learns the physical
        stress tensor divided by the constant ``S`` (a pure scale, NOT the affine z-score). Because
        von Mises is homogeneous of degree 1 and relative-L2 is jointly scale-invariant,
            relative_l2(von_mises(out), sigma_phys/S) == relative_l2(S*von_mises(out), sigma_phys)
        i.e. the reported data loss EQUALS the physical relative-L2 of the prediction ``S*von_mises(out)``
        against the physical target (verified numerically). Applying the scalar normalizer's
        ``decode`` here would be WRONG: it would add the von-Mises mean (~187) to every tensor
        component. The divergence residual below is therefore in scaled units (= physical/S^2);
        it is converted to physical units only for reporting (see final_metrics).
        """
        if approach == "A":
            vm = von_mises(out)                      # (N,)  von Mises of the scaled tensor
            data_loss = relative_l2(vm.unsqueeze(0).unsqueeze(-1), (sigma_phys / S).unsqueeze(0), reduction="mean")
            Gx_s, Gy_s, mask = ops_i
            phys = _sparse_div_residual(out, Gx_s, Gy_s, mask)  # scaled units (physical/S^2)
        else:  # B: scalar output + Laplacian smoothness
            pred = normalizer.decode(out)            # (N,1) physical
            data_loss = relative_l2(pred.unsqueeze(0), sigma_phys.unsqueeze(0), reduction="mean")
            L_s = ops_i
            Lf = torch.sparse.mm(L_s, out[:, :1])    # operate on the (normalized) scalar field
            phys = (Lf ** 2).mean()
        return data_loss, phys

    @torch.no_grad()
    def evaluate() -> Tuple[float, float]:
        model.eval()
        tot_data, tot_phys = 0.0, 0.0
        for i in range(ntest):
            out = model(test_coords[i:i + 1], None)[0]   # (N, C)
            dloss, phys = data_and_phys(out, test_sigma[i], test_ops[i])
            tot_data += dloss.item()
            tot_phys += phys.item()
        return tot_data / ntest, tot_phys / ntest

    t0 = time.time()
    best_rel = float("inf")
    test_rel = float("nan")
    test_phys = float("nan")
    history = []
    for ep in range(epochs):
        model.train()
        perm = torch.randperm(n_train, generator=gen).tolist()
        run_data = 0.0
        for i in perm:
            optimizer.zero_grad()
            out = model(train_coords[i:i + 1], None)[0]   # (N, C)
            dloss, phys = data_and_phys(out, train_sigma[i], train_ops[i])
            loss = dloss + lam * phys
            loss.backward()
            if max_grad_norm is not None:
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            optimizer.step()
            run_data += dloss.item()
        scheduler.step()
        train_rel = run_data / n_train

        if (ep % eval_every == 0) or (ep >= epochs - 5):
            test_rel, test_phys = evaluate()
            best_rel = min(best_rel, test_rel)
        history.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"[eqreg 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)
    # Report the residual in physical units. Approach A computes div on the S-scaled tensor, so the
    # physical residual is test_phys * S^2; Approach B's Laplacian residual is already in the
    # (normalized) field's units (S == 1).
    test_residual_phys = test_phys * (S ** 2) if approach == "A" else test_phys
    final_metrics = {
        "test_rel_l2": round(test_rel, 6),
        "best_test_rel_l2": round(best_rel, 6),
        "test_residual": test_residual_phys,
        "test_residual_scaled": test_phys,
        "scale_S": S,
        "train_rel_l2": round(train_rel, 6),
        "n_params": n_params,
        "epochs": epochs,
        "approach": approach,
        "lambda": lam,
    }
    if ckpt_path is not None:
        os.makedirs(os.path.dirname(ckpt_path) or ".", exist_ok=True)
        torch.save(
            {"state_dict": model.state_dict(),
             "normalizer": {"mean": normalizer.mean.detach().cpu(), "std": normalizer.std.detach().cpu()},
             "scale_S": S, "config": config, "seed": seed, "metrics": final_metrics},
            ckpt_path,
        )
    if results_path is None:
        results_path = os.path.join("results", f"{config.get('name','eqreg')}_seed{seed}.json")
    write_run_log(results_path, config, seed, final_metrics, wall, gpu_name, wall * rate,
                  extra={"history_tail": history[-5:]})
    print(f"[eqreg 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 final_metrics