from pathlib import Path
import re
from typing import Any, List, Mapping

import torch


class ShaperNet(torch.nn.Module):
    """Neural network used to infer shaper order probability and modal parameters.

    Args:
        in_features: Number of input features per axis.
        hidden: Hidden layer widths.
        order_features: Number of order-head outputs.
        mode_features: Number of mode-head outputs.
    Side Effects:
        Initializes trainable torch layers.
    Raises:
        ValueError: If any configured feature width is non-positive.
    Preconditions:
        `in_features > 0`.
    """

    def __init__(
        self,
        in_features: int,
        hidden: List[int],
        order_features: int = 1,
        mode_features: int = 4,
    ) -> None:
        """Build the per-axis inference network.

        Args:
            in_features: Number of input features.
            hidden: Hidden layer widths.
            order_features: Number of order-head outputs.
            mode_features: Number of mode-head outputs.
        Returns:
            `None`.
        Side Effects:
            Allocates torch module parameters.
        Raises:
            ValueError: If any configured feature width is non-positive.
        Preconditions:
            `in_features > 0`.
        """
        if in_features <= 0:
            raise ValueError("in_features must be positive")
        if any(h <= 0 for h in hidden):
            raise ValueError("hidden widths must be positive")
        if order_features <= 0:
            raise ValueError("order_features must be positive")
        if mode_features <= 0:
            raise ValueError("mode_features must be positive")
        super().__init__()
        layers: list[torch.nn.Module] = []
        last = in_features
        for h in hidden:
            layers.append(torch.nn.Linear(last, h))
            layers.append(torch.nn.ReLU())
            last = h
        self.body = torch.nn.Sequential(*layers)
        self.order_head = torch.nn.Linear(last, order_features)
        self.mode_head = torch.nn.Linear(last, mode_features)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Run one forward pass and clamp probability output.

        Args:
            x: Feature tensor shaped `[batch, in_features]`.
        Returns:
            `tuple[torch.Tensor, torch.Tensor]` with `(order_prob, mode_params)`.
        Side Effects:
            None.
        Raises:
            RuntimeError: If torch linear layers fail due to shape mismatch.
        Preconditions:
            `x.shape[-1] == in_features`.
        """
        x = self.body(x)
        order_logits = self.order_head(x)
        order = torch.sigmoid(order_logits)
        order = torch.where(order.isnan(), torch.full_like(order, 0.5), order)
        order = torch.clamp(order, 1e-6, 1.0 - 1e-6)
        modes = self.mode_head(x)
        return order, modes


class MapFitterPy:
    """Load one model per axis and run batch inference.

    Args:
        axes: Number of shaped axes.
        in_features: Number of features per axis.
        hidden: Hidden layer widths.
        directory: Directory containing `axis*_model.pt`.
    Side Effects:
        Loads model files from disk.
    Raises:
        RuntimeError: If model file type is unsupported.
    Preconditions:
        `directory` exists and contains one file per axis.
    """

    def __init__(
        self, axes: int, in_features: int, hidden: List[int], directory: str
    ) -> None:
        """Load all per-axis model checkpoints.

        Args:
            axes: Number of axes.
            in_features: Number of model input features.
            hidden: Hidden layer widths.
            directory: Model directory path.
        Returns:
            `None`.
        Side Effects:
            Loads model checkpoints into memory.
        Raises:
            RuntimeError: If a loaded model object is unsupported.
        Preconditions:
            All `axis{idx}_model.pt` files exist for `idx in [0, axes)`.
        """
        self.axes = axes
        self.in_features = in_features
        self.hidden = hidden
        self.models: list[torch.nn.Module] = []
        directory_path = Path(directory)
        for axis in range(axes):
            fn = directory_path / f"axis{axis}_model.pt"
            model = self._load_axis_model(fn)
            model.eval()
            self.models.append(model)

    def _load_axis_model(self, fn: Path) -> torch.nn.Module:
        """Load one axis model from disk, supporting scripted and state-dict formats."""
        try:
            return torch.jit.load(str(fn), map_location="cpu")
        except Exception:
            pass

        loaded = torch.load(fn, map_location="cpu", weights_only=False)
        if isinstance(
            loaded,
            (torch.jit.ScriptModule, torch.jit.RecursiveScriptModule),
        ):
            return loaded
        if isinstance(loaded, torch.nn.Module):
            return loaded

        state_dict = self._extract_state_dict(loaded, fn)
        remapped = self._remap_state_dict_keys(state_dict)
        topology = self._infer_shapernet_topology(remapped)
        if topology is None:
            model = ShaperNet(self.in_features, self.hidden)
        else:
            in_features, hidden, order_features, mode_features = topology
            model = ShaperNet(
                in_features=in_features,
                hidden=hidden,
                order_features=order_features,
                mode_features=mode_features,
            )

        model.load_state_dict(remapped, strict=False)
        return model

    @staticmethod
    def _extract_state_dict(loaded: Any, fn: Path) -> Mapping[str, Any]:
        """Return a flat state-dict from common checkpoint formats."""
        if isinstance(loaded, Mapping):
            if "state_dict" in loaded and isinstance(loaded["state_dict"], Mapping):
                return loaded["state_dict"]
            return loaded
        raise RuntimeError(f"Unsupported model type in {fn}: {type(loaded)}")

    @staticmethod
    def _remap_state_dict_keys(state_dict: Mapping[str, Any]) -> dict[str, Any]:
        """Normalize key names used by C++ and Python shaper checkpoints."""
        key_map: dict[str, str] = {
            "order.weight": "order_head.weight",
            "order.bias": "order_head.bias",
            "modes.weight": "mode_head.weight",
            "modes.bias": "mode_head.bias",
        }
        remapped: dict[str, Any] = {}
        for key, value in state_dict.items():
            normalized = str(key)
            if normalized.startswith("module."):
                normalized = normalized[len("module.") :]
            remapped[key_map.get(normalized, normalized)] = value
        return remapped

    @staticmethod
    def _infer_shapernet_topology(
        state_dict: Mapping[str, Any],
    ) -> tuple[int, list[int], int, int] | None:
        """Infer `(in_features, hidden, order_features, mode_features)` from weights."""
        body_weights: list[tuple[int, torch.Tensor]] = []
        for key, value in state_dict.items():
            match = re.fullmatch(r"body\.(\d+)\.weight", key)
            if match and isinstance(value, torch.Tensor) and value.dim() == 2:
                body_weights.append((int(match.group(1)), value))
        body_weights.sort(key=lambda x: x[0])

        order_w = state_dict.get("order_head.weight")
        mode_w = state_dict.get("mode_head.weight")
        if not isinstance(order_w, torch.Tensor) or order_w.dim() != 2:
            return None
        if not isinstance(mode_w, torch.Tensor) or mode_w.dim() != 2:
            return None

        if body_weights:
            hidden = [int(w.shape[0]) for _, w in body_weights]
            in_features = int(body_weights[0][1].shape[1])
            body_out = hidden[-1]
        else:
            hidden = []
            in_features = int(order_w.shape[1])
            body_out = in_features

        if int(order_w.shape[1]) != body_out or int(mode_w.shape[1]) != body_out:
            return None
        order_features = int(order_w.shape[0])
        mode_features = int(mode_w.shape[0])
        return in_features, hidden, order_features, mode_features

    def infer_batch(self, features: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Infer shaper outputs for all configured axes.

        Args:
            features: Feature tensor shaped `[axes, in_features]`.
        Returns:
            `tuple[torch.Tensor, torch.Tensor]` as `(orders, modes)`.
        Side Effects:
            None.
        Raises:
            ValueError: If `features` shape does not match model configuration.
        Preconditions:
            Number of rows in `features` is at least `axes`.
        """
        if features.dim() != 2 or features.size(0) < self.axes:
            raise ValueError("features must be [axes, in_features]")
        orders: list[torch.Tensor] = []
        modes: list[torch.Tensor] = []
        for axis, model in enumerate(self.models):
            x = features[axis : axis + 1]
            with torch.no_grad():
                order, mode = model(x)
            orders.append(order)
            modes.append(mode)
        return torch.cat(orders, dim=0), torch.cat(modes, dim=0)


class ModelLoader:
    """Factory for loading map-fitter model bundles."""

    @staticmethod
    def load(
        directory: str, axes: int, input_features: int, hidden: List[int]
    ) -> MapFitterPy:
        """Instantiate and return a `MapFitterPy`.

        Args:
            directory: Directory containing axis model checkpoints.
            axes: Number of axes.
            input_features: Input feature count.
            hidden: Hidden layer widths.
        Returns:
            `MapFitterPy` configured with loaded models.
        Side Effects:
            Loads model files from disk.
        Raises:
            RuntimeError: If checkpoint loading fails.
        Preconditions:
            `directory` contains expected model files.
        """
        return MapFitterPy(axes, input_features, hidden, directory)
