from __future__ import annotations

from dataclasses import dataclass
import math
from pathlib import Path
from typing import Any, Sequence

import numpy as np
import torch
from torch import nn
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import DataLoader, TensorDataset

from .callbacks import Callback, CallbackList


def _default_device(device: str | torch.device | None = None) -> torch.device:
    if device is None:
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(device)


def _to_numpy(x: Any) -> np.ndarray:
    if isinstance(x, np.ndarray):
        return x
    if isinstance(x, torch.Tensor):
        return x.detach().cpu().numpy()
    return np.asarray(x)


def _as_float_tensor(x: Any) -> torch.Tensor:
    if isinstance(x, torch.Tensor):
        return x.detach().to(dtype=torch.float32, device="cpu")
    return torch.as_tensor(np.asarray(x), dtype=torch.float32)


def _as_lengths_tensor(lengths: Any, *, num_samples: int, max_len: int) -> torch.Tensor:
    if isinstance(lengths, torch.Tensor):
        lengths_tensor = lengths.detach().to(dtype=torch.long, device="cpu")
    else:
        lengths_tensor = torch.as_tensor(np.asarray(lengths), dtype=torch.long)
    if lengths_tensor.ndim != 1:
        lengths_tensor = lengths_tensor.reshape(-1)
    if int(lengths_tensor.numel()) != int(num_samples):
        raise ValueError(f"Expected {num_samples} sequence lengths, got {int(lengths_tensor.numel())}.")
    if torch.any(lengths_tensor < 1):
        raise ValueError("All sequence lengths must be >= 1.")
    if torch.any(lengths_tensor > int(max_len)):
        raise ValueError(f"Sequence lengths must be <= padded sequence length ({max_len}).")
    return lengths_tensor


def _as_sample_weight_tensor(sample_weight: Any, *, num_samples: int) -> torch.Tensor:
    if isinstance(sample_weight, torch.Tensor):
        weight_tensor = sample_weight.detach().to(dtype=torch.float32, device="cpu")
    else:
        weight_tensor = torch.as_tensor(np.asarray(sample_weight), dtype=torch.float32)
    if weight_tensor.ndim != 1:
        weight_tensor = weight_tensor.reshape(-1)
    if int(weight_tensor.numel()) != int(num_samples):
        raise ValueError(f"Expected {num_samples} sample weights, got {int(weight_tensor.numel())}.")
    if torch.any(weight_tensor < 0):
        raise ValueError("Sample weights must be >= 0.")
    return weight_tensor


def _clone_history_dict(history: dict[str, Any] | None) -> dict[str, list[float]]:
    if not isinstance(history, dict):
        return {}
    cloned: dict[str, list[float]] = {}
    for key, value in history.items():
        if isinstance(value, list):
            cloned[str(key)] = [float(v) for v in value]
    return cloned


def _update_confusion_matrix(
    conf_mat: np.ndarray,
    y_true: torch.Tensor,
    y_pred: torch.Tensor,
    *,
    num_classes: int,
) -> None:
    idx = (y_true.to(torch.long) * num_classes + y_pred.to(torch.long)).detach().cpu()
    counts = torch.bincount(idx, minlength=num_classes * num_classes).reshape(num_classes, num_classes)
    conf_mat += counts.numpy()


def _classification_metrics_from_confusion(conf_mat: np.ndarray) -> dict[str, float]:
    conf = conf_mat.astype(np.float64, copy=False)
    total = float(conf.sum())
    correct = float(np.trace(conf))
    tp = np.diag(conf)
    fp = conf.sum(axis=0) - tp
    fn = conf.sum(axis=1) - tp

    precision_per_class = np.divide(tp, tp + fp, out=np.zeros_like(tp), where=(tp + fp) > 0)
    recall_per_class = np.divide(tp, tp + fn, out=np.zeros_like(tp), where=(tp + fn) > 0)
    f1_per_class = np.divide(
        2.0 * precision_per_class * recall_per_class,
        precision_per_class + recall_per_class,
        out=np.zeros_like(precision_per_class),
        where=(precision_per_class + recall_per_class) > 0,
    )
    return {
        "accuracy": correct / total if total > 0 else 0.0,
        "precision": float(np.mean(precision_per_class)) if precision_per_class.size else 0.0,
        "recall": float(np.mean(recall_per_class)) if recall_per_class.size else 0.0,
        "f1": float(np.mean(f1_per_class)) if f1_per_class.size else 0.0,
    }


def _r2_from_sums(*, sq_error_sum: float, y_sum: float, y_sq_sum: float, count: int) -> float:
    if count <= 0:
        return 0.0
    ss_tot = y_sq_sum - (y_sum * y_sum) / float(count)
    if ss_tot <= 0.0:
        return 0.0 if sq_error_sum > 0.0 else 1.0
    return 1.0 - (sq_error_sum / ss_tot)


@dataclass
class TrainingConfig:
    epochs: int = 10
    batch_size: int = 32
    lr: float = 1e-3
    weight_decay: float = 0.0
    shuffle: bool = True
    verbose: bool = True


