from __future__ import annotations

import copy
from typing import Any, Callable, Dict, Mapping, Optional, Tuple, Union, Literal
import warnings

import numpy as np
import torch
import torch.nn as nn

try:  # Optional scikit-learn import for API compatibility
    from sklearn.base import BaseEstimator, RegressorMixin  # type: ignore
    from sklearn.metrics import r2_score as _sk_r2_score  # type: ignore
except Exception:  # Fallbacks if sklearn isn't installed at runtime
    class BaseEstimator:  # minimal stub
        def get_params(self, deep: bool = True):
            # Return non-private, non-callable attributes
            params = {}
            for k, v in self.__dict__.items():
                if k.endswith("_"):
                    continue
                if not k.startswith("_") and not callable(v):
                    params[k] = v
            return params

        def set_params(self, **params):
            for k, v in params.items():
                setattr(self, k, v)
            return self

    class RegressorMixin:
        pass

    def _sk_r2_score(y_true, y_pred):
            y_true = np.asarray(y_true)
            y_pred = np.asarray(y_pred)
            u = ((y_true - y_pred) ** 2).sum()
            v = ((y_true - y_true.mean()) ** 2).sum()
            return 1.0 - (u / v if v != 0 else np.nan)

from .nn import PSANNNet, WithPreprocessor, ResidualPSANNNet
from .models import WaveResNet
from .conv import PSANNConv1dNet, PSANNConv2dNet, PSANNConv3dNet, ResidualPSANNConv2dNet
from .utils import choose_device, seed_all
from .types import ActivationConfig, LossLike, NoiseSpec, ScalerSpec
from .preproc import PreprocessorLike, build_preprocessor
from .state import StateConfig, ensure_state_config

from ._aliases import resolve_int_alias

ValidationDataLike = Union[
    Tuple[np.ndarray, np.ndarray],
    Tuple[np.ndarray, np.ndarray, np.ndarray],
]
from .estimators._fit_utils import (
    FitVariantHooks,
    HISSOTrainingPlan,
    ModelBuildRequest,
    NormalisedFitArgs,
    PreparedInputState,
    build_model_from_hooks,
    build_hisso_training_plan,
    maybe_run_hisso,
    normalise_fit_args,
    prepare_inputs_and_scaler,
    run_supervised_training,
    _build_optimizer as _build_optimizer_helper,
)



