"""CPU inference helpers for the demo (loads a trained checkpoint; handles 1- and 3-channel heads).

Kept dependency-light so it also works inside a Hugging Face Space. The model architecture is
imported from the installed ``stress_operator`` package; the HF export bundles a self-contained
copy (see scripts/05_export_to_hf.sh).
"""
from __future__ import annotations

import os
import sys

import numpy as np
import torch

# Make `stress_operator` importable in BOTH layouts: the local repo (package under ../src) and a
# bundled Hugging Face Space (package is a sibling of this file).
_HERE = os.path.dirname(os.path.abspath(__file__))
for _p in (_HERE, os.path.join(_HERE, "..", "src")):
    if _p not in sys.path:
        sys.path.insert(0, _p)

from stress_operator.models.transolver import build_model  # noqa: E402


def load_checkpoint(ckpt_path: str, device: str = "cpu"):
    """Load a trained model. Supports two formats:

    - ``*.safetensors`` (HF deployment): weights from safetensors + ``config.json`` (sibling)
      holding the model config, normalizer mean/std, and scale_S.
    - ``*.pt`` (local training checkpoint): a dict with state_dict / normalizer / config.
    """
    if ckpt_path.endswith(".safetensors"):
        import json

        from safetensors.torch import load_file

        cfg_path = os.path.join(os.path.dirname(ckpt_path) or ".", "config.json")
        with open(cfg_path) as f:
            cfg = json.load(f)
        model_cfg = cfg["model"]
        state_dict = load_file(ckpt_path, device=device)
        mean = torch.tensor(float(cfg["normalizer"]["mean"]), device=device).reshape(1, 1, 1)
        std = torch.tensor(float(cfg["normalizer"]["std"]), device=device).reshape(1, 1, 1)
        scale_S = cfg.get("scale_S", None)
        metrics = cfg.get("metrics", {})
    else:
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
        model_cfg = ckpt["config"]["model"]
        state_dict = ckpt["state_dict"]
        mean = ckpt["normalizer"]["mean"].to(device)
        std = ckpt["normalizer"]["std"].to(device)
        scale_S = ckpt.get("scale_S", None)
        metrics = ckpt.get("metrics", {})

    model = build_model(model_cfg).to(device)
    model.load_state_dict(state_dict)
    model.eval()
    info = {
        "out_dim": model_cfg.get("out_dim", 1),
        "scale_S": scale_S,
        "attention": model_cfg.get("attention", "physics"),
        "metrics": metrics,
    }
    return model, (mean, std), info


@torch.no_grad()
def predict_stress(model, coords: np.ndarray, norm, info, device: str = "cpu") -> np.ndarray:
    """coords: (N, 2) -> per-node von Mises stress (N,) in physical units."""
    mean, std = norm
    x = torch.as_tensor(coords, dtype=torch.float32, device=device).unsqueeze(0)  # (1,N,2)
    out = model(x, None)[0]  # (N, out_dim)
    if info["out_dim"] == 3:
        from stress_operator.losses.equilibrium import von_mises

        S = info["scale_S"] or 1.0
        stress = von_mises(out) * S          # 3-channel tensor -> von Mises, undo target scale
    else:
        # de-normalize scalar prediction; reshape so the broadcast of a (1,1[,1]) normalizer
        # against out[:,0] (N,) cannot leak a leading axis (audit bug 1). Always returns (N,).
        stress = (out[:, 0] * std.reshape(()) + mean.reshape(()))
    return stress.reshape(-1).detach().cpu().numpy()