class _BaseClassifier:
    """Minimal sklearn-like wrapper around a PyTorch classifier."""

    def __init__(
        self,
        *,
        device: str | torch.device | None = None,
        epochs: int = 10,
        batch_size: int = 32,
        lr: float = 1e-3,
        weight_decay: float = 0.0,
        random_state: int | None = None,
    ) -> None:
        self.device = _default_device(device)
        self.training = TrainingConfig(
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            weight_decay=weight_decay,
        )
        self.random_state = random_state
        self.model_: nn.Module | None = None
        self.classes_: np.ndarray | None = None
        self.history_: dict[str, list[float]] = {}
        self._built_input_shape: tuple[int, ...] | None = None
        self._criterion = nn.CrossEntropyLoss()
        self.best_epoch_: int | None = None
        self.best_score_: float | None = None
        self.monitor_: str | None = None
        self.stopped_epoch_: int | None = None
        self.early_stopped_: bool = False
        if random_state is not None:
            np.random.seed(random_state)
            torch.manual_seed(random_state)

    def _prepare_X(self, X: Any) -> torch.Tensor:
        raise NotImplementedError

    def _build_model(self, input_shape: tuple[int, ...], num_classes: int) -> nn.Module:
        raise NotImplementedError

    def get_config(self) -> dict[str, Any]:
        return {}

    @staticmethod
    def _monitor_mode(monitor: str) -> str:
        if monitor.endswith(("loss", "mse", "mae", "rmse")):
            return "min"
        if monitor.endswith(("accuracy", "precision", "recall", "f1", "r2")):
            return "max"
        raise ValueError(f"Unsupported monitor: {monitor!r}")

    @staticmethod
    def _is_improvement(current: float, best: float | None, *, mode: str, min_delta: float) -> bool:
        if best is None:
            return True
        if mode == "min":
            return current < (best - min_delta)
        if mode == "max":
            return current > (best + min_delta)
        raise ValueError(f"Unsupported monitor mode: {mode!r}")

    @staticmethod
    def _clone_state_dict_cpu(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        return {k: v.detach().cpu().clone() for k, v in state_dict.items()}

    def _serialization_payload(self, *, extra: dict[str, Any] | None = None) -> dict[str, Any]:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        payload = {
            "state_dict": self.model_.state_dict(),
            "classes": None if self.classes_ is None else self.classes_.tolist(),
            "config": self.get_config(),
            "model_type": self.__class__.__name__,
            "input_shape": self._built_input_shape,
        }
        if extra:
            payload.update(extra)
        return payload

    @staticmethod
    def _load_checkpoint_payload(path: str | Path, *, device: str | torch.device | None = None) -> dict[str, Any]:
        ckpt_path = Path(path)
        map_location = _default_device(device)
        try:
            payload = torch.load(str(ckpt_path), map_location=map_location, weights_only=False)
        except TypeError:
            payload = torch.load(str(ckpt_path), map_location=map_location)
        if not isinstance(payload, dict):
            raise ValueError(f"Invalid checkpoint format in {ckpt_path!s}: expected dict payload.")
        return payload

    def _validate_checkpoint_payload_for_self(self, payload: dict[str, Any]) -> None:
        saved_model_type = payload.get("model_type")
        if saved_model_type is not None and saved_model_type != self.__class__.__name__:
            raise ValueError(
                f"Checkpoint model type {saved_model_type!r} does not match model class {self.__class__.__name__!r}."
            )
        input_shape = payload.get("input_shape")
        if input_shape is not None and self._built_input_shape is not None:
            saved_input_shape = tuple(int(v) for v in input_shape)
            if tuple(saved_input_shape) != tuple(self._built_input_shape):
                raise ValueError(
                    f"Checkpoint input shape {saved_input_shape!r} does not match current model input shape "
                    f"{self._built_input_shape!r}."
                )

    def _encode_y_fit(self, y: Any) -> torch.Tensor:
        y_np = _to_numpy(y)
        if y_np.ndim != 1:
            y_np = y_np.reshape(-1)
        classes, encoded = np.unique(y_np, return_inverse=True)
        self.classes_ = classes
        return torch.as_tensor(encoded, dtype=torch.long)

    def _encode_y_existing(self, y: Any) -> torch.Tensor:
        if self.classes_ is None:
            raise RuntimeError("Model has not been fitted.")
        y_np = _to_numpy(y)
        if y_np.ndim != 1:
            y_np = y_np.reshape(-1)
        class_to_index = {cls.item() if hasattr(cls, "item") else cls: idx for idx, cls in enumerate(self.classes_)}
        encoded = []
        for label in y_np.tolist():
            key = label.item() if hasattr(label, "item") else label
            if key not in class_to_index:
                raise ValueError(f"Unknown label during evaluation: {key!r}")
            encoded.append(class_to_index[key])
        return torch.as_tensor(encoded, dtype=torch.long)

    def _ensure_built(self, X_tensor: torch.Tensor, y_tensor: torch.Tensor) -> None:
        input_shape = tuple(X_tensor.shape[1:])
        num_classes = int(torch.unique(y_tensor).numel())
        if num_classes < 2:
            raise ValueError("At least 2 classes are required for classification.")
        if self.model_ is None:
            self.model_ = self._build_model(input_shape, num_classes).to(self.device)
            self._built_input_shape = input_shape

    def _prepare_lengths(self, lengths: Any | None, X_tensor: torch.Tensor) -> torch.Tensor | None:
        if lengths is None:
            return None
        raise ValueError(f"{self.__class__.__name__} does not support sequence lengths.")

    def _prepare_sample_weight(self, sample_weight: Any | None, X_tensor: torch.Tensor) -> torch.Tensor | None:
        if sample_weight is None:
            return None
        return _as_sample_weight_tensor(sample_weight, num_samples=int(X_tensor.shape[0]))

    def _prepare_class_weight(self, class_weight: Any | None) -> torch.Tensor | None:
        if class_weight is None:
            return None
        if self.classes_ is None:
            raise RuntimeError("Model classes are not initialized.")

        num_classes = int(len(self.classes_))
        if isinstance(class_weight, dict):
            weights = np.ones(num_classes, dtype=np.float32)
            class_to_index = {cls.item() if hasattr(cls, "item") else cls: idx for idx, cls in enumerate(self.classes_)}
            for label, value in class_weight.items():
                key = label.item() if hasattr(label, "item") else label
                if key not in class_to_index:
                    raise ValueError(f"Unknown class in class_weight: {label!r}")
                weights[class_to_index[key]] = float(value)
            arr = weights
        else:
            arr = np.asarray(class_weight, dtype=np.float32).reshape(-1)
            if int(arr.size) != num_classes:
                raise ValueError(f"class_weight must have length {num_classes}, got {int(arr.size)}.")
        if np.any(arr < 0):
            raise ValueError("class_weight values must be >= 0.")
        return torch.as_tensor(arr, dtype=torch.float32)

    def _forward_model(self, xb: torch.Tensor, lengths: torch.Tensor | None = None) -> torch.Tensor:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        if lengths is not None:
            raise ValueError(f"{self.__class__.__name__} does not support sequence lengths.")
        return self.model_(xb)

    @staticmethod
    def _unpack_batch(
        batch: Sequence[torch.Tensor],
        *,
        has_y: bool,
        has_lengths: bool,
        has_sample_weight: bool,
    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
        idx = 0
        xb = batch[idx]
        idx += 1
        yb = None
        lengths_b = None
        sample_weight_b = None
        if has_y:
            if idx >= len(batch):
                raise RuntimeError("Unexpected batch structure: missing target tensor.")
            yb = batch[idx]
            idx += 1
        if has_lengths:
            if idx >= len(batch):
                raise RuntimeError("Unexpected batch structure: missing lengths tensor.")
            lengths_b = batch[idx]
            idx += 1
        if has_sample_weight:
            if idx >= len(batch):
                raise RuntimeError("Unexpected batch structure: missing sample-weight tensor.")
            sample_weight_b = batch[idx]
            idx += 1
        if idx != len(batch):
            raise RuntimeError("Unexpected batch structure: extra tensors encountered.")
        return xb, yb, lengths_b, sample_weight_b

    def _create_scheduler(
        self,
        optimizer: torch.optim.Optimizer,
        scheduler: Any | None,
        *,
        scheduler_kwargs: dict[str, Any] | None = None,
        epochs: int | None = None,
        monitor_mode: str | None = None,
    ) -> Any | None:
        if scheduler is None:
            return None
        kwargs = dict(scheduler_kwargs or {})
        sched_obj: Any
        if isinstance(scheduler, dict):
            cfg = dict(scheduler)
            name = str(cfg.pop("name", cfg.pop("type", ""))).lower()
            kwargs = {**cfg, **kwargs}
            sched_obj = self._create_scheduler(
                optimizer,
                name,
                scheduler_kwargs=kwargs,
                epochs=epochs,
                monitor_mode=monitor_mode,
            )
            return sched_obj
        if isinstance(scheduler, str):
            name = scheduler.lower()
            if name in {"step", "steplr"}:
                kwargs.setdefault("step_size", 10)
                kwargs.setdefault("gamma", 0.1)
                return torch.optim.lr_scheduler.StepLR(optimizer, **kwargs)
            if name in {"exp", "exponential", "exponentiallr"}:
                kwargs.setdefault("gamma", 0.99)
                return torch.optim.lr_scheduler.ExponentialLR(optimizer, **kwargs)
            if name in {"cosine", "cosineannealing", "cosineannealinglr"}:
                kwargs.setdefault("T_max", int(epochs or 10))
                return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **kwargs)
            if name in {"plateau", "reduceonplateau", "reducelronplateau"}:
                if monitor_mode is not None:
                    kwargs.setdefault("mode", monitor_mode)
                return ReduceLROnPlateau(optimizer, **kwargs)
            raise ValueError(
                f"Unsupported scheduler {scheduler!r}. "
                "Use one of {'step','exponential','cosine','plateau'}, a dict, a scheduler instance, or a callable."
            )
        if callable(scheduler) and not hasattr(scheduler, "step"):
            sched_obj = scheduler(optimizer)
            if not hasattr(sched_obj, "step"):
                raise ValueError("Scheduler factory must return an object with a .step(...) method.")
            return sched_obj
        if hasattr(scheduler, "step"):
            return scheduler
        raise ValueError("Invalid scheduler argument.")

    @staticmethod
    def _scheduler_is_plateau(scheduler: Any | None) -> bool:
        return isinstance(scheduler, ReduceLROnPlateau) or (
            scheduler is not None and scheduler.__class__.__name__ == "ReduceLROnPlateau"
        )

    def _step_scheduler(self, scheduler: Any | None, *, metric: float | None = None) -> None:
        if scheduler is None:
            return
        if self._scheduler_is_plateau(scheduler):
            if metric is None:
                raise ValueError("ReduceLROnPlateau requires a monitor metric.")
            scheduler.step(metric)
            return
        scheduler.step()

    @staticmethod
    def _current_lr(optimizer: torch.optim.Optimizer) -> float:
        if not optimizer.param_groups:
            return 0.0
        return float(optimizer.param_groups[0].get("lr", 0.0))

    def _make_loader(
        self,
        X: torch.Tensor,
        y: torch.Tensor | None = None,
        *,
        shuffle: bool = False,
        batch_size: int | None = None,
        extra_tensors: Sequence[torch.Tensor] | None = None,
    ) -> DataLoader:
        tensors: list[torch.Tensor] = [X]
        if y is not None:
            tensors.append(y)
        if extra_tensors:
            tensors.extend(extra_tensors)
        dataset = TensorDataset(*tensors)
        return DataLoader(dataset, batch_size=batch_size or self.training.batch_size, shuffle=shuffle)

    def fit(
        self,
        X: Any,
        y: Any,
        *,
        epochs: int | None = None,
        batch_size: int | None = None,
        lr: float | None = None,
        weight_decay: float | None = None,
        lengths: Any | None = None,
        val_data: tuple[Any, ...] | None = None,
        early_stopping: bool = False,
        patience: int = 5,
        min_delta: float = 0.0,
        restore_best_weights: bool | None = None,
        monitor: str | None = None,
        checkpoint_path: str | Path | None = None,
        resume_from_checkpoint: str | Path | None = None,
        class_weight: Any | None = None,
        sample_weight: Any | None = None,
        scheduler: Any | None = None,
        scheduler_kwargs: dict[str, Any] | None = None,
        scheduler_monitor: str | None = None,
        callbacks: Sequence[Callback] | None = None,
        verbose: bool | None = None,
    ) -> "_BaseClassifier":
        X_tensor = self._prepare_X(X)
        y_tensor = self._encode_y_fit(y)
        lengths_tensor = self._prepare_lengths(lengths, X_tensor)
        sample_weight_tensor = self._prepare_sample_weight(sample_weight, X_tensor)
        if X_tensor.shape[0] != y_tensor.shape[0]:
            raise ValueError("X and y must have the same number of samples.")
        if patience < 1:
            raise ValueError("patience must be >= 1.")
        if min_delta < 0:
            raise ValueError("min_delta must be >= 0.")

        self._ensure_built(X_tensor, y_tensor)
        assert self.model_ is not None

        cfg = TrainingConfig(
            epochs=epochs if epochs is not None else self.training.epochs,
            batch_size=batch_size if batch_size is not None else self.training.batch_size,
            lr=lr if lr is not None else self.training.lr,
            weight_decay=weight_decay if weight_decay is not None else self.training.weight_decay,
            shuffle=True,
            verbose=self.training.verbose if verbose is None else verbose,
        )
        self.training = cfg

        train_loader = self._make_loader(
            X_tensor,
            y_tensor,
            shuffle=cfg.shuffle,
            extra_tensors=[
                t for t in (lengths_tensor, sample_weight_tensor) if t is not None
            ] or None,
        )
        optimizer = torch.optim.Adam(self.model_.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)

        history: dict[str, list[float]] = {
            "loss": [],
            "accuracy": [],
            "precision": [],
            "recall": [],
            "f1": [],
        }
        val_X = None
        val_y = None
        val_lengths = None
        if val_data is not None:
            if len(val_data) == 2:
                val_X, val_y = val_data
            elif len(val_data) == 3:
                val_X, val_y, val_lengths = val_data
            else:
                raise ValueError("val_data must be (X_val, y_val) or (X_val, y_val, lengths_val).")
            history["val_loss"] = []
            history["val_accuracy"] = []
            history["val_precision"] = []
            history["val_recall"] = []
            history["val_f1"] = []

        resolved_monitor = monitor or ("val_loss" if val_data is not None else "loss")
        valid_monitors = {"loss", "accuracy", "precision", "recall", "f1"}
        if val_data is not None:
            valid_monitors.update({"val_loss", "val_accuracy", "val_precision", "val_recall", "val_f1"})
        if resolved_monitor not in valid_monitors:
            raise ValueError(
                f"Unsupported monitor {resolved_monitor!r}. "
                f"Expected one of {sorted(valid_monitors)!r}."
            )
        monitor_mode = self._monitor_mode(resolved_monitor)
        class_weight_tensor = self._prepare_class_weight(class_weight)
        scheduler_metric_name = scheduler_monitor or resolved_monitor
        if scheduler_metric_name not in valid_monitors:
            raise ValueError(
                f"Unsupported scheduler_monitor {scheduler_metric_name!r}. "
                f"Expected one of {sorted(valid_monitors)!r}."
            )
        scheduler_monitor_mode = self._monitor_mode(scheduler_metric_name)
        train_loss_fn = nn.CrossEntropyLoss(
            weight=class_weight_tensor.to(self.device) if class_weight_tensor is not None else None,
            reduction="none",
        )
        scheduler_obj = self._create_scheduler(
            optimizer,
            scheduler,
            scheduler_kwargs=scheduler_kwargs,
            epochs=cfg.epochs,
            monitor_mode=scheduler_monitor_mode,
        )
        ckpt_path = Path(checkpoint_path) if checkpoint_path is not None else None
        if ckpt_path is not None:
            ckpt_path.parent.mkdir(parents=True, exist_ok=True)
        effective_restore_best = early_stopping if restore_best_weights is None else restore_best_weights

        if resume_from_checkpoint is not None:
            resume_payload = self._load_checkpoint_payload(resume_from_checkpoint, device=self.device)
            self._validate_checkpoint_payload_for_self(resume_payload)

            saved_classes = resume_payload.get("classes")
            if saved_classes is not None and self.classes_ is not None:
                if not np.array_equal(np.asarray(saved_classes), self.classes_):
                    raise ValueError("Checkpoint classes do not match labels inferred from current y.")

            state_dict = resume_payload.get("state_dict")
            if state_dict is None:
                raise ValueError("Checkpoint is missing 'state_dict'.")
            self.model_.load_state_dict(state_dict)

            opt_state = resume_payload.get("optimizer_state_dict")
            if isinstance(opt_state, dict):
                optimizer.load_state_dict(opt_state)
            if scheduler_obj is not None:
                sched_state = resume_payload.get("scheduler_state_dict")
                if isinstance(sched_state, dict):
                    try:
                        scheduler_obj.load_state_dict(sched_state)
                    except Exception:
                        pass

            saved_history = _clone_history_dict(resume_payload.get("history"))
            if saved_history:
                for key in history.keys():
                    if key in saved_history:
                        history[key] = saved_history[key]

        best_metric: float | None = None
        best_epoch: int | None = None
        best_state_dict: dict[str, torch.Tensor] | None = None
        wait = 0
        start_epoch = 0
        if resume_from_checkpoint is not None:
            checkpoint_meta = resume_payload.get("checkpoint")
            if isinstance(checkpoint_meta, dict) and checkpoint_meta.get("monitor") == resolved_monitor:
                if "best_score" in checkpoint_meta:
                    best_metric = float(checkpoint_meta["best_score"])
                if "best_epoch" in checkpoint_meta:
                    best_epoch = int(checkpoint_meta["best_epoch"])
            if isinstance(checkpoint_meta, dict) and "epoch" in checkpoint_meta:
                start_epoch = int(checkpoint_meta["epoch"])
            elif history.get("loss"):
                start_epoch = len(history["loss"])
        self.monitor_ = resolved_monitor
        self.best_epoch_ = None
        self.best_score_ = None
        self.stopped_epoch_ = None
        self.early_stopped_ = False
        cb_list = CallbackList(callbacks)
        cb_list.set_model(self)
        cb_list.set_context(optimizer=optimizer, scheduler=scheduler_obj, training_config=cfg)
        cb_list.on_train_begin(logs={"history": history, "start_epoch": start_epoch, "lr": self._current_lr(optimizer)})

        total_epochs_target = start_epoch + cfg.epochs
        for epoch_idx in range(cfg.epochs):
            epoch_num = start_epoch + epoch_idx + 1
            cb_list.on_epoch_begin(epoch_num, logs={"lr": self._current_lr(optimizer)})
            self.model_.train()
            total_loss = 0.0
            total_samples = 0
            total_loss_denom = 0.0
            num_classes = int(len(self.classes_)) if self.classes_ is not None else 0
            conf_mat = np.zeros((num_classes, num_classes), dtype=np.int64)
            has_lengths = lengths_tensor is not None
            has_sample_weight = sample_weight_tensor is not None
            class_weights_device = class_weight_tensor.to(self.device) if class_weight_tensor is not None else None

            for batch in train_loader:
                xb, yb, lengths_b, sample_weight_b = self._unpack_batch(
                    batch,
                    has_y=True,
                    has_lengths=has_lengths,
                    has_sample_weight=has_sample_weight,
                )
                assert yb is not None
                xb = xb.to(self.device)
                yb = yb.to(self.device)
                optimizer.zero_grad(set_to_none=True)
                logits = self._forward_model(xb, lengths_b)
                per_sample_loss = train_loss_fn(logits, yb)
                denom = torch.ones_like(per_sample_loss)
                if class_weights_device is not None:
                    denom = denom * class_weights_device[yb]
                if sample_weight_b is not None:
                    sw = sample_weight_b.to(self.device)
                    per_sample_loss = per_sample_loss * sw
                    denom = denom * sw
                denom_sum = denom.sum().clamp_min(1e-12)
                loss = per_sample_loss.sum() / denom_sum
                loss.backward()
                optimizer.step()

                total_loss += float(per_sample_loss.sum().item())
                total_loss_denom += float(denom_sum.item())
                preds = logits.argmax(dim=1)
                _update_confusion_matrix(conf_mat, yb, preds, num_classes=num_classes)
                total_samples += xb.size(0)

            epoch_loss = total_loss / max(total_loss_denom, 1e-12)
            cls_metrics = _classification_metrics_from_confusion(conf_mat)
            epoch_acc = cls_metrics["accuracy"]
            history["loss"].append(epoch_loss)
            history["accuracy"].append(epoch_acc)
            history["precision"].append(cls_metrics["precision"])
            history["recall"].append(cls_metrics["recall"])
            history["f1"].append(cls_metrics["f1"])

            if val_X is not None and val_y is not None:
                metrics = self.evaluate(
                    val_X,
                    val_y,
                    lengths=val_lengths,
                    batch_size=cfg.batch_size,
                    _return_dict=True,
                )
                history["val_loss"].append(metrics["loss"])
                history["val_accuracy"].append(metrics["accuracy"])
                history["val_precision"].append(metrics["precision"])
                history["val_recall"].append(metrics["recall"])
                history["val_f1"].append(metrics["f1"])

            if cfg.verbose:
                if val_data is None:
                    print(
                        f"Epoch {epoch_num}/{total_epochs_target} - "
                        f"loss: {epoch_loss:.4f} - acc: {epoch_acc:.4f}"
                    )
                else:
                    print(
                        f"Epoch {epoch_num}/{total_epochs_target} - "
                        f"loss: {epoch_loss:.4f} - acc: {epoch_acc:.4f} - "
                        f"val_loss: {history['val_loss'][-1]:.4f} - val_acc: {history['val_accuracy'][-1]:.4f}"
                    )

            if scheduler_obj is not None:
                sched_metric = float(history[scheduler_metric_name][-1]) if scheduler_metric_name in history else None
                self._step_scheduler(scheduler_obj, metric=sched_metric)

            epoch_logs: dict[str, Any] = {k: float(v[-1]) for k, v in history.items() if isinstance(v, list) and v}
            epoch_logs["lr"] = self._current_lr(optimizer)
            cb_list.on_epoch_end(epoch_num, logs=epoch_logs)

            current_metric = float(history[resolved_monitor][-1])
            improved = self._is_improvement(
                current_metric,
                best_metric,
                mode=monitor_mode,
                min_delta=min_delta,
            )
            if improved:
                best_metric = current_metric
                best_epoch = epoch_num
                wait = 0
                if effective_restore_best or early_stopping or ckpt_path is not None:
                    best_state_dict = self._clone_state_dict_cpu(self.model_.state_dict())
                if ckpt_path is not None:
                    checkpoint_meta = {
                        "checkpoint": {
                            "epoch": epoch_num,
                            "best_epoch": best_epoch,
                            "best_score": best_metric,
                            "monitor": resolved_monitor,
                            "monitor_mode": monitor_mode,
                        },
                        "history": _clone_history_dict(history),
                        "optimizer_state_dict": optimizer.state_dict(),
                    }
                    if scheduler_obj is not None and hasattr(scheduler_obj, "state_dict"):
                        checkpoint_meta["scheduler_state_dict"] = scheduler_obj.state_dict()
                    torch.save(self._serialization_payload(extra=checkpoint_meta), str(ckpt_path))
            else:
                wait += 1

            if cb_list.stop_training:
                self.stopped_epoch_ = epoch_num
                self.early_stopped_ = True
                if cfg.verbose:
                    print(f"Training stopped by callback at epoch {epoch_num}.")
                break

            if early_stopping and wait >= patience:
                self.stopped_epoch_ = epoch_num
                self.early_stopped_ = True
                if cfg.verbose:
                    print(
                        f"Early stopping at epoch {epoch_num} "
                        f"(best {resolved_monitor}={best_metric:.4f} at epoch {best_epoch})."
                    )
                break

        if effective_restore_best and best_state_dict is not None:
            self.model_.load_state_dict(best_state_dict)
        self.best_epoch_ = best_epoch
        self.best_score_ = best_metric
        self.model_.eval()

        self.history_ = history
        cb_list.on_train_end(logs={"history": history, "best_epoch": self.best_epoch_, "best_score": self.best_score_})
        return self

    @torch.no_grad()
    def predict_proba(
        self,
        X: Any,
        *,
        lengths: Any | None = None,
        batch_size: int | None = None,
    ) -> np.ndarray:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        X_tensor = self._prepare_X(X)
        lengths_tensor = self._prepare_lengths(lengths, X_tensor)
        loader = self._make_loader(
            X_tensor,
            shuffle=False,
            batch_size=batch_size,
            extra_tensors=[lengths_tensor] if lengths_tensor is not None else None,
        )
        self.model_.eval()
        probs = []
        for batch in loader:
            if len(batch) == 1:
                (xb,) = batch
                lengths_b = None
            elif len(batch) == 2:
                xb, lengths_b = batch
            else:
                raise RuntimeError("Unexpected batch structure in classifier predict_proba loop.")
            xb = xb.to(self.device)
            logits = self._forward_model(xb, lengths_b)
            probs.append(torch.softmax(logits, dim=1).cpu())
        return torch.cat(probs, dim=0).numpy()

    def predict(
        self,
        X: Any,
        *,
        lengths: Any | None = None,
        batch_size: int | None = None,
    ) -> np.ndarray:
        probs = self.predict_proba(X, lengths=lengths, batch_size=batch_size)
        idx = probs.argmax(axis=1)
        if self.classes_ is None:
            raise RuntimeError("Model has not been fitted.")
        return self.classes_[idx]

    @torch.no_grad()
    def evaluate(
        self,
        X: Any,
        y: Any,
        *,
        lengths: Any | None = None,
        batch_size: int | None = None,
        _return_dict: bool = False,
    ) -> dict[str, float]:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        X_tensor = self._prepare_X(X)
        y_tensor = self._encode_y_existing(y)
        lengths_tensor = self._prepare_lengths(lengths, X_tensor)
        if X_tensor.shape[0] != y_tensor.shape[0]:
            raise ValueError("X and y must have the same number of samples.")
        loader = self._make_loader(
            X_tensor,
            y_tensor,
            shuffle=False,
            batch_size=batch_size,
            extra_tensors=[lengths_tensor] if lengths_tensor is not None else None,
        )
        self.model_.eval()

        total_loss = 0.0
        total_samples = 0
        num_classes = int(len(self.classes_)) if self.classes_ is not None else 0
        conf_mat = np.zeros((num_classes, num_classes), dtype=np.int64)
        for batch in loader:
            if len(batch) == 2:
                xb, yb = batch
                lengths_b = None
            elif len(batch) == 3:
                xb, yb, lengths_b = batch
            else:
                raise RuntimeError("Unexpected batch structure in classifier evaluation loop.")
            xb = xb.to(self.device)
            yb = yb.to(self.device)
            logits = self._forward_model(xb, lengths_b)
            loss = self._criterion(logits, yb)
            total_loss += float(loss.item()) * xb.size(0)
            preds = logits.argmax(dim=1)
            _update_confusion_matrix(conf_mat, yb, preds, num_classes=num_classes)
            total_samples += xb.size(0)

        cls_metrics = _classification_metrics_from_confusion(conf_mat)
        result = {
            "loss": total_loss / max(total_samples, 1),
            "accuracy": cls_metrics["accuracy"],
            "precision": cls_metrics["precision"],
            "recall": cls_metrics["recall"],
            "f1": cls_metrics["f1"],
        }
        if _return_dict:
            return result
        return result

    def save(self, path: str | Path) -> None:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(self._serialization_payload(), str(path))

    @classmethod
    def load(
        cls,
        path: str | Path,
        *,
        device: str | torch.device | None = None,
    ) -> "_BaseClassifier":
        """Load a saved classifier checkpoint produced by ``save()``."""
        path = Path(path)
        map_location = _default_device(device)

        try:
            payload = torch.load(str(path), map_location=map_location, weights_only=False)
        except TypeError:
            # Compatibility with older torch versions that do not support `weights_only`.
            payload = torch.load(str(path), map_location=map_location)

        if not isinstance(payload, dict):
            raise ValueError(f"Invalid checkpoint format in {path!s}: expected dict payload.")

        saved_model_type = payload.get("model_type")
        if saved_model_type is not None and saved_model_type != cls.__name__:
            raise ValueError(
                f"Checkpoint model type {saved_model_type!r} does not match loader class {cls.__name__!r}."
            )

        input_shape = payload.get("input_shape")
        if input_shape is None:
            raise ValueError("Checkpoint is missing 'input_shape'.")
        input_shape = tuple(int(v) for v in input_shape)

        classes = payload.get("classes")
        if classes is None:
            raise ValueError("Checkpoint is missing 'classes'.")
        classes_arr = np.asarray(classes)
        num_classes = int(len(classes_arr))
        if num_classes < 2:
            raise ValueError("Checkpoint must contain at least 2 classes.")

        config = payload.get("config") or {}
        if not isinstance(config, dict):
            raise ValueError("Checkpoint 'config' must be a dict.")

        init_kwargs: dict[str, Any] = {}
        training_cfg = config.get("training", {})
        if isinstance(training_cfg, dict):
            for key in ("epochs", "batch_size", "lr", "weight_decay"):
                if key in training_cfg:
                    init_kwargs[key] = training_cfg[key]
        for key, value in config.items():
            if key != "training":
                init_kwargs[key] = value
        if device is not None:
            init_kwargs["device"] = device

        obj = cls(**init_kwargs)
        obj.classes_ = classes_arr
        obj._built_input_shape = input_shape
        obj.model_ = obj._build_model(input_shape, num_classes).to(obj.device)

        state_dict = payload.get("state_dict")
        if state_dict is None:
            raise ValueError("Checkpoint is missing 'state_dict'.")
        obj.model_.load_state_dict(state_dict)
        obj.model_.eval()
        return obj


class _BaseRegressor(_BaseClassifier):
    """Minimal sklearn-like wrapper around a PyTorch regressor."""

    def __init__(
        self,
        *,
        device: str | torch.device | None = None,
        epochs: int = 10,
        batch_size: int = 32,
        lr: float = 1e-3,
        weight_decay: float = 0.0,
        random_state: int | None = None,
    ) -> None:
        super().__init__(
            device=device,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            weight_decay=weight_decay,
            random_state=random_state,
        )
        self._criterion = nn.MSELoss()
        self.target_dim_: int | None = None
        self._squeeze_output_: bool = True
        # Class labels are not used for regression.
        self.classes_ = None

    def _serialization_payload(self, *, extra: dict[str, Any] | None = None) -> dict[str, Any]:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        payload = {
            "state_dict": self.model_.state_dict(),
            "config": self.get_config(),
            "model_type": self.__class__.__name__,
            "input_shape": self._built_input_shape,
            "target_dim": self.target_dim_,
            "squeeze_output": self._squeeze_output_,
        }
        if extra:
            payload.update(extra)
        return payload

    def _prepare_y_fit(self, y: Any) -> torch.Tensor:
        y_tensor = _as_float_tensor(y)
        original_ndim = y_tensor.ndim
        if original_ndim == 1:
            y_tensor = y_tensor.unsqueeze(1)
        if y_tensor.ndim != 2:
            raise ValueError("Regression targets must have shape (N,) or (N, D).")
        if y_tensor.shape[1] < 1:
            raise ValueError("Regression targets must have at least one output dimension.")
        self.target_dim_ = int(y_tensor.shape[1])
        self._squeeze_output_ = original_ndim == 1
        return y_tensor

    def _prepare_y_existing(self, y: Any) -> torch.Tensor:
        if self.target_dim_ is None:
            raise RuntimeError("Model has not been fitted.")
        y_tensor = _as_float_tensor(y)
        if y_tensor.ndim == 1:
            if self.target_dim_ != 1:
                raise ValueError(f"Expected multi-output targets with shape (N, {self.target_dim_}).")
            y_tensor = y_tensor.unsqueeze(1)
        if y_tensor.ndim != 2:
            raise ValueError("Regression targets must have shape (N,) or (N, D).")
        if int(y_tensor.shape[1]) != self.target_dim_:
            raise ValueError(f"Expected target dimension {self.target_dim_}, got {int(y_tensor.shape[1])}.")
        return y_tensor

    def _ensure_built(self, X_tensor: torch.Tensor, y_tensor: torch.Tensor) -> None:
        input_shape = tuple(X_tensor.shape[1:])
        if y_tensor.ndim != 2:
            raise ValueError("Regression targets must be a 2D tensor after preprocessing.")
        output_dim = int(y_tensor.shape[1])
        if output_dim < 1:
            raise ValueError("Regression output dimension must be >= 1.")
        if self.model_ is None:
            self.model_ = self._build_model(input_shape, output_dim).to(self.device)
            self._built_input_shape = input_shape

    def fit(
        self,
        X: Any,
        y: Any,
        *,
        epochs: int | None = None,
        batch_size: int | None = None,
        lr: float | None = None,
        weight_decay: float | None = None,
        lengths: Any | None = None,
        val_data: tuple[Any, ...] | None = None,
        early_stopping: bool = False,
        patience: int = 5,
        min_delta: float = 0.0,
        restore_best_weights: bool | None = None,
        monitor: str | None = None,
        checkpoint_path: str | Path | None = None,
        resume_from_checkpoint: str | Path | None = None,
        sample_weight: Any | None = None,
        scheduler: Any | None = None,
        scheduler_kwargs: dict[str, Any] | None = None,
        scheduler_monitor: str | None = None,
        callbacks: Sequence[Callback] | None = None,
        verbose: bool | None = None,
    ) -> "_BaseRegressor":
        X_tensor = self._prepare_X(X)
        y_tensor = self._prepare_y_fit(y)
        lengths_tensor = self._prepare_lengths(lengths, X_tensor)
        sample_weight_tensor = self._prepare_sample_weight(sample_weight, X_tensor)
        if X_tensor.shape[0] != y_tensor.shape[0]:
            raise ValueError("X and y must have the same number of samples.")
        if patience < 1:
            raise ValueError("patience must be >= 1.")
        if min_delta < 0:
            raise ValueError("min_delta must be >= 0.")

        self._ensure_built(X_tensor, y_tensor)
        assert self.model_ is not None

        cfg = TrainingConfig(
            epochs=epochs if epochs is not None else self.training.epochs,
            batch_size=batch_size if batch_size is not None else self.training.batch_size,
            lr=lr if lr is not None else self.training.lr,
            weight_decay=weight_decay if weight_decay is not None else self.training.weight_decay,
            shuffle=True,
            verbose=self.training.verbose if verbose is None else verbose,
        )
        self.training = cfg

        train_loader = self._make_loader(
            X_tensor,
            y_tensor,
            shuffle=cfg.shuffle,
            extra_tensors=[
                t for t in (lengths_tensor, sample_weight_tensor) if t is not None
            ] or None,
        )
        optimizer = torch.optim.Adam(self.model_.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)

        history: dict[str, list[float]] = {"loss": [], "mse": [], "mae": [], "rmse": [], "r2": []}
        val_X = None
        val_y = None
        val_lengths = None
        if val_data is not None:
            if len(val_data) == 2:
                val_X, val_y = val_data
            elif len(val_data) == 3:
                val_X, val_y, val_lengths = val_data
            else:
                raise ValueError("val_data must be (X_val, y_val) or (X_val, y_val, lengths_val).")
            history["val_loss"] = []
            history["val_mse"] = []
            history["val_mae"] = []
            history["val_rmse"] = []
            history["val_r2"] = []

        resolved_monitor = monitor or ("val_loss" if val_data is not None else "loss")
        valid_monitors = {"loss", "mse", "mae", "rmse", "r2"}
        if val_data is not None:
            valid_monitors.update({"val_loss", "val_mse", "val_mae", "val_rmse", "val_r2"})
        if resolved_monitor not in valid_monitors:
            raise ValueError(
                f"Unsupported monitor {resolved_monitor!r}. "
                f"Expected one of {sorted(valid_monitors)!r}."
            )
        monitor_mode = self._monitor_mode(resolved_monitor)
        scheduler_metric_name = scheduler_monitor or resolved_monitor
        if scheduler_metric_name not in valid_monitors:
            raise ValueError(
                f"Unsupported scheduler_monitor {scheduler_metric_name!r}. "
                f"Expected one of {sorted(valid_monitors)!r}."
            )
        scheduler_monitor_mode = self._monitor_mode(scheduler_metric_name)
        scheduler_obj = self._create_scheduler(
            optimizer,
            scheduler,
            scheduler_kwargs=scheduler_kwargs,
            epochs=cfg.epochs,
            monitor_mode=scheduler_monitor_mode,
        )
        ckpt_path = Path(checkpoint_path) if checkpoint_path is not None else None
        if ckpt_path is not None:
            ckpt_path.parent.mkdir(parents=True, exist_ok=True)
        effective_restore_best = early_stopping if restore_best_weights is None else restore_best_weights

        if resume_from_checkpoint is not None:
            resume_payload = self._load_checkpoint_payload(resume_from_checkpoint, device=self.device)
            self._validate_checkpoint_payload_for_self(resume_payload)

            saved_target_dim = resume_payload.get("target_dim")
            if saved_target_dim is not None and self.target_dim_ is not None:
                if int(saved_target_dim) != int(self.target_dim_):
                    raise ValueError("Checkpoint target_dim does not match current regression target dimension.")

            state_dict = resume_payload.get("state_dict")
            if state_dict is None:
                raise ValueError("Checkpoint is missing 'state_dict'.")
            self.model_.load_state_dict(state_dict)

            opt_state = resume_payload.get("optimizer_state_dict")
            if isinstance(opt_state, dict):
                optimizer.load_state_dict(opt_state)
            if scheduler_obj is not None:
                sched_state = resume_payload.get("scheduler_state_dict")
                if isinstance(sched_state, dict):
                    try:
                        scheduler_obj.load_state_dict(sched_state)
                    except Exception:
                        pass

            saved_history = _clone_history_dict(resume_payload.get("history"))
            if saved_history:
                for key in history.keys():
                    if key in saved_history:
                        history[key] = saved_history[key]

        best_metric: float | None = None
        best_epoch: int | None = None
        best_state_dict: dict[str, torch.Tensor] | None = None
        wait = 0
        start_epoch = 0
        if resume_from_checkpoint is not None:
            checkpoint_meta = resume_payload.get("checkpoint")
            if isinstance(checkpoint_meta, dict) and checkpoint_meta.get("monitor") == resolved_monitor:
                if "best_score" in checkpoint_meta:
                    best_metric = float(checkpoint_meta["best_score"])
                if "best_epoch" in checkpoint_meta:
                    best_epoch = int(checkpoint_meta["best_epoch"])
            if isinstance(checkpoint_meta, dict) and "epoch" in checkpoint_meta:
                start_epoch = int(checkpoint_meta["epoch"])
            elif history.get("loss"):
                start_epoch = len(history["loss"])
        self.monitor_ = resolved_monitor
        self.best_epoch_ = None
        self.best_score_ = None
        self.stopped_epoch_ = None
        self.early_stopped_ = False
        cb_list = CallbackList(callbacks)
        cb_list.set_model(self)
        cb_list.set_context(optimizer=optimizer, scheduler=scheduler_obj, training_config=cfg)
        cb_list.on_train_begin(logs={"history": history, "start_epoch": start_epoch, "lr": self._current_lr(optimizer)})

        total_epochs_target = start_epoch + cfg.epochs
        for epoch_idx in range(cfg.epochs):
            epoch_num = start_epoch + epoch_idx + 1
            cb_list.on_epoch_begin(epoch_num, logs={"lr": self._current_lr(optimizer)})
            self.model_.train()
            total_loss = 0.0
            total_loss_denom = 0.0
            total_samples = 0
            sq_error_sum = 0.0
            abs_error_sum = 0.0
            total_targets = 0
            y_sum = 0.0
            y_sq_sum = 0.0
            has_lengths = lengths_tensor is not None
            has_sample_weight = sample_weight_tensor is not None

            for batch in train_loader:
                xb, yb, lengths_b, sample_weight_b = self._unpack_batch(
                    batch,
                    has_y=True,
                    has_lengths=has_lengths,
                    has_sample_weight=has_sample_weight,
                )
                assert yb is not None
                xb = xb.to(self.device)
                yb = yb.to(self.device)
                optimizer.zero_grad(set_to_none=True)
                preds = self._forward_model(xb, lengths_b)
                if preds.ndim == 1:
                    preds = preds.unsqueeze(1)
                diff = preds - yb
                per_sample_mse = (diff * diff).mean(dim=1)
                if sample_weight_b is not None:
                    sw = sample_weight_b.to(self.device)
                    denom_sum = sw.sum().clamp_min(1e-12)
                    loss = (per_sample_mse * sw).sum() / denom_sum
                    total_loss += float((per_sample_mse * sw).sum().item())
                    total_loss_denom += float(denom_sum.item())
                else:
                    loss = per_sample_mse.mean()
                    total_loss += float(per_sample_mse.sum().item())
                    total_loss_denom += float(per_sample_mse.numel())
                loss.backward()
                optimizer.step()

                sq_error_sum += float((diff * diff).sum().item())
                abs_error_sum += float(diff.abs().sum().item())
                total_samples += xb.size(0)
                total_targets += int(yb.numel())
                y_sum += float(yb.sum().item())
                y_sq_sum += float((yb * yb).sum().item())

            epoch_loss = total_loss / max(total_loss_denom, 1e-12)
            epoch_mse = sq_error_sum / max(total_targets, 1)
            epoch_mae = abs_error_sum / max(total_targets, 1)
            epoch_rmse = math.sqrt(epoch_mse)
            epoch_r2 = _r2_from_sums(
                sq_error_sum=sq_error_sum,
                y_sum=y_sum,
                y_sq_sum=y_sq_sum,
                count=total_targets,
            )
            history["loss"].append(epoch_loss)
            history["mse"].append(epoch_mse)
            history["mae"].append(epoch_mae)
            history["rmse"].append(epoch_rmse)
            history["r2"].append(epoch_r2)

            if val_X is not None and val_y is not None:
                metrics = self.evaluate(
                    val_X,
                    val_y,
                    lengths=val_lengths,
                    batch_size=cfg.batch_size,
                    _return_dict=True,
                )
                history["val_loss"].append(metrics["loss"])
                history["val_mse"].append(metrics["mse"])
                history["val_mae"].append(metrics["mae"])
                history["val_rmse"].append(metrics["rmse"])
                history["val_r2"].append(metrics["r2"])

            if cfg.verbose:
                if val_data is None:
                    print(
                        f"Epoch {epoch_num}/{total_epochs_target} - "
                        f"loss: {epoch_loss:.4f} - mae: {epoch_mae:.4f}"
                    )
                else:
                    print(
                        f"Epoch {epoch_num}/{total_epochs_target} - "
                        f"loss: {epoch_loss:.4f} - mae: {epoch_mae:.4f} - "
                        f"val_loss: {history['val_loss'][-1]:.4f} - val_mae: {history['val_mae'][-1]:.4f}"
                    )

            if scheduler_obj is not None:
                sched_metric = float(history[scheduler_metric_name][-1]) if scheduler_metric_name in history else None
                self._step_scheduler(scheduler_obj, metric=sched_metric)

            epoch_logs: dict[str, Any] = {k: float(v[-1]) for k, v in history.items() if isinstance(v, list) and v}
            epoch_logs["lr"] = self._current_lr(optimizer)
            cb_list.on_epoch_end(epoch_num, logs=epoch_logs)

            current_metric = float(history[resolved_monitor][-1])
            improved = self._is_improvement(
                current_metric,
                best_metric,
                mode=monitor_mode,
                min_delta=min_delta,
            )
            if improved:
                best_metric = current_metric
                best_epoch = epoch_num
                wait = 0
                if effective_restore_best or early_stopping or ckpt_path is not None:
                    best_state_dict = self._clone_state_dict_cpu(self.model_.state_dict())
                if ckpt_path is not None:
                    checkpoint_meta = {
                        "checkpoint": {
                            "epoch": epoch_num,
                            "best_epoch": best_epoch,
                            "best_score": best_metric,
                            "monitor": resolved_monitor,
                            "monitor_mode": monitor_mode,
                        },
                        "history": _clone_history_dict(history),
                        "optimizer_state_dict": optimizer.state_dict(),
                    }
                    if scheduler_obj is not None and hasattr(scheduler_obj, "state_dict"):
                        checkpoint_meta["scheduler_state_dict"] = scheduler_obj.state_dict()
                    torch.save(self._serialization_payload(extra=checkpoint_meta), str(ckpt_path))
            else:
                wait += 1

            if cb_list.stop_training:
                self.stopped_epoch_ = epoch_num
                self.early_stopped_ = True
                if cfg.verbose:
                    print(f"Training stopped by callback at epoch {epoch_num}.")
                break

            if early_stopping and wait >= patience:
                self.stopped_epoch_ = epoch_num
                self.early_stopped_ = True
                if cfg.verbose:
                    print(
                        f"Early stopping at epoch {epoch_num} "
                        f"(best {resolved_monitor}={best_metric:.4f} at epoch {best_epoch})."
                    )
                break

        if effective_restore_best and best_state_dict is not None:
            self.model_.load_state_dict(best_state_dict)
        self.best_epoch_ = best_epoch
        self.best_score_ = best_metric
        self.model_.eval()

        self.history_ = history
        cb_list.on_train_end(logs={"history": history, "best_epoch": self.best_epoch_, "best_score": self.best_score_})
        return self

    @torch.no_grad()
    def predict(
        self,
        X: Any,
        *,
        lengths: Any | None = None,
        batch_size: int | None = None,
    ) -> np.ndarray:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        X_tensor = self._prepare_X(X)
        lengths_tensor = self._prepare_lengths(lengths, X_tensor)
        loader = self._make_loader(
            X_tensor,
            shuffle=False,
            batch_size=batch_size,
            extra_tensors=[lengths_tensor] if lengths_tensor is not None else None,
        )
        self.model_.eval()
        preds: list[torch.Tensor] = []
        for batch in loader:
            if len(batch) == 1:
                (xb,) = batch
                lengths_b = None
            elif len(batch) == 2:
                xb, lengths_b = batch
            else:
                raise RuntimeError("Unexpected batch structure in regressor predict loop.")
            xb = xb.to(self.device)
            out = self._forward_model(xb, lengths_b)
            if out.ndim == 1:
                out = out.unsqueeze(1)
            preds.append(out.cpu())
        arr = torch.cat(preds, dim=0).numpy()
        if self._squeeze_output_ and arr.ndim == 2 and arr.shape[1] == 1:
            return arr[:, 0]
        return arr

    def predict_proba(self, X: Any, *, batch_size: int | None = None) -> np.ndarray:
        raise NotImplementedError("predict_proba is not available for regressors.")

    @torch.no_grad()
    def evaluate(
        self,
        X: Any,
        y: Any,
        *,
        lengths: Any | None = None,
        batch_size: int | None = None,
        _return_dict: bool = False,
    ) -> dict[str, float]:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        X_tensor = self._prepare_X(X)
        y_tensor = self._prepare_y_existing(y)
        lengths_tensor = self._prepare_lengths(lengths, X_tensor)
        if X_tensor.shape[0] != y_tensor.shape[0]:
            raise ValueError("X and y must have the same number of samples.")
        loader = self._make_loader(
            X_tensor,
            y_tensor,
            shuffle=False,
            batch_size=batch_size,
            extra_tensors=[lengths_tensor] if lengths_tensor is not None else None,
        )
        self.model_.eval()

        total_loss = 0.0
        total_samples = 0
        sq_error_sum = 0.0
        abs_error_sum = 0.0
        total_targets = 0
        y_sum = 0.0
        y_sq_sum = 0.0
        for batch in loader:
            if len(batch) == 2:
                xb, yb = batch
                lengths_b = None
            elif len(batch) == 3:
                xb, yb, lengths_b = batch
            else:
                raise RuntimeError("Unexpected batch structure in regressor evaluation loop.")
            xb = xb.to(self.device)
            yb = yb.to(self.device)
            preds = self._forward_model(xb, lengths_b)
            if preds.ndim == 1:
                preds = preds.unsqueeze(1)
            loss = self._criterion(preds, yb)
            diff = preds - yb
            total_loss += float(loss.item()) * xb.size(0)
            sq_error_sum += float((diff * diff).sum().item())
            abs_error_sum += float(diff.abs().sum().item())
            total_samples += xb.size(0)
            total_targets += int(yb.numel())
            y_sum += float(yb.sum().item())
            y_sq_sum += float((yb * yb).sum().item())

        mse = sq_error_sum / max(total_targets, 1)
        result = {
            "loss": total_loss / max(total_samples, 1),
            "mse": mse,
            "mae": abs_error_sum / max(total_targets, 1),
            "rmse": math.sqrt(mse),
            "r2": _r2_from_sums(
                sq_error_sum=sq_error_sum,
                y_sum=y_sum,
                y_sq_sum=y_sq_sum,
                count=total_targets,
            ),
        }
        if _return_dict:
            return result
        return result

    @classmethod
    def load(
        cls,
        path: str | Path,
        *,
        device: str | torch.device | None = None,
    ) -> "_BaseRegressor":
        """Load a saved regressor checkpoint produced by ``save()``."""
        path = Path(path)
        map_location = _default_device(device)

        try:
            payload = torch.load(str(path), map_location=map_location, weights_only=False)
        except TypeError:
            payload = torch.load(str(path), map_location=map_location)

        if not isinstance(payload, dict):
            raise ValueError(f"Invalid checkpoint format in {path!s}: expected dict payload.")

        saved_model_type = payload.get("model_type")
        if saved_model_type is not None and saved_model_type != cls.__name__:
            raise ValueError(
                f"Checkpoint model type {saved_model_type!r} does not match loader class {cls.__name__!r}."
            )

        input_shape = payload.get("input_shape")
        if input_shape is None:
            raise ValueError("Checkpoint is missing 'input_shape'.")
        input_shape = tuple(int(v) for v in input_shape)

        target_dim = payload.get("target_dim")
        if target_dim is None:
            raise ValueError("Checkpoint is missing 'target_dim'.")
        target_dim = int(target_dim)
        if target_dim < 1:
            raise ValueError("Checkpoint target_dim must be >= 1.")

        config = payload.get("config") or {}
        if not isinstance(config, dict):
            raise ValueError("Checkpoint 'config' must be a dict.")

        init_kwargs: dict[str, Any] = {}
        training_cfg = config.get("training", {})
        if isinstance(training_cfg, dict):
            for key in ("epochs", "batch_size", "lr", "weight_decay"):
                if key in training_cfg:
                    init_kwargs[key] = training_cfg[key]
        for key, value in config.items():
            if key != "training":
                init_kwargs[key] = value
        if device is not None:
            init_kwargs["device"] = device

        obj = cls(**init_kwargs)
        obj.target_dim_ = target_dim
        obj._squeeze_output_ = bool(payload.get("squeeze_output", target_dim == 1))
        obj._built_input_shape = input_shape
        obj.model_ = obj._build_model(input_shape, target_dim).to(obj.device)

        state_dict = payload.get("state_dict")
        if state_dict is None:
            raise ValueError("Checkpoint is missing 'state_dict'.")
        obj.model_.load_state_dict(state_dict)
        obj.model_.eval()
        return obj


class _MLPNet(nn.Module):
    def __init__(
        self,
        input_dim: int,
        num_classes: int,
        *,
        hidden_layers: Sequence[int],
        dropout: float,
        activation: str,
    ) -> None:
        super().__init__()
        activations = {
            "relu": nn.ReLU,
            "tanh": nn.Tanh,
            "gelu": nn.GELU,
        }
        if activation.lower() not in activations:
            raise ValueError(f"Unsupported activation: {activation}")
        act_cls = activations[activation.lower()]

        layers: list[nn.Module] = [nn.Flatten()]
        in_features = input_dim
        for hidden in hidden_layers:
            layers.append(nn.Linear(in_features, int(hidden)))
            layers.append(act_cls())
            if dropout > 0:
                layers.append(nn.Dropout(dropout))
            in_features = int(hidden)
        layers.append(nn.Linear(in_features, num_classes))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class MLPClassifier(_BaseClassifier):
    """Feedforward network for tabular or flattened inputs."""

    def __init__(
        self,
        *,
        hidden_layers: Sequence[int] = (128, 64),
        dropout: float = 0.0,
        activation: str = "relu",
        device: str | torch.device | None = None,
        epochs: int = 10,
        batch_size: int = 32,
        lr: float = 1e-3,
        weight_decay: float = 0.0,
        random_state: int | None = None,
    ) -> None:
        super().__init__(
            device=device,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            weight_decay=weight_decay,
            random_state=random_state,
        )
        self.hidden_layers = tuple(int(v) for v in hidden_layers)
        self.dropout = float(dropout)
        self.activation = activation

    def _prepare_X(self, X: Any) -> torch.Tensor:
        x = _as_float_tensor(X)
        if x.ndim == 1:
            x = x.unsqueeze(1)
        if x.ndim < 2:
            raise ValueError("MLP input must have at least 2 dimensions (batch, features).")
        return x

    def _build_model(self, input_shape: tuple[int, ...], num_classes: int) -> nn.Module:
        input_dim = int(np.prod(input_shape))
        return _MLPNet(
            input_dim=input_dim,
            num_classes=num_classes,
            hidden_layers=self.hidden_layers,
            dropout=self.dropout,
            activation=self.activation,
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "hidden_layers": self.hidden_layers,
            "dropout": self.dropout,
            "activation": self.activation,
            "training": self.training.__dict__,
        }


class _CNNNet(nn.Module):
    def __init__(
        self,
        in_channels: int,
        num_classes: int,
        *,
        conv_channels: Sequence[int],
        kernel_size: int,
        dropout: float,
        batch_norm: bool,
        classifier_hidden: int,
    ) -> None:
        super().__init__()
        padding = kernel_size // 2
        features: list[nn.Module] = []
        current_channels = in_channels
        for out_channels in conv_channels:
            features.append(nn.Conv2d(current_channels, int(out_channels), kernel_size=kernel_size, padding=padding))
            if batch_norm:
                features.append(nn.BatchNorm2d(int(out_channels)))
            features.append(nn.ReLU())
            features.append(nn.MaxPool2d(kernel_size=2))
            current_channels = int(out_channels)

        self.features = nn.Sequential(*features) if features else nn.Identity()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(current_channels, classifier_hidden),
            nn.ReLU(),
            nn.Dropout(dropout) if dropout > 0 else nn.Identity(),
            nn.Linear(classifier_hidden, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = self.pool(x)
        return self.classifier(x)


class CNNClassifier(_BaseClassifier):
    """Simple CNN classifier for image tensors."""

    def __init__(
        self,
        *,
        conv_channels: Sequence[int] = (16, 32),
        kernel_size: int = 3,
        dropout: float = 0.1,
        batch_norm: bool = True,
        classifier_hidden: int = 128,
        channels_last: bool = False,
        device: str | torch.device | None = None,
        epochs: int = 10,
        batch_size: int = 32,
        lr: float = 1e-3,
        weight_decay: float = 0.0,
        random_state: int | None = None,
    ) -> None:
        super().__init__(
            device=device,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            weight_decay=weight_decay,
            random_state=random_state,
        )
        self.conv_channels = tuple(int(v) for v in conv_channels)
        self.kernel_size = int(kernel_size)
        self.dropout = float(dropout)
        self.batch_norm = bool(batch_norm)
        self.classifier_hidden = int(classifier_hidden)
        self.channels_last = bool(channels_last)

    def _prepare_X(self, X: Any) -> torch.Tensor:
        x = _as_float_tensor(X)
        original_ndim = x.ndim
        if original_ndim == 3:
            # grayscale images: (N, H, W) -> (N, 1, H, W)
            x = x.unsqueeze(1)
        if x.ndim != 4:
            raise ValueError("CNN input must have shape (N, C, H, W) or (N, H, W).")
        if self.channels_last and original_ndim == 4:
            x = x.permute(0, 3, 1, 2).contiguous()
        return x

    def _build_model(self, input_shape: tuple[int, ...], num_classes: int) -> nn.Module:
        if len(input_shape) != 3:
            raise ValueError("CNN input shape must be (C, H, W).")
        in_channels = int(input_shape[0])
        return _CNNNet(
            in_channels=in_channels,
            num_classes=num_classes,
            conv_channels=self.conv_channels,
            kernel_size=self.kernel_size,
            dropout=self.dropout,
            batch_norm=self.batch_norm,
            classifier_hidden=self.classifier_hidden,
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "conv_channels": self.conv_channels,
            "kernel_size": self.kernel_size,
            "dropout": self.dropout,
            "batch_norm": self.batch_norm,
            "classifier_hidden": self.classifier_hidden,
            "channels_last": self.channels_last,
            "training": self.training.__dict__,
        }


class _RNNNet(nn.Module):
    def __init__(
        self,
        input_size: int,
        num_classes: int,
        *,
        hidden_size: int,
        num_layers: int,
        dropout: float,
        bidirectional: bool,
        rnn_type: str,
        aggregate: str,
    ) -> None:
        super().__init__()
        rnn_type = rnn_type.lower()
        rnn_map: dict[str, type[nn.Module]] = {
            "rnn": nn.RNN,
            "gru": nn.GRU,
            "lstm": nn.LSTM,
        }
        if rnn_type not in rnn_map:
            raise ValueError(f"Unsupported rnn_type: {rnn_type!r}")
        if aggregate not in {"last", "mean"}:
            raise ValueError("aggregate must be 'last' or 'mean'")

        rnn_dropout = dropout if num_layers > 1 else 0.0
        self.rnn_type = rnn_type
        self.aggregate = aggregate
        self.bidirectional = bidirectional
        self.num_layers = num_layers
        self.hidden_size = hidden_size

        rnn_cls = rnn_map[rnn_type]
        self.rnn = rnn_cls(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=rnn_dropout,
            bidirectional=bidirectional,
        )

        out_dim = hidden_size * (2 if bidirectional else 1)
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        self.fc = nn.Linear(out_dim, num_classes)

    def _last_hidden(self, hidden: torch.Tensor) -> torch.Tensor:
        if self.bidirectional:
            forward = hidden[-2]
            backward = hidden[-1]
            return torch.cat([forward, backward], dim=1)
        return hidden[-1]

    def forward(self, x: torch.Tensor, lengths: torch.Tensor | None = None) -> torch.Tensor:
        outputs: torch.Tensor | None
        if lengths is None:
            outputs, hidden = self.rnn(x)
            lengths_cpu = None
        else:
            lengths_cpu = lengths.detach().to(dtype=torch.long, device="cpu")
            if lengths_cpu.ndim != 1 or int(lengths_cpu.numel()) != int(x.size(0)):
                raise ValueError("lengths must be a 1D tensor with one length per sequence.")
            if torch.any(lengths_cpu < 1):
                raise ValueError("All sequence lengths must be >= 1.")
            if torch.any(lengths_cpu > x.size(1)):
                raise ValueError("Sequence lengths cannot exceed the padded sequence length.")
            packed = nn.utils.rnn.pack_padded_sequence(
                x,
                lengths_cpu,
                batch_first=True,
                enforce_sorted=False,
            )
            packed_outputs, hidden = self.rnn(packed)
            outputs = None
            if self.aggregate == "mean":
                outputs, _ = nn.utils.rnn.pad_packed_sequence(
                    packed_outputs,
                    batch_first=True,
                    total_length=x.size(1),
                )

        if self.aggregate == "mean":
            assert outputs is not None
            if lengths_cpu is None:
                feat = outputs.mean(dim=1)
            else:
                lengths_dev = lengths_cpu.to(outputs.device)
                time_index = torch.arange(outputs.size(1), device=outputs.device).unsqueeze(0)
                mask = (time_index < lengths_dev.unsqueeze(1)).unsqueeze(-1).to(outputs.dtype)
                denom = lengths_dev.clamp_min(1).to(outputs.dtype).unsqueeze(1)
                feat = (outputs * mask).sum(dim=1) / denom
        else:
            if self.rnn_type == "lstm":
                hidden = hidden[0]
            feat = self._last_hidden(hidden)
        feat = self.dropout(feat)
        return self.fc(feat)


class RNNClassifier(_BaseClassifier):
    """Sequence classifier using RNN/GRU/LSTM backends."""

    def __init__(
        self,
        *,
        hidden_size: int = 64,
        num_layers: int = 1,
        dropout: float = 0.0,
        bidirectional: bool = False,
        rnn_type: str = "rnn",
        aggregate: str = "last",
        device: str | torch.device | None = None,
        epochs: int = 10,
        batch_size: int = 32,
        lr: float = 1e-3,
        weight_decay: float = 0.0,
        random_state: int | None = None,
    ) -> None:
        super().__init__(
            device=device,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            weight_decay=weight_decay,
            random_state=random_state,
        )
        self.hidden_size = int(hidden_size)
        self.num_layers = int(num_layers)
        self.dropout = float(dropout)
        self.bidirectional = bool(bidirectional)
        self.rnn_type = rnn_type.lower()
        self.aggregate = aggregate

    def _prepare_X(self, X: Any) -> torch.Tensor:
        x = _as_float_tensor(X)
        if x.ndim == 2:
            # Univariate sequence: (N, T) -> (N, T, 1)
            x = x.unsqueeze(-1)
        if x.ndim != 3:
            raise ValueError("RNN input must have shape (N, T, F) or (N, T).")
        return x

    def _prepare_lengths(self, lengths: Any | None, X_tensor: torch.Tensor) -> torch.Tensor | None:
        if lengths is None:
            return None
        return _as_lengths_tensor(
            lengths,
            num_samples=int(X_tensor.shape[0]),
            max_len=int(X_tensor.shape[1]),
        )

    def _forward_model(self, xb: torch.Tensor, lengths: torch.Tensor | None = None) -> torch.Tensor:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        return self.model_(xb, lengths=lengths)

    def _build_model(self, input_shape: tuple[int, ...], num_classes: int) -> nn.Module:
        if len(input_shape) != 2:
            raise ValueError("RNN input shape must be (T, F).")
        _, input_size = input_shape
        return _RNNNet(
            input_size=int(input_size),
            num_classes=num_classes,
            hidden_size=self.hidden_size,
            num_layers=self.num_layers,
            dropout=self.dropout,
            bidirectional=self.bidirectional,
            rnn_type=self.rnn_type,
            aggregate=self.aggregate,
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "hidden_size": self.hidden_size,
            "num_layers": self.num_layers,
            "dropout": self.dropout,
            "bidirectional": self.bidirectional,
            "rnn_type": self.rnn_type,
            "aggregate": self.aggregate,
            "training": self.training.__dict__,
        }


class MLPRegressor(_BaseRegressor):
    """Feedforward network for tabular or flattened regression inputs."""

    def __init__(
        self,
        *,
        hidden_layers: Sequence[int] = (128, 64),
        dropout: float = 0.0,
        activation: str = "relu",
        device: str | torch.device | None = None,
        epochs: int = 10,
        batch_size: int = 32,
        lr: float = 1e-3,
        weight_decay: float = 0.0,
        random_state: int | None = None,
    ) -> None:
        super().__init__(
            device=device,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            weight_decay=weight_decay,
            random_state=random_state,
        )
        self.hidden_layers = tuple(int(v) for v in hidden_layers)
        self.dropout = float(dropout)
        self.activation = activation

    def _prepare_X(self, X: Any) -> torch.Tensor:
        x = _as_float_tensor(X)
        if x.ndim == 1:
            x = x.unsqueeze(1)
        if x.ndim < 2:
            raise ValueError("MLP input must have at least 2 dimensions (batch, features).")
        return x

    def _build_model(self, input_shape: tuple[int, ...], output_dim: int) -> nn.Module:
        input_dim = int(np.prod(input_shape))
        return _MLPNet(
            input_dim=input_dim,
            num_classes=output_dim,
            hidden_layers=self.hidden_layers,
            dropout=self.dropout,
            activation=self.activation,
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "hidden_layers": self.hidden_layers,
            "dropout": self.dropout,
            "activation": self.activation,
            "training": self.training.__dict__,
        }


class CNNRegressor(_BaseRegressor):
    """Simple CNN regressor for image tensors."""

    def __init__(
        self,
        *,
        conv_channels: Sequence[int] = (16, 32),
        kernel_size: int = 3,
        dropout: float = 0.1,
        batch_norm: bool = True,
        classifier_hidden: int = 128,
        channels_last: bool = False,
        device: str | torch.device | None = None,
        epochs: int = 10,
        batch_size: int = 32,
        lr: float = 1e-3,
        weight_decay: float = 0.0,
        random_state: int | None = None,
    ) -> None:
        super().__init__(
            device=device,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            weight_decay=weight_decay,
            random_state=random_state,
        )
        self.conv_channels = tuple(int(v) for v in conv_channels)
        self.kernel_size = int(kernel_size)
        self.dropout = float(dropout)
        self.batch_norm = bool(batch_norm)
        self.classifier_hidden = int(classifier_hidden)
        self.channels_last = bool(channels_last)

    def _prepare_X(self, X: Any) -> torch.Tensor:
        x = _as_float_tensor(X)
        original_ndim = x.ndim
        if original_ndim == 3:
            x = x.unsqueeze(1)
        if x.ndim != 4:
            raise ValueError("CNN input must have shape (N, C, H, W) or (N, H, W).")
        if self.channels_last and original_ndim == 4:
            x = x.permute(0, 3, 1, 2).contiguous()
        return x

    def _build_model(self, input_shape: tuple[int, ...], output_dim: int) -> nn.Module:
        if len(input_shape) != 3:
            raise ValueError("CNN input shape must be (C, H, W).")
        in_channels = int(input_shape[0])
        return _CNNNet(
            in_channels=in_channels,
            num_classes=output_dim,
            conv_channels=self.conv_channels,
            kernel_size=self.kernel_size,
            dropout=self.dropout,
            batch_norm=self.batch_norm,
            classifier_hidden=self.classifier_hidden,
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "conv_channels": self.conv_channels,
            "kernel_size": self.kernel_size,
            "dropout": self.dropout,
            "batch_norm": self.batch_norm,
            "classifier_hidden": self.classifier_hidden,
            "channels_last": self.channels_last,
            "training": self.training.__dict__,
        }


class RNNRegressor(_BaseRegressor):
    """Sequence regressor using RNN/GRU/LSTM backends."""

    def __init__(
        self,
        *,
        hidden_size: int = 64,
        num_layers: int = 1,
        dropout: float = 0.0,
        bidirectional: bool = False,
        rnn_type: str = "rnn",
        aggregate: str = "last",
        device: str | torch.device | None = None,
        epochs: int = 10,
        batch_size: int = 32,
        lr: float = 1e-3,
        weight_decay: float = 0.0,
        random_state: int | None = None,
    ) -> None:
        super().__init__(
            device=device,
            epochs=epochs,
            batch_size=batch_size,
            lr=lr,
            weight_decay=weight_decay,
            random_state=random_state,
        )
        self.hidden_size = int(hidden_size)
        self.num_layers = int(num_layers)
        self.dropout = float(dropout)
        self.bidirectional = bool(bidirectional)
        self.rnn_type = rnn_type.lower()
        self.aggregate = aggregate

    def _prepare_X(self, X: Any) -> torch.Tensor:
        x = _as_float_tensor(X)
        if x.ndim == 2:
            x = x.unsqueeze(-1)
        if x.ndim != 3:
            raise ValueError("RNN input must have shape (N, T, F) or (N, T).")
        return x

    def _prepare_lengths(self, lengths: Any | None, X_tensor: torch.Tensor) -> torch.Tensor | None:
        if lengths is None:
            return None
        return _as_lengths_tensor(
            lengths,
            num_samples=int(X_tensor.shape[0]),
            max_len=int(X_tensor.shape[1]),
        )

    def _forward_model(self, xb: torch.Tensor, lengths: torch.Tensor | None = None) -> torch.Tensor:
        if self.model_ is None:
            raise RuntimeError("Model has not been fitted.")
        return self.model_(xb, lengths=lengths)

    def _build_model(self, input_shape: tuple[int, ...], output_dim: int) -> nn.Module:
        if len(input_shape) != 2:
            raise ValueError("RNN input shape must be (T, F).")
        _, input_size = input_shape
        return _RNNNet(
            input_size=int(input_size),
            num_classes=output_dim,
            hidden_size=self.hidden_size,
            num_layers=self.num_layers,
            dropout=self.dropout,
            bidirectional=self.bidirectional,
            rnn_type=self.rnn_type,
            aggregate=self.aggregate,
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "hidden_size": self.hidden_size,
            "num_layers": self.num_layers,
            "dropout": self.dropout,
            "bidirectional": self.bidirectional,
            "rnn_type": self.rnn_type,
            "aggregate": self.aggregate,
            "training": self.training.__dict__,
        }