class PSANNRegressor(BaseEstimator, RegressorMixin):
    """Sklearn-style regressor wrapper around a PSANN network (PyTorch).

    Parameters mirror the README's proposed API.
    """

    def __init__(
        self,
        *,
        hidden_layers: int = 2,
        hidden_width: Optional[int] = None,
        hidden_units: Optional[int] = None,
        epochs: int = 200,
        batch_size: int = 128,
        lr: float = 1e-3,
        optimizer: str = "adam",
        weight_decay: float = 0.0,
        activation: Optional[ActivationConfig] = None,
        device: str | torch.device = "auto",
        random_state: Optional[int] = None,
        early_stopping: bool = False,
        patience: int = 20,
        num_workers: int = 0,
        loss: LossLike = "mse",
        loss_params: Optional[Dict[str, Any]] = None,
        loss_reduction: str = "mean",
        w0: float = 30.0,
        preserve_shape: bool = False,
        data_format: str = "channels_first",
        conv_kernel_size: int = 1,
        conv_channels: Optional[int] = None,
        per_element: bool = False,
        activation_type: str = "psann",
        stateful: bool = False,
        state: Optional[Union[StateConfig, Mapping[str, Any]]] = None,
        state_reset: str = "batch",  # 'batch' | 'epoch' | 'none'
        stream_lr: Optional[float] = None,
        output_shape: Optional[Tuple[int, ...]] = None,
        lsm: Optional[PreprocessorLike] = None,
        lsm_train: bool = False,
        lsm_pretrain_epochs: int = 0,
        lsm_lr: Optional[float] = None,
        warm_start: bool = False,
        scaler: Optional[ScalerSpec] = None,
        scaler_params: Optional[Dict[str, Any]] = None,
    ) -> None:
        self.hidden_layers = int(hidden_layers)

        hidden_units_res = resolve_int_alias(
            primary_value=hidden_units,
            alias_value=hidden_width,
            primary_name="hidden_units",
            alias_name="hidden_width",
            context="PSANNRegressor",
            default=64,
        )
        user_set_hidden_units = hidden_units_res.provided_primary
        units = hidden_units_res.value if hidden_units_res.value is not None else 64
        self.hidden_units = units
        self.hidden_width = units

        user_set_conv = conv_channels is not None
        if user_set_conv and not preserve_shape:
            warnings.warn(
                "`conv_channels` has no effect when preserve_shape=False; ignoring value.",
                UserWarning,
                stacklevel=2,
            )

        conv_val = conv_channels if user_set_conv else units
        if conv_val is None:
            conv_val = units
        conv_val = int(conv_val)
        if user_set_conv and user_set_hidden_units and conv_val != units:
            warnings.warn(
                "`conv_channels` differs from `hidden_units`; using `conv_channels` for convolutional paths.",
                UserWarning,
                stacklevel=2,
            )
        self.conv_channels = conv_val

        self.epochs = int(epochs)
        self.batch_size = int(batch_size)
        self.lr = float(lr)
        self.optimizer = str(optimizer)
        self.weight_decay = float(weight_decay)
        self.activation = activation or {}
        self.device = device
        self.random_state = random_state
        self.early_stopping = bool(early_stopping)
        self.patience = int(patience)
        self.num_workers = int(num_workers)
        self.loss = loss
        self.loss_params = loss_params
        self.loss_reduction = loss_reduction
        self.w0 = float(w0)
        self.preserve_shape = bool(preserve_shape)
        self.data_format = str(data_format)
        self.conv_kernel_size = int(conv_kernel_size)
        self.per_element = bool(per_element)
        self.activation_type = activation_type
        self.stateful = bool(stateful)
        self.state = ensure_state_config(state)
        self.state_reset = state_reset
        self.stream_lr = stream_lr
        self.output_shape = output_shape
        self.lsm = lsm
        self.lsm_train = bool(lsm_train)
        self.lsm_pretrain_epochs = int(lsm_pretrain_epochs)
        self.lsm_lr = lsm_lr
        self.warm_start = bool(warm_start)
        # Optional input scaler (minmax/standard or custom object with fit/transform)
        self.scaler = scaler
        self.scaler_params = scaler_params or None
        self._use_channel_first_train_inputs_ = False
        self._preproc_cfg_ = {'lsm': lsm, 'train': bool(lsm_train), 'pretrain_epochs': int(lsm_pretrain_epochs)}
        self._lsm_module_ = None
        self._hisso_cache_: Optional[np.ndarray] = None
        self._hisso_trainer_: Optional[Any] = None
        self._hisso_options_: Optional[Any] = None
        self._hisso_reward_fn_: Optional[Any] = None
        self._hisso_context_extractor_: Optional[Any] = None
        self._hisso_cfg_: Optional[Any] = None
        self._hisso_trained_ = False

        # Training state caches
        self._optimizer_: Optional[torch.optim.Optimizer] = None
        self._lr_scheduler_: Optional[Any] = None
        self._amp_scaler_: Optional[Any] = None
        self._training_state_token_: int = 0
        self._stream_opt_: Optional[torch.optim.Optimizer] = None
        self._stream_loss_: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None
        self._stream_model_token_: Optional[int] = None
        self._stream_last_lr_: Optional[float] = None

        # Fitted scaler state (set during fit)
        self._scaler_kind_: Optional[str] = None
        self._scaler_state_: Optional[Dict[str, Any]] = None

        # Inference metadata
        self._train_inputs_layout_: str = "flat"
        self._primary_dim_: Optional[int] = None
        self._output_dim_: Optional[int] = None
        self._target_cf_shape_: Optional[Tuple[int, ...]] = None
        self._target_vector_dim_: Optional[int] = None
        self._output_shape_tuple_: Optional[Tuple[int, ...]] = tuple(output_shape) if output_shape is not None else None

    @staticmethod
    def _normalize_param_aliases(params: Dict[str, Any]) -> Dict[str, Any]:
        out = dict(params)
        hidden_width_raw = out.pop("hidden_width", None)
        hidden_units_res = resolve_int_alias(
            primary_value=out.pop("hidden_units", None),
            alias_value=hidden_width_raw,
            primary_name="hidden_units",
            alias_name="hidden_width",
            context="PSANNRegressor.set_params",
        )
        if hidden_units_res.value is not None:
            out["hidden_units"] = hidden_units_res.value
            out["hidden_width"] = hidden_units_res.value
        elif hidden_width_raw is not None:
            out["hidden_width"] = hidden_width_raw

        conv_channels_res = resolve_int_alias(
            primary_value=out.pop("conv_channels", None),
            alias_value=out.pop("hidden_channels", None),
            primary_name="conv_channels",
            alias_name="hidden_channels",
            context="PSANNRegressor.set_params",
            default=out.get("hidden_units"),
        )
        if conv_channels_res.value is not None:
            out["conv_channels"] = conv_channels_res.value
        else:
            out.pop("conv_channels", None)
        return out



    # ------------------------- Scaling helpers -------------------------
    def _make_internal_scaler(self) -> Optional[Dict[str, Any]]:
        kind = self.scaler
        if kind is None:
            return None
        if isinstance(kind, str):
            key = kind.lower()
            if key not in {"standard", "minmax"}:
                raise ValueError("Unsupported scaler string. Use 'standard', 'minmax', or provide an object with fit/transform.")
            return {"type": key, "state": {}}
        # Custom object: must implement fit/transform; inverse_transform optional
        if not hasattr(kind, "fit") or not hasattr(kind, "transform"):
            raise ValueError("Custom scaler must implement fit(X) and transform(X). Optional inverse_transform(X).")
        return {"type": "custom", "obj": kind}

    def _scaler_fit_update(self, X2d: np.ndarray) -> Optional[Callable[[np.ndarray], np.ndarray]]:
        """Fit or update scaler on 2D array and return a transform function.

        - For built-in scalers, supports incremental update when warm_start=True.
        - For custom object, calls .fit on first time, else attempts partial_fit if available, else refit on concat.
        """
        if self.scaler is None:
            self._scaler_kind_ = None
            self._scaler_state_ = None
            return None
        spec = getattr(self, "_scaler_spec_", None)
        if spec is None:
            spec = self._make_internal_scaler()
            self._scaler_spec_ = spec
        if spec is None:
            self._scaler_kind_ = None
            self._scaler_state_ = None
            return None

        if spec.get("type") == "standard":
            self._scaler_kind_ = "standard"
            st = self._scaler_state_ or {"n": 0, "mean": None, "M2": None}
            n0 = int(st["n"])
            X = np.asarray(X2d, dtype=np.float32)
            # Welford online update per feature
            if n0 == 0:
                mean = X.mean(axis=0)
                diff = X - mean
                M2 = (diff * diff).sum(axis=0)
                n = X.shape[0]
            else:
                mean0 = st["mean"]
                M20 = st["M2"]
                n = n0 + X.shape[0]
                delta = X.mean(axis=0) - mean0
                mean = (mean0 * n0 + X.sum(axis=0)) / n
                # Update M2 across batches: combine variances
                # M2_total = M2_a + M2_b + delta^2 * n_a * n_b / n_total
                M2a = M20
                xa = n0
                xb = X.shape[0]
                X_centered = X - X.mean(axis=0)
                M2b = (X_centered * X_centered).sum(axis=0)
                M2 = M2a + M2b + (delta * delta) * (xa * xb) / max(n, 1)
            self._scaler_state_ = {"n": int(n), "mean": mean, "M2": M2}

            def _xfm(Z: np.ndarray) -> np.ndarray:
                st2 = self._scaler_state_
                assert st2 is not None
                mean2 = st2["mean"]
                var = st2["M2"] / max(st2["n"], 1)
                std = np.sqrt(np.maximum(var, 1e-8)).astype(np.float32)
                return (Z - mean2) / std

            return _xfm

        if spec.get("type") == "minmax":
            self._scaler_kind_ = "minmax"
            st = self._scaler_state_ or {"min": None, "max": None}
            X = np.asarray(X2d, dtype=np.float32)
            mn = X.min(axis=0) if st["min"] is None else np.minimum(st["min"], X.min(axis=0))
            mx = X.max(axis=0) if st["max"] is None else np.maximum(st["max"], X.max(axis=0))
            self._scaler_state_ = {"min": mn, "max": mx}

            def _xfm(Z: np.ndarray) -> np.ndarray:
                st2 = self._scaler_state_
                assert st2 is not None
                mn2 = st2["min"]
                mx2 = st2["max"]
                scale = np.where((mx2 - mn2) > 1e-8, (mx2 - mn2), 1.0)
                return (Z - mn2) / scale

            return _xfm

        # Custom object
        obj = spec.get("obj")
        self._scaler_kind_ = "custom"
        if not hasattr(self, "_scaler_fitted_") or not getattr(self, "_scaler_fitted_", False):
            # Fit once
            try:
                obj.fit(X2d, **(self.scaler_params or {}))
            except TypeError:
                obj.fit(X2d)
            self._scaler_fitted_ = True
        else:
            if hasattr(obj, "partial_fit"):
                obj.partial_fit(X2d)
            else:
                # Fallback: refit on concatenation of small cache if available
                pass

        def _xfm(Z: np.ndarray) -> np.ndarray:
            return obj.transform(Z)

        return _xfm

    def _scaler_inverse_tensor(self, X_ep: torch.Tensor, *, feature_dim: int = -1) -> torch.Tensor:
        """Inverse-transform a torch tensor episode if scaler is active.

        Expects features along last dim by default (B,T,D) or (N,D).
        """
        kind = getattr(self, "_scaler_kind_", None)
        st = getattr(self, "_scaler_state_", None)
        if kind is None:
            return X_ep
        if kind == "standard" and st is not None:
            mean = torch.as_tensor(st["mean"], device=X_ep.device, dtype=X_ep.dtype)
            var = torch.as_tensor(st["M2"] / max(st["n"], 1), device=X_ep.device, dtype=X_ep.dtype)
            std = torch.sqrt(torch.clamp(var, min=1e-8))
            return X_ep * std + mean
        if kind == "minmax" and st is not None:
            mn = torch.as_tensor(st["min"], device=X_ep.device, dtype=X_ep.dtype)
            mx = torch.as_tensor(st["max"], device=X_ep.device, dtype=X_ep.dtype)
            scale = torch.where((mx - mn) > 1e-8, (mx - mn), torch.ones_like(mx))
            return X_ep * scale + mn
        if kind == "custom" and hasattr(self.scaler, "inverse_transform"):
            # Fallback via CPU numpy; small overhead acceptable for context extraction
            X_np = X_ep.detach().cpu().numpy()
            X_inv = self.scaler.inverse_transform(X_np)
            return torch.as_tensor(X_inv, device=X_ep.device, dtype=X_ep.dtype)
        return X_ep

    def _apply_fitted_scaler(self, X2d: np.ndarray) -> np.ndarray:
        kind = getattr(self, "_scaler_kind_", None)
        if kind is None:
            return X2d.astype(np.float32, copy=False)
        st = getattr(self, "_scaler_state_", None)
        if kind == "standard" and st is not None:
            n = max(int(st.get("n", 0)), 1)
            mean = np.asarray(st.get("mean"), dtype=np.float32)
            var = np.asarray(st.get("M2"), dtype=np.float32) / float(n)
            std = np.sqrt(np.maximum(var, 1e-8)).astype(np.float32, copy=False)
            return ((X2d - mean) / std).astype(np.float32, copy=False)
        if kind == "minmax" and st is not None:
            mn = np.asarray(st.get("min"), dtype=np.float32)
            mx = np.asarray(st.get("max"), dtype=np.float32)
            scale = np.where((mx - mn) > 1e-8, (mx - mn), 1.0).astype(np.float32, copy=False)
            return ((X2d - mn) / scale).astype(np.float32, copy=False)
        if kind == "custom" and hasattr(self, "scaler") and hasattr(self.scaler, "transform"):
            transformed = self.scaler.transform(X2d)
            return np.asarray(transformed, dtype=np.float32)
        return X2d.astype(np.float32, copy=False)

    def _ensure_fitted(self) -> None:
        model = getattr(self, "model_", None)
        if model is None:
            raise RuntimeError("Estimator is not fitted yet; call 'fit' before inference.")
        if getattr(self, "input_shape_", None) is None:
            raise RuntimeError("Estimator is missing fitted input shape; inference is unavailable.")

    def _prepare_inference_inputs(self, X: np.ndarray) -> Tuple[np.ndarray, Dict[str, Any]]:
        self._ensure_fitted()
        X_arr = np.asarray(X, dtype=np.float32)
        if X_arr.ndim == len(self.input_shape_):
            X_arr = X_arr.reshape((1,) + tuple(self.input_shape_))
        if X_arr.ndim != len(self.input_shape_) + 1:
            raise ValueError(f"Expected input with {len(self.input_shape_) + 1} dimensions; received shape {X_arr.shape}.")
        expected = tuple(self.input_shape_)
        if tuple(X_arr.shape[1:]) != expected:
            raise ValueError(f"Input shape {X_arr.shape[1:]} does not match fitted shape {expected}.")

        meta: Dict[str, Any] = {
            "n_samples": int(X_arr.shape[0]),
            "data_format": self.data_format,
            "per_element": bool(self.per_element),
        }

        if not self.preserve_shape:
            X2d = self._flatten(X_arr)
            X2d = self._apply_fitted_scaler(X2d)
            meta["layout"] = "flat"
            return X2d, meta

        X_cf = np.moveaxis(X_arr, -1, 1) if self.data_format == "channels_last" else X_arr.copy()
        meta["cf_shape"] = tuple(X_cf.shape[1:])
        N, C = X_cf.shape[0], int(X_cf.shape[1])
        X2d = X_cf.reshape(N, C, -1).transpose(0, 2, 1).reshape(-1, C)
        X2d_scaled = self._apply_fitted_scaler(X2d)
        if X2d_scaled is not X2d:
            X_cf = X2d_scaled.reshape(N, -1, C).transpose(0, 2, 1).reshape(X_cf.shape)
        use_cf_inputs = bool(self.per_element or getattr(self, "_use_channel_first_train_inputs_", False))
        meta["layout"] = "cf" if use_cf_inputs else "flat"
        if use_cf_inputs:
            inputs_np = X_cf.astype(np.float32, copy=False)
        else:
            inputs_np = X2d_scaled.reshape(N, -1).astype(np.float32, copy=False)
        return inputs_np, meta

    def _reshape_predictions(self, preds: np.ndarray, meta: Dict[str, Any]) -> np.ndarray:
        n_samples = int(meta["n_samples"])
        preds = preds.astype(np.float32, copy=False)

        if self.preserve_shape and self.per_element:
            if self._target_cf_shape_ is not None:
                cf_shape = (n_samples,) + tuple(self._target_cf_shape_)
            elif getattr(self, "_internal_input_shape_cf_", None) is not None:
                spatial = tuple(self._internal_input_shape_cf_[1:])
                channels = preds.shape[1]
                cf_shape = (n_samples, channels) + spatial
            else:
                cf_shape = (n_samples, preds.shape[1])
            preds_cf = preds.reshape(cf_shape)
            if self.data_format == "channels_last" and preds_cf.ndim >= 3:
                preds_cf = np.moveaxis(preds_cf, 1, -1)
            return preds_cf

        if self._output_shape_tuple_ is not None:
            return preds.reshape((n_samples,) + self._output_shape_tuple_)

        if not self._keep_column_output_ and preds.shape[1] == 1:
            return preds.reshape(n_samples)

        return preds

    def _run_model(self, inputs_np: np.ndarray, *, state_updates: bool = False) -> np.ndarray:
        self._ensure_fitted()
        state_updates = bool(state_updates and self.stateful)
        model = self.model_
        if model is None:
            raise RuntimeError("Estimator is not fitted yet; no model available.")

        device = self._device()
        prev_training = model.training
        try:
            if state_updates:
                model.train(True)
                if hasattr(model, "set_state_updates"):
                    model.set_state_updates(True)
            else:
                model.eval()
                if hasattr(model, "set_state_updates"):
                    model.set_state_updates(False)

            model.to(device)
            with torch.no_grad():
                tensor = torch.from_numpy(inputs_np.astype(np.float32, copy=False)).to(device)
                outputs = model(tensor)
                return outputs.detach().cpu().numpy()
        finally:
            model.train(prev_training)
            if hasattr(model, "set_state_updates"):
                model.set_state_updates(bool(prev_training))
    # Internal helpers
    def _device(self) -> torch.device:
        return choose_device(self.device)

    def _infer_input_shape(self, X: np.ndarray) -> tuple:
        if X.ndim < 2:
            raise ValueError("X must be at least 2D (batch, features...)")
        return tuple(X.shape[1:])

    def _flatten(self, X: np.ndarray) -> np.ndarray:
        return X.reshape(X.shape[0], -1).astype(np.float32, copy=False)

    def _resolve_lsm_module(
        self,
        data: Any,
        *,
        preserve_shape: bool,
    ) -> Tuple[Optional[nn.Module], Optional[int]]:
        if self.lsm is None:
            self._lsm_module_ = None
            return None, None

        preproc, module = build_preprocessor(
            self.lsm,
            allow_train=data is not None,
            pretrain_epochs=self.lsm_pretrain_epochs,
            data=data,
        )
        if preproc is None:
            self._lsm_module_ = None
            return None, None

        self.lsm = preproc
        lsm_module = module if module is not None else (preproc if isinstance(preproc, nn.Module) else None)
        if lsm_module is None or not hasattr(lsm_module, 'forward'):
            raise RuntimeError(
                "Provided lsm preprocessor must expose a torch.nn.Module. Fit the expander or pass an nn.Module."
            )

        self._lsm_module_ = lsm_module
        if lsm_module is not None and preproc is not None:
            if hasattr(preproc, "score_reconstruction") and not hasattr(lsm_module, "score_reconstruction"):
                setattr(lsm_module, "score_reconstruction", preproc.score_reconstruction)
        attr = 'out_channels' if preserve_shape else 'output_dim'
        dim = getattr(lsm_module, attr, getattr(preproc, attr, None))
        return lsm_module, int(dim) if dim is not None else None

    def _build_dense_core(
        self,
        input_dim: int,
        output_dim: int,
        *,
        state_cfg: Optional[Dict[str, Any]] = None,
    ) -> nn.Module:
        return PSANNNet(
            int(input_dim),
            int(output_dim),
            hidden_layers=self.hidden_layers,
            hidden_units=self.hidden_units,
            hidden_width=self.hidden_width,
            act_kw=self.activation,
            state_cfg=state_cfg,
            activation_type=self.activation_type,
            w0=self.w0,
        )

    def _build_conv_core(
        self,
        spatial_ndim: int,
        in_channels: int,
        output_dim: int,
        *,
        segmentation_head: bool,
    ) -> nn.Module:
        conv_map = {
            1: PSANNConv1dNet,
            2: PSANNConv2dNet,
            3: PSANNConv3dNet,
        }
        conv_cls = conv_map.get(int(spatial_ndim))
        if conv_cls is None:
            raise ValueError(f"Unsupported spatial dimensionality {spatial_ndim}; expected 1, 2, or 3.")
        return conv_cls(
            int(in_channels),
            int(output_dim),
            hidden_layers=self.hidden_layers,
            conv_channels=self.conv_channels,
            hidden_channels=self.conv_channels,
            kernel_size=self.conv_kernel_size,
            act_kw=self.activation,
            activation_type=self.activation_type,
            w0=self.w0,
            segmentation_head=bool(segmentation_head),
        )


    def _make_optimizer(self, model: torch.nn.Module, lr: Optional[float] = None):
        lr = float(self.lr if lr is None else lr)
        if self.optimizer.lower() == "adamw":
            return torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=self.weight_decay)
        if self.optimizer.lower() == "sgd":
            return torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
        return torch.optim.Adam(model.parameters(), lr=lr, weight_decay=self.weight_decay)

    def _build_optimizer(self, model: torch.nn.Module) -> torch.optim.Optimizer:
        """Compatibility helper for warm-start flows expecting estimator-owned builders."""

        return _build_optimizer_helper(self, model)

    def _make_loss(self):
        # Built-in strings
        if isinstance(self.loss, str):
            name = self.loss.lower()
            params = self.loss_params or {}
            reduction = self.loss_reduction
            if name in ("l1", "mae"):
                return torch.nn.L1Loss(reduction=reduction)
            if name in ("mse", "l2"):
                return torch.nn.MSELoss(reduction=reduction)
            if name in ("smooth_l1", "huber_smooth"):
                beta = float(params.get("beta", 1.0))
                return torch.nn.SmoothL1Loss(beta=beta, reduction=reduction)
            if name in ("huber",):
                delta = float(params.get("delta", 1.0))
                return torch.nn.HuberLoss(delta=delta, reduction=reduction)
            raise ValueError(f"Unknown loss '{self.loss}'. Supported: mse, l1/mae, smooth_l1, huber, or a callable.")

        # Callable custom loss; may return tensor (any shape) or float
        if callable(self.loss):
            user_fn = self.loss
            params = self.loss_params or {}
            reduction = self.loss_reduction

            def _loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                out = user_fn(pred, target, **params) if params else user_fn(pred, target)
                if not isinstance(out, torch.Tensor):
                    out = torch.as_tensor(out, dtype=pred.dtype, device=pred.device)
                if out.ndim == 0:
                    return out
                if reduction == "mean":
                    return out.mean()
                if reduction == "sum":
                    return out.sum()
                if reduction == "none":
                    return out
                raise ValueError(f"Unsupported reduction '{reduction}' for custom loss")

            return _loss

        raise TypeError("loss must be a string or a callable returning a scalar tensor")

    def _make_per_element_fit_hooks(self) -> FitVariantHooks:
        def build_model(request: ModelBuildRequest) -> nn.Module:
            prepared = request.prepared
            X_cf = prepared.train_inputs
            if not isinstance(X_cf, np.ndarray):
                X_cf = prepared.X_cf
            if X_cf is None:
                raise ValueError("PreparedInputState missing channel-first inputs for per-element training.")
            nd = int(X_cf.ndim) - 2
            if nd < 1:
                raise ValueError("per_element=True expects inputs with spatial dimensions.")

            in_channels = int(X_cf.shape[1])
            lsm_module = request.lsm_module
            if lsm_module is not None:
                if request.lsm_output_dim is not None:
                    in_channels = int(request.lsm_output_dim)
                elif hasattr(lsm_module, "out_channels"):
                    in_channels = int(getattr(lsm_module, "out_channels"))
                elif hasattr(self.lsm, "out_channels"):
                    in_channels = int(getattr(self.lsm, "out_channels"))

            core = self._build_conv_core(
                nd,
                in_channels,
                int(prepared.output_dim),
                segmentation_head=True,
            )

            preproc = lsm_module
            if preproc is not None and not self.lsm_train:
                for param in preproc.parameters():
                    param.requires_grad = False
            return WithPreprocessor(preproc, core)

        return FitVariantHooks(build_model=build_model)

    def _make_conv_fit_hooks(
        self,
        *,
        prepared: PreparedInputState,
        verbose: int,
    ) -> FitVariantHooks:
        internal_shape = prepared.internal_shape_cf
        if internal_shape is None:
            raise ValueError("PreparedInputState missing channels-first shape for conv training.")
        spatial_ndim = max(1, len(internal_shape) - 1)
        base_channels = int(internal_shape[0])

        def build_model(request: ModelBuildRequest) -> nn.Module:
            lsm_module = request.lsm_module
            in_channels = base_channels
            if lsm_module is not None:
                if request.lsm_output_dim is not None:
                    in_channels = int(request.lsm_output_dim)
                elif hasattr(lsm_module, "out_channels"):
                    in_channels = int(getattr(lsm_module, "out_channels"))
                elif hasattr(self.lsm, "out_channels"):
                    in_channels = int(getattr(self.lsm, "out_channels"))

            core = self._build_conv_core(
                spatial_ndim,
                in_channels,
                int(prepared.output_dim),
                segmentation_head=bool(self.per_element),
            )

            preproc = lsm_module
            if preproc is not None and not self.lsm_train:
                for param in preproc.parameters():
                    param.requires_grad = False
            return WithPreprocessor(preproc, core)

        def build_hisso_plan(
            estimator_ref: "ResConvPSANNRegressor",
            request: ModelBuildRequest,
            *,
            fit_args: NormalisedFitArgs,
        ) -> Optional[HISSOTrainingPlan]:
            if estimator_ref.per_element:
                raise ValueError("hisso=True currently supports per_element=False for ResConvPSANNRegressor")
            prepared_local = request.prepared
            inputs_cf = prepared_local.train_inputs
            if not isinstance(inputs_cf, np.ndarray):
                inputs_cf = prepared_local.X_cf
            if inputs_cf is None:
                raise ValueError("PreparedInputState missing channel-first inputs for conv HISSO training.")
            if fit_args.hisso_options is None:
                raise ValueError("HISSO options were not prepared despite hisso=True.")
            return build_hisso_training_plan(
                estimator_ref,
                train_inputs=inputs_cf,
                primary_dim=int(request.primary_dim),
                fit_args=fit_args,
                options=fit_args.hisso_options,
                lsm_module=request.lsm_module,
            )

        return FitVariantHooks(
            build_model=build_model,
            build_hisso_plan=build_hisso_plan,
        )

    def _make_flatten_fit_hooks(
        self,
        *,
        prepared: PreparedInputState,
        verbose: int,
    ) -> FitVariantHooks:
        train_inputs = prepared.train_inputs
        if train_inputs is None:
            raise ValueError("PreparedInputState missing flattened training inputs.")
        if not isinstance(train_inputs, np.ndarray):
            raise ValueError("PreparedInputState.train_inputs must be a numpy array for flat training.")

        def build_model(request: ModelBuildRequest) -> nn.Module:
            prepared_local = request.prepared
            inputs_arr = prepared_local.train_inputs
            if inputs_arr is None or not isinstance(inputs_arr, np.ndarray):
                raise ValueError("PreparedInputState missing train_inputs for model construction.")
            if request.lsm_output_dim is not None:
                input_dim = int(request.lsm_output_dim)
            else:
                input_dim = int(inputs_arr.shape[1])
            core = self._build_dense_core(
                input_dim,
                int(prepared_local.output_dim),
                state_cfg=(self.state if self.stateful else None),
            )
            preproc = request.lsm_module
            if preproc is not None and not self.lsm_train:
                for param in preproc.parameters():
                    param.requires_grad = False
            if preproc is None:
                return core
            return WithPreprocessor(preproc, core)

        def build_hisso_plan(
            estimator_ref: "PSANNRegressor",
            request: ModelBuildRequest,
            *,
            fit_args: NormalisedFitArgs,
        ) -> Optional[HISSOTrainingPlan]:
            if not fit_args.hisso:
                return None
            if fit_args.hisso_options is None:
                raise ValueError("HISSO options were not prepared despite hisso=True.")
            inputs_arr = request.prepared.train_inputs
            if inputs_arr is None or not isinstance(inputs_arr, np.ndarray):
                inputs_arr = request.prepared.X_flat
                if inputs_arr is None:
                    raise ValueError("PreparedInputState missing inputs for HISSO planning.")
            return build_hisso_training_plan(
                estimator_ref,
                train_inputs=inputs_arr,
                primary_dim=int(request.primary_dim),
                fit_args=fit_args,
                options=fit_args.hisso_options,
                lsm_module=request.lsm_module,
            )

        return FitVariantHooks(
            build_model=build_model,
            build_hisso_plan=build_hisso_plan,
        )

    def _make_fit_hooks(
        self,
        *,
        prepared: PreparedInputState,
        verbose: int,
    ) -> FitVariantHooks:
        if self.preserve_shape and self.per_element:
            return self._make_per_element_fit_hooks()
        if self.preserve_shape and getattr(self, "_use_channel_first_train_inputs_", False):
            return self._make_conv_fit_hooks(prepared=prepared, verbose=verbose)
        return self._make_flatten_fit_hooks(prepared=prepared, verbose=verbose)

    # Estimator API

    def fit(
        self,
        X: np.ndarray,
        y: np.ndarray | None,
        *,
        validation_data: Optional[ValidationDataLike] = None,
        verbose: int = 0,
        noisy: Optional[NoiseSpec] = None,
        hisso: bool = False,
        hisso_window: Optional[int] = None,
        hisso_reward_fn: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
        hisso_context_extractor: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        hisso_primary_transform: Optional[str] = None,
        hisso_transition_penalty: Optional[float] = None,
        hisso_trans_cost: Optional[float] = None,
        hisso_supervised: Optional[Mapping[str, Any] | bool] = None,
        lr_max: Optional[float] = None,
        lr_min: Optional[float] = None,
    ):
        """Fit the estimator.

        Parameters
        - X: np.ndarray
            Training inputs. Shapes:
              - MLP/flattened: (N, F1, ..., Fk) flattened internally to (N, prod(F*))
              - preserve_shape=True: (N, C, ...) or (N, ..., C) depending on data_format
        - y: np.ndarray
            Targets. Shapes:
              - vector/pooled head: (N, T) or (N,) where T=prod(output_shape) if provided
              - per_element=True: (N, C_out, ...) or (N, ..., C_out) matching X spatial dims
        - validation_data: optional (X_val, y_val) tuple for early stopping/logging
        - verbose: 0/1 to control epoch logging
        - noisy: optional Gaussian input noise std; scalar, per-feature vector, or tensor matching input shape
        - hisso_supervised: optional bool or dict to run a supervised warm start before HISSO (requires providing 'y')
        - hisso: if True, train via Horizon-Informed Sampling Strategy Optimization (episodic reward)
        - hisso_window: episode/window length for HISSO (default 64)
        - hisso_primary_transform: optional transform ('identity' | 'softmax' | 'tanh') applied to primary outputs before reward evaluation
        - hisso_transition_penalty: optional float penalty applied between HISSO steps (deprecated alias `hisso_trans_cost` still accepted)
        """
        seed_all(self.random_state)

        if hisso and self.per_element:
            raise ValueError("hisso=True currently supports per_element=False.")

        # Reset HISSO-specific runtime artefacts before launching a new fit.
        self._hisso_options_ = None
        self._hisso_reward_fn_ = None
        self._hisso_context_extractor_ = None
        self._hisso_cfg_ = None
        self._hisso_trainer_ = None
        self._hisso_trained_ = False

        fit_args = normalise_fit_args(
            self,
            X,
            y,
            validation_data=validation_data,
            noisy=noisy,
            verbose=verbose,
            lr_max=lr_max,
            lr_min=lr_min,
            hisso=hisso,
            hisso_kwargs={
                "hisso_window": hisso_window,
                "hisso_reward_fn": hisso_reward_fn,
                "hisso_context_extractor": hisso_context_extractor,
                "hisso_primary_transform": hisso_primary_transform,
                "hisso_transition_penalty": hisso_transition_penalty,
                "hisso_trans_cost": hisso_trans_cost,
                "hisso_supervised": hisso_supervised,
            },
        )

        verbose = fit_args.verbose
        hisso = fit_args.hisso

        X = fit_args.X
        y = fit_args.y
        self._keep_column_output_ = bool(y is not None and y.ndim > 1)

        prepared_state, primary_dim = prepare_inputs_and_scaler(
            self,
            fit_args,
        )
        primary_dim = int(primary_dim)
        self._primary_dim_ = primary_dim
        self._output_dim_ = int(prepared_state.output_dim)
        layout = "cf" if (self.preserve_shape and (self.per_element or getattr(self, "_use_channel_first_train_inputs_", False))) else "flat"
        self._train_inputs_layout_ = layout
        self._target_cf_shape_ = (
            tuple(prepared_state.y_cf.shape[1:]) if prepared_state.y_cf is not None else self._target_cf_shape_
        )
        self._target_vector_dim_ = (
            int(prepared_state.y_vector.shape[1]) if prepared_state.y_vector is not None else self._target_vector_dim_
        )

        preserve_inputs = bool(self.preserve_shape and self.per_element)
        lsm_data = prepared_state.train_inputs
        lsm_model, lsm_dim = self._resolve_lsm_module(lsm_data, preserve_shape=preserve_inputs)

        hooks = self._make_fit_hooks(prepared=prepared_state, verbose=verbose)

        request = ModelBuildRequest(
            estimator=self,
            prepared=prepared_state,
            primary_dim=primary_dim,
            lsm_module=lsm_model,
            lsm_output_dim=lsm_dim,
            preserve_shape=bool(self.preserve_shape),
        )

        rebuild = not (self.warm_start and isinstance(getattr(self, "model_", None), nn.Module))
        if rebuild:
            self.model_ = build_model_from_hooks(hooks, request)

        device = self._device()
        self.model_.to(device)

        if hisso:
            result = maybe_run_hisso(hooks, request, fit_args=fit_args)
            if result is None:
                raise RuntimeError("HISSO requested but no variant hook was provided.")
            return self

        run_supervised_training(
            self,
            self.model_,
            prepared_state,
            fit_args=fit_args,
        )
        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        inputs_np, meta = self._prepare_inference_inputs(X)
        preds = self._run_model(inputs_np, state_updates=False)
        return self._reshape_predictions(preds, meta)

    def score(self, X: np.ndarray, y: np.ndarray) -> float:
        y_true = np.asarray(y, dtype=np.float32)
        y_pred = self.predict(X)
        if y_true.ndim == 1 and y_pred.ndim == 2 and y_pred.shape[1] == 1:
            y_pred = y_pred.reshape(-1)
        elif y_true.shape != y_pred.shape:
            y_pred = y_pred.reshape(y_true.shape)
        return float(_sk_r2_score(y_true, y_pred))

    def reset_state(self) -> None:
        self._ensure_fitted()
        if hasattr(self.model_, "reset_state"):
            self.model_.reset_state()

    def step(
        self,
        x: np.ndarray,
        *,
        target: Optional[np.ndarray] = None,
        update_params: bool = False,
        update_state: bool = True,
    ) -> Any:
        batch = np.asarray(x, dtype=np.float32)
        if batch.ndim == len(self.input_shape_):
            batch = batch.reshape((1,) + tuple(self.input_shape_))
        elif batch.ndim == len(self.input_shape_) + 1 and batch.shape[0] == 1:
            batch = batch.reshape((1,) + tuple(self.input_shape_))
        elif batch.ndim != len(self.input_shape_) + 1:
            raise ValueError(f"Expected input with {len(self.input_shape_) + 1} dims; received shape {batch.shape}.")

        inputs_np, meta = self._prepare_inference_inputs(batch)
        preds = self._run_model(inputs_np, state_updates=bool(update_state))
        reshaped = self._reshape_predictions(preds, meta)
        if update_params:
            if target is None:
                raise ValueError("step(..., update_params=True) requires a target array.")
            self._apply_stream_update(inputs_np, target=target)
        if isinstance(reshaped, np.ndarray):
            if reshaped.shape[0] == 1:
                return reshaped[0]
            if reshaped.ndim == 0:
                return float(reshaped)
        return reshaped

    def predict_sequence(
        self,
        X: np.ndarray,
        *,
        reset_state: bool = False,
        return_sequence: bool = False,
        update_state: bool = True,
    ) -> Any:
        return self._sequence_rollout(
            X,
            targets=None,
            reset_state=reset_state,
            update_params=False,
            update_state=update_state,
            return_sequence=return_sequence,
        )

    def predict_sequence_online(
        self,
        X: np.ndarray,
        y: np.ndarray,
        *,
        reset_state: bool = True,
        return_sequence: bool = True,
        update_state: bool = True,
    ) -> np.ndarray:
        """Teacher-forced rollout with per-step streaming updates."""

        return self._sequence_rollout(
            X,
            targets=y,
            reset_state=reset_state,
            update_params=True,
            update_state=update_state,
            return_sequence=return_sequence,
        )

    # ------------------------------------------------------------------
    # Internal helpers for stateful rollouts
    # ------------------------------------------------------------------

    def _sequence_rollout(
        self,
        X_seq: np.ndarray,
        *,
        targets: Optional[np.ndarray],
        reset_state: bool,
        update_params: bool,
        update_state: bool,
        return_sequence: bool,
    ) -> Any:
        self._ensure_fitted()
        sequence = self._coerce_sequence_inputs(X_seq)
        steps = int(sequence.shape[0])
        if steps == 0:
            raise ValueError("predict_sequence requires at least one timestep.")

        targets_arr: Optional[np.ndarray] = None
        if targets is not None:
            targets_arr = self._coerce_sequence_targets(targets, steps)
        if update_params and targets_arr is None:
            raise ValueError("Streaming rollouts require targets when update_params=True.")

        if reset_state:
            self.reset_state()

        outputs: list[Any] = []
        for idx in range(steps):
            tgt_step = None if targets_arr is None else targets_arr[idx]
            outputs.append(
                self.step(
                    sequence[idx],
                    target=tgt_step,
                    update_params=bool(update_params and targets_arr is not None),
                    update_state=update_state,
                )
            )

        if not return_sequence:
            return outputs[-1]

        stacked_inputs = [np.asarray(out, dtype=np.float32) for out in outputs]
        try:
            return np.stack(stacked_inputs, axis=0)
        except ValueError as exc:
            raise RuntimeError("Sequence outputs have inconsistent shapes; cannot stack results.") from exc

    def _coerce_sequence_inputs(self, sequence: np.ndarray) -> np.ndarray:
        seq = np.asarray(sequence, dtype=np.float32)
        expected_shape = tuple(self.input_shape_)

        if seq.ndim == len(expected_shape):
            seq = seq.reshape((1,) + expected_shape)
        elif seq.ndim == len(expected_shape) + 2 and seq.shape[0] == 1:
            seq = seq.reshape((-1,) + expected_shape)
        elif seq.ndim != len(expected_shape) + 1:
            raise ValueError(
                "Expected sequence shaped (T, ...) optionally preceded by a singleton batch; "
                f"received array with shape {seq.shape}."
            )

        if seq.shape[1:] != expected_shape:
            raise ValueError(
                f"Sequence feature layout {seq.shape[1:]} does not match fitted shape {expected_shape}."
            )

        return seq

    def _coerce_sequence_targets(self, targets: np.ndarray, steps: int) -> np.ndarray:
        arr = np.asarray(targets, dtype=np.float32)
        if arr.ndim == 0:
            if steps != 1:
                raise ValueError("Scalar targets are only valid for single-step rollouts.")
            return arr.reshape(1, 1)
        if arr.ndim >= 2 and arr.shape[0] == 1 and arr.shape[1] == steps:
            arr = arr.reshape(steps, *arr.shape[2:])
        elif arr.shape[0] != steps:
            raise ValueError(
                f"Targets length {arr.shape[0]} does not match sequence length {steps}."
            )
        return arr

    def _ensure_streaming_ready(self) -> None:
        if self.stream_lr is None or float(self.stream_lr) <= 0.0:
            raise RuntimeError("Streaming updates require 'stream_lr' > 0. Configure the estimator accordingly.")
        self._ensure_fitted()
        model = self.model_
        if model is None:
            raise RuntimeError("Estimator is not fitted yet; cannot initialise streaming optimiser.")

        needs_rebuild = self._stream_opt_ is None or self._stream_model_token_ != id(model)
        if needs_rebuild:
            self._stream_opt_ = self._make_optimizer(model, lr=float(self.stream_lr))
            self._stream_loss_ = self._make_loss()
            self._stream_model_token_ = id(model)
            self._stream_last_lr_ = float(self.stream_lr)
        elif self._stream_last_lr_ is None or self._stream_last_lr_ != float(self.stream_lr):
            assert self._stream_opt_ is not None
            for group in self._stream_opt_.param_groups:
                group["lr"] = float(self.stream_lr)
            self._stream_last_lr_ = float(self.stream_lr)

        if self._stream_loss_ is None:
            self._stream_loss_ = self._make_loss()

    def _coerce_stream_target(
        self,
        target: np.ndarray,
        reference: torch.Tensor,
        device: torch.device,
    ) -> torch.Tensor:
        arr = np.asarray(target, dtype=np.float32)
        expected_shape = tuple(reference.shape)

        if arr.shape != expected_shape:
            if reference.ndim == 2 and reference.shape[0] == 1:
                feat_dim = int(reference.shape[1])
                if arr.ndim == 1 and arr.shape[0] == feat_dim:
                    arr = arr.reshape(1, feat_dim)
                elif arr.ndim == 0 and feat_dim == 1:
                    arr = np.asarray([arr], dtype=np.float32).reshape(1, 1)
                elif arr.size == reference.numel():
                    arr = arr.reshape(expected_shape)
            elif arr.size == reference.numel():
                arr = arr.reshape(expected_shape)

        if arr.shape != expected_shape:
            raise ValueError(
                f"Streaming target shape {arr.shape} does not match prediction shape {expected_shape}."
            )

        return torch.from_numpy(arr.astype(np.float32, copy=False)).to(device)

    def _apply_stream_update(self, inputs_np: np.ndarray, *, target: np.ndarray) -> None:
        self._ensure_streaming_ready()
        model = self.model_
        if model is None:
            raise RuntimeError("Estimator is not fitted yet; cannot apply streaming update.")

        device = self._device()
        model.to(device)
        optimizer = self._stream_opt_
        loss_fn = self._stream_loss_
        if optimizer is None or loss_fn is None:
            raise RuntimeError("Streaming optimiser state is missing; call _ensure_streaming_ready first.")

        prev_mode = model.training
        prev_state_updates = None
        if hasattr(model, "set_state_updates"):
            prev_state_updates = getattr(model, "enable_state_updates", None)
            model.set_state_updates(False)

        try:
            model.train(True)
            xb = torch.from_numpy(inputs_np.astype(np.float32, copy=False)).to(device)
            pred = model(xb)
            target_t = self._coerce_stream_target(target, pred, device)

            optimizer.zero_grad(set_to_none=True)
            loss = loss_fn(pred, target_t)
            loss.backward()
            optimizer.step()

            if hasattr(model, "commit_state_updates"):
                model.commit_state_updates()
        finally:
            if hasattr(model, "set_state_updates"):
                if prev_state_updates is None:
                    model.set_state_updates(True)
                else:
                    model.set_state_updates(bool(prev_state_updates))
            model.train(prev_mode)

    def save(self, path: str) -> None:
        self._ensure_fitted()
        model = self.model_
        orig_device = torch.device("cpu")
        for param in model.parameters():
            orig_device = param.device
            break
        model_cpu = copy.deepcopy(model).cpu()
        payload = {
            "class": self.__class__.__name__,
            "params": self.get_params(deep=True),
            "model": model_cpu,
            "scaler_kind": getattr(self, "_scaler_kind_", None),
            "scaler_state": getattr(self, "_scaler_state_", None),
            "scaler_spec": getattr(self, "_scaler_spec_", None),
            "scaler_obj": self.scaler if getattr(self, "_scaler_kind_", None) == "custom" else None,
            "input_shape": tuple(self.input_shape_) if getattr(self, "input_shape_", None) is not None else None,
            "internal_shape_cf": (
                tuple(self._internal_input_shape_cf_) if getattr(self, "_internal_input_shape_cf_", None) is not None else None
            ),
            "primary_dim": self._primary_dim_,
            "output_dim": self._output_dim_,
            "keep_column_output": bool(getattr(self, "_keep_column_output_", False)),
            "train_layout": self._train_inputs_layout_,
            "target_cf_shape": self._target_cf_shape_,
            "target_vector_dim": self._target_vector_dim_,
            "output_shape_tuple": self._output_shape_tuple_,
        }
        torch.save(payload, path)
        model.to(orig_device)

    @classmethod
    def load(
        cls,
        path: str,
        *,
        map_location: Optional[Union[str, torch.device]] = "cpu",
    ) -> "PSANNRegressor":
        try:
            payload = torch.load(path, map_location=map_location, weights_only=False)
        except TypeError:
            payload = torch.load(path, map_location=map_location)
        class_name = payload.get("class")
        if class_name is not None and class_name != cls.__name__:
            raise ValueError(f"Checkpoint was created for '{class_name}', cannot load into '{cls.__name__}'.")
        params = payload.get("params", {})
        estimator = cls(**params)
        if "model" not in payload:
            raise RuntimeError("Checkpoint is missing model weights.")
        estimator.model_ = payload["model"]
        estimator.model_.to(estimator._device())
        estimator.model_.eval()

        estimator._scaler_kind_ = payload.get("scaler_kind")
        estimator._scaler_state_ = payload.get("scaler_state")
        estimator._scaler_spec_ = payload.get("scaler_spec")
        scaler_obj = payload.get("scaler_obj")
        if scaler_obj is not None:
            estimator.scaler = scaler_obj
            estimator._scaler_fitted_ = True

        input_shape = payload.get("input_shape")
        estimator.input_shape_ = tuple(input_shape) if input_shape is not None else None
        internal_cf = payload.get("internal_shape_cf")
        estimator._internal_input_shape_cf_ = tuple(internal_cf) if internal_cf is not None else None
        estimator._primary_dim_ = payload.get("primary_dim")
        estimator._output_dim_ = payload.get("output_dim")
        estimator._keep_column_output_ = bool(payload.get("keep_column_output", False))
        estimator._train_inputs_layout_ = payload.get("train_layout", "flat")
        target_cf = payload.get("target_cf_shape")
        estimator._target_cf_shape_ = tuple(target_cf) if target_cf is not None else None
        estimator._target_vector_dim_ = payload.get("target_vector_dim")
        output_shape_tuple = payload.get("output_shape_tuple")
        estimator._output_shape_tuple_ = tuple(output_shape_tuple) if output_shape_tuple is not None else None

        estimator._hisso_trainer_ = None
        return estimator

class WaveResNetRegressor(PSANNRegressor):
    """Sklearn-style regressor that wraps the WaveResNet backbone."""

    def __init__(
        self,
        *,
        hidden_layers: int = 6,
        hidden_width: Optional[int] = None,
        hidden_units: Optional[int] = None,
        epochs: int = 200,
        batch_size: int = 128,
        lr: float = 1e-3,
        optimizer: str = "adam",
        weight_decay: float = 0.0,
        activation: Optional[ActivationConfig] = None,
        device: str | torch.device = "auto",
        random_state: Optional[int] = None,
        early_stopping: bool = False,
        patience: int = 20,
        num_workers: int = 0,
        loss: LossLike = "mse",
        loss_params: Optional[Dict[str, Any]] = None,
        loss_reduction: str = "mean",
        w0: float = 30.0,
        preserve_shape: bool = False,
        data_format: str = "channels_first",
        conv_kernel_size: int = 1,
        conv_channels: Optional[int] = None,
        per_element: bool = False,
        activation_type: str = "psann",
        stateful: bool = False,
        state: Optional[Union[StateConfig, Mapping[str, Any]]] = None,
        state_reset: str = "batch",
        stream_lr: Optional[float] = None,
        output_shape: Optional[Tuple[int, ...]] = None,
        lsm: Optional[PreprocessorLike] = None,
        lsm_train: bool = False,
        lsm_pretrain_epochs: int = 0,
        lsm_lr: Optional[float] = None,
        warm_start: bool = False,
        scaler: Optional[ScalerSpec] = None,
        scaler_params: Optional[Dict[str, Any]] = None,
        first_layer_w0: float = 30.0,
        hidden_w0: float = 1.0,
        norm: Literal["none", "weight", "rms"] = "none",
        use_film: bool = True,
        use_phase_shift: bool = True,
        dropout: float = 0.0,
        context_dim: Optional[int] = None,
    ) -> None:
        if preserve_shape:
            raise ValueError("WaveResNetRegressor currently supports preserve_shape=False.")
        if per_element:
            raise ValueError("WaveResNetRegressor does not support per_element=True.")
        if conv_channels is not None:
            warnings.warn(
                "conv_channels has no effect for WaveResNetRegressor; ignoring value.",
                RuntimeWarning,
                stacklevel=2,
            )
        if conv_kernel_size != 1:
            warnings.warn(
                "conv_kernel_size has no effect for WaveResNetRegressor; ignoring value.",
                RuntimeWarning,
                stacklevel=2,
            )

        norm_value = str(norm).lower()
        if norm_value not in {"none", "weight", "rms"}:
            raise ValueError("norm must be one of {'none', 'weight', 'rms'}.")

        if context_dim is None:
            context_val = None
        else:
            context_val = int(context_dim)
            if context_val <= 0:
                raise ValueError("context_dim must be positive when provided.")

        if stateful or state is not None:
            warnings.warn(
                "WaveResNetRegressor does not support stateful configurations; ignoring state/stateful.",
                RuntimeWarning,
                stacklevel=2,
            )
        stateful_flag = False
        state_cfg = None

        super().__init__(
            hidden_layers=hidden_layers,
            hidden_width=hidden_width,
            hidden_units=hidden_units,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            optimizer=optimizer,
            weight_decay=weight_decay,
            activation=activation,
            device=device,
            random_state=random_state,
            early_stopping=early_stopping,
            patience=patience,
            num_workers=num_workers,
            loss=loss,
            loss_params=loss_params,
            loss_reduction=loss_reduction,
            w0=w0,
            preserve_shape=False,
            data_format=data_format,
            conv_kernel_size=1,
            conv_channels=None,
            per_element=False,
            activation_type=activation_type,
            stateful=stateful_flag,
            state=state_cfg,
            state_reset=state_reset,
            stream_lr=stream_lr,
            output_shape=output_shape,
            lsm=lsm,
            lsm_train=lsm_train,
            lsm_pretrain_epochs=lsm_pretrain_epochs,
            lsm_lr=lsm_lr,
            warm_start=warm_start,
            scaler=scaler,
            scaler_params=scaler_params,
        )

        self.first_layer_w0 = float(first_layer_w0)
        self.hidden_w0 = float(hidden_w0)
        self.norm = norm_value
        self.use_film = bool(use_film)
        self.use_phase_shift = bool(use_phase_shift)
        self.dropout = float(dropout)
        self.context_dim = context_val

        self._wave_hidden_dim = int(self.hidden_units)

    def _build_dense_core(
        self,
        input_dim: int,
        output_dim: int,
        *,
        state_cfg: Optional[Dict[str, Any]] = None,
    ) -> nn.Module:
        if state_cfg is not None:
            warnings.warn(
                "WaveResNetRegressor ignores state_cfg; WaveResNet does not expose external state.",
                RuntimeWarning,
                stacklevel=2,
            )
        return WaveResNet(
            input_dim=int(input_dim),
            hidden_dim=int(self._wave_hidden_dim),
            depth=int(self.hidden_layers),
            output_dim=int(output_dim),
            first_layer_w0=self.first_layer_w0,
            hidden_w0=self.hidden_w0,
            context_dim=self.context_dim,
            norm=self.norm,
            use_film=self.use_film,
            use_phase_shift=self.use_phase_shift,
            dropout=self.dropout,
        )


class ResPSANNRegressor(PSANNRegressor):
    """Sklearn-style regressor using ResidualPSANNNet core.

    Adds residual-specific args while keeping .fit/.predict API identical,
    including HISSO training.
    """

    def __init__(
        self,
        *,
        hidden_layers: int = 8,
        hidden_width: Optional[int] = None,
        hidden_units: Optional[int] = None,
        epochs: int = 200,
        batch_size: int = 128,
        lr: float = 1e-3,
        optimizer: str = "adam",
        weight_decay: float = 0.0,
        activation: Optional[ActivationConfig] = None,
        device: str | torch.device = "auto",
        random_state: Optional[int] = None,
        early_stopping: bool = False,
        patience: int = 20,
        num_workers: int = 0,
        loss: LossLike = "mse",
        loss_params: Optional[Dict[str, Any]] = None,
        loss_reduction: str = "mean",
        # maintained for parity; not used in residual core
        w0: float = 30.0,
        preserve_shape: bool = False,
        data_format: str = "channels_first",
        conv_kernel_size: int = 1,
        conv_channels: Optional[int] = None,
        per_element: bool = False,
        activation_type: str = "psann",
        stateful: bool = False,
        state: Optional[Union[StateConfig, Mapping[str, Any]]] = None,
        state_reset: str = "batch",
        stream_lr: Optional[float] = None,
        output_shape: Optional[Tuple[int, ...]] = None,
        lsm: Optional[PreprocessorLike] = None,
        lsm_train: bool = False,
        lsm_pretrain_epochs: int = 0,
        lsm_lr: Optional[float] = None,
        warm_start: bool = False,
        scaler: Optional[ScalerSpec] = None,
        scaler_params: Optional[Dict[str, Any]] = None,
        # residual-specific
        w0_first: float = 12.0,
        w0_hidden: float = 1.0,
        norm: str = "rms",
        drop_path_max: float = 0.0,
        residual_alpha_init: float = 0.0,
    ) -> None:
        super().__init__(
            hidden_layers=hidden_layers,
            hidden_width=hidden_width,
            hidden_units=hidden_units,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            optimizer=optimizer,
            weight_decay=weight_decay,
            activation=activation,
            device=device,
            random_state=random_state,
            early_stopping=early_stopping,
            patience=patience,
            num_workers=num_workers,
            loss=loss,
            loss_params=loss_params,
            loss_reduction=loss_reduction,
            w0=w0,
            preserve_shape=preserve_shape,
            data_format=data_format,
            conv_kernel_size=conv_kernel_size,
            conv_channels=conv_channels,
            per_element=per_element,
            activation_type=activation_type,
            stateful=stateful,
            state=state,
            state_reset=state_reset,
            stream_lr=stream_lr,
            output_shape=output_shape,
            lsm=lsm,
            lsm_train=lsm_train,
            lsm_pretrain_epochs=lsm_pretrain_epochs,
            lsm_lr=lsm_lr,
            warm_start=warm_start,
            scaler=scaler,
            scaler_params=scaler_params,
        )
        self.w0_first = float(w0_first)
        self.w0_hidden = float(w0_hidden)
        self.norm = str(norm)
        self.drop_path_max = float(drop_path_max)
        self.residual_alpha_init = float(residual_alpha_init)
        if self.preserve_shape:
            self._use_channel_first_train_inputs_ = True

    def _build_dense_core(
        self,
        input_dim: int,
        output_dim: int,
        *,
        state_cfg: Optional[Dict[str, Any]] = None,
    ) -> nn.Module:
        if state_cfg is not None:
            warnings.warn(
                "ResidualPSANNNet does not currently support stateful configurations; ignoring state_cfg.",
                RuntimeWarning,
                stacklevel=2,
            )
        return ResidualPSANNNet(
            int(input_dim),
            int(output_dim),
            hidden_layers=self.hidden_layers,
            hidden_units=self.hidden_units,
            hidden_width=self.hidden_width,
            act_kw=self.activation,
            activation_type=self.activation_type,
            w0_first=self.w0_first,
            w0_hidden=self.w0_hidden,
            norm=self.norm,
            drop_path_max=self.drop_path_max,
            residual_alpha_init=self.residual_alpha_init,
        )

class ResConvPSANNRegressor(ResPSANNRegressor):
    """Residual 2D convolutional PSANN regressor with HISSO support."""

    def __init__(
        self,
        *,
        hidden_layers: int = 6,
        hidden_width: Optional[int] = None,
        hidden_units: Optional[int] = None,
        epochs: int = 200,
        batch_size: int = 64,
        lr: float = 1e-3,
        optimizer: str = "adam",
        weight_decay: float = 0.0,
        activation: Optional[ActivationConfig] = None,
        device: str | torch.device = "auto",
        random_state: Optional[int] = None,
        early_stopping: bool = False,
        patience: int = 20,
        num_workers: int = 0,
        loss: LossLike = "mse",
        loss_params: Optional[Dict[str, Any]] = None,
        loss_reduction: str = "mean",
        w0: float = 30.0,
        preserve_shape: bool = True,
        data_format: str = "channels_first",
        conv_kernel_size: int = 3,
        conv_channels: Optional[int] = None,
        per_element: bool = False,
        activation_type: str = "psann",
        stateful: bool = False,
        state: Optional[Union[StateConfig, Mapping[str, Any]]] = None,
        state_reset: str = "batch",
        stream_lr: Optional[float] = None,
        output_shape: Optional[Tuple[int, ...]] = None,
        lsm: Optional[PreprocessorLike] = None,
        lsm_train: bool = False,
        lsm_pretrain_epochs: int = 0,
        lsm_lr: Optional[float] = None,
        warm_start: bool = False,
        scaler: Optional[ScalerSpec] = None,
        scaler_params: Optional[Dict[str, Any]] = None,
        w0_first: float = 12.0,
        w0_hidden: float = 1.0,
        norm: str = "rms",
        drop_path_max: float = 0.0,
        residual_alpha_init: float = 0.0,
    ) -> None:
        if not preserve_shape:
            warnings.warn(
                "ResConvPSANNRegressor forces preserve_shape=True; overriding provided value.",
                UserWarning,
                stacklevel=2,
            )
        super().__init__(
            hidden_layers=hidden_layers,
            hidden_width=hidden_width,
            hidden_units=hidden_units,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            optimizer=optimizer,
            weight_decay=weight_decay,
            activation=activation,
            device=device,
            random_state=random_state,
            early_stopping=early_stopping,
            patience=patience,
            num_workers=num_workers,
            loss=loss,
            loss_params=loss_params,
            loss_reduction=loss_reduction,
            w0=w0,
            preserve_shape=True,
            data_format=data_format,
            conv_kernel_size=conv_kernel_size,
            conv_channels=conv_channels,
            per_element=per_element,
            activation_type=activation_type,
            stateful=stateful,
            state=state,
            state_reset=state_reset,
            stream_lr=stream_lr,
            output_shape=output_shape,
            lsm=lsm,
            lsm_train=lsm_train,
            lsm_pretrain_epochs=lsm_pretrain_epochs,
            lsm_lr=lsm_lr,
            warm_start=warm_start,
            scaler=scaler,
            scaler_params=scaler_params,
            w0_first=w0_first,
            w0_hidden=w0_hidden,
            norm=norm,
            drop_path_max=drop_path_max,
            residual_alpha_init=residual_alpha_init,
        )
        self._use_channel_first_train_inputs_ = True
        self.data_format = data_format

    def _build_conv_core(
        self,
        spatial_ndim: int,
        in_channels: int,
        output_dim: int,
        *,
        segmentation_head: bool,
    ) -> nn.Module:
        if spatial_ndim != 2:
            raise ValueError("ResConvPSANNRegressor currently supports 2D inputs only")
        return ResidualPSANNConv2dNet(
            int(in_channels),
            int(output_dim),
            hidden_layers=self.hidden_layers,
            conv_channels=self.conv_channels,
            hidden_channels=self.conv_channels,
            kernel_size=self.conv_kernel_size,
            act_kw=self.activation,
            activation_type=self.activation_type,
            w0_first=self.w0_first,
            w0_hidden=self.w0_hidden,
            norm=self.norm,
            drop_path_max=self.drop_path_max,
            residual_alpha_init=self.residual_alpha_init,
            segmentation_head=bool(segmentation_head),
        )

    def fit(
        self,
        X: np.ndarray,
        y: np.ndarray | None,
        *,
        validation_data: Optional[ValidationDataLike] = None,
        verbose: int = 0,
        noisy: Optional[NoiseSpec] = None,
        hisso: bool = False,
        hisso_window: Optional[int] = None,
        hisso_reward_fn: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
        hisso_context_extractor: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        hisso_primary_transform: Optional[str] = None,
        hisso_transition_penalty: Optional[float] = None,
        hisso_trans_cost: Optional[float] = None,
        hisso_supervised: Optional[Mapping[str, Any] | bool] = None,
        lr_max: Optional[float] = None,
        lr_min: Optional[float] = None,
    ) -> "ResConvPSANNRegressor":
        if not hisso:
            return super().fit(
                X,
                y,
                validation_data=validation_data,
                verbose=verbose,
                noisy=noisy,
                hisso=False,
                hisso_window=hisso_window,
                hisso_reward_fn=hisso_reward_fn,
                hisso_context_extractor=hisso_context_extractor,
                hisso_primary_transform=hisso_primary_transform,
                hisso_transition_penalty=hisso_transition_penalty,
                hisso_trans_cost=hisso_trans_cost,
                hisso_supervised=hisso_supervised,
                lr_max=lr_max,
                lr_min=lr_min,
            )

        if self.per_element:
            raise ValueError("hisso=True currently supports per_element=False.")

        seed_all(self.random_state)

        fit_args = normalise_fit_args(
            self,
            X,
            y,
            validation_data=validation_data,
            noisy=noisy,
            verbose=verbose,
            lr_max=lr_max,
            lr_min=lr_min,
            hisso=True,
            hisso_kwargs={
                "hisso_window": hisso_window,
                "hisso_reward_fn": hisso_reward_fn,
                "hisso_context_extractor": hisso_context_extractor,
                "hisso_primary_transform": hisso_primary_transform,
                "hisso_transition_penalty": hisso_transition_penalty,
                "hisso_trans_cost": hisso_trans_cost,
                "hisso_supervised": hisso_supervised,
            },
        )

        verbose = fit_args.verbose
        self._keep_column_output_ = bool(fit_args.y is not None and fit_args.y.ndim > 1)

        prepared_state, primary_dim = prepare_inputs_and_scaler(self, fit_args)
        primary_dim = int(primary_dim)
        self._primary_dim_ = primary_dim
        self._output_dim_ = int(prepared_state.output_dim)
        self._train_inputs_layout_ = "cf"
        self._target_cf_shape_ = (
            tuple(prepared_state.y_cf.shape[1:]) if prepared_state.y_cf is not None else self._target_cf_shape_
        )
        self._target_vector_dim_ = (
            int(prepared_state.y_vector.shape[1]) if prepared_state.y_vector is not None else self._target_vector_dim_
        )

        lsm_data = prepared_state.train_inputs
        lsm_model, lsm_channels = self._resolve_lsm_module(lsm_data, preserve_shape=True)

        hooks = self._make_fit_hooks(prepared=prepared_state, verbose=verbose)

        request = ModelBuildRequest(
            estimator=self,
            prepared=prepared_state,
            primary_dim=primary_dim,
            lsm_module=lsm_model,
            lsm_output_dim=lsm_channels,
            preserve_shape=True,
        )

        rebuild = not (self.warm_start and isinstance(getattr(self, "model_", None), nn.Module))
        if rebuild:
            self.model_ = build_model_from_hooks(hooks, request)

        device = self._device()
        self.model_.to(device)

        result = maybe_run_hisso(hooks, request, fit_args=fit_args)
        if result is None:
            raise RuntimeError("HISSO requested but no variant hook was provided.")
        return self


