PSO/pso/optimizer.py

import collections
import copy
import csv
import json
import math
import os
from typing import Any, Literal, Sequence
import torch
import torch.nn as nn

from ._version import __version__
from ._weights import ParameterCodec
from .particle import Particle
from .plugins import (
    BasePlugin,
    ConvergencePlugin,
    EvaluationPlugin,
    FitContext,
    InitializationPlugin,
    IterationContext,
    MovementPlugin,
    PluginMetadata,
    RefinementPlugin,
    SwarmState,
    _is_at_least_delta,
    _is_better_score,
    get_plugin,
)


def resolve_device(device: str | torch.device | None = None) -> torch.device:
    """
    Resolves execution device.
    If device is None, auto-selects mps if available & built, else cuda, else cpu.
    Explicit device requirement ('mps', 'cuda', 'cpu') checks availability or raises RuntimeError.
    """
    if device is None:
        if (
            hasattr(torch.backends, "mps")
            and torch.backends.mps.is_built()
            and torch.backends.mps.is_available()
        ):
            return torch.device("mps")
        elif torch.cuda.is_available():
            return torch.device("cuda")
        else:
            return torch.device("cpu")

    if isinstance(device, str):
        try:
            dev = torch.device(device)
        except RuntimeError as e:
            raise ValueError(f"Unsupported device type: '{device}'") from e
    elif isinstance(device, torch.device):
        dev = device
    else:
        raise TypeError(
            f"device must be a string, torch.device, or None, got {type(device)}"
        )

    if dev.type == "mps":
        built = hasattr(torch.backends, "mps") and torch.backends.mps.is_built()
        avail = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
        if not (built and avail):
            raise RuntimeError(
                f"Explicit MPS device requested ('{device}'), but PyTorch MPS backend is not available (built={built}, available={avail})."
            )
    elif dev.type == "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError(
                f"Explicit CUDA device requested ('{device}'), but CUDA is not available."
            )
    elif dev.type == "cpu":
        pass
    else:
        raise ValueError(
            f"Unsupported device type: '{dev.type}'. Only 'cpu', 'cuda', and 'mps' are supported."
        )

    return dev


def _is_better_score(
    new_score: tuple[float, float, float],
    best_score: tuple[float, float, float] | None,
    renewal: str = "acc",
) -> bool:
    """
    Returns True if new_score is strictly better than best_score according to renewal
    metric and deterministic tie breaks (loss asc, accuracy desc, mse asc).
    """
    if best_score is None:
        return True

    n_loss, n_acc, n_mse = new_score
    b_loss, b_acc, b_mse = best_score

    if renewal in ("acc", "accuracy"):
        if n_acc != b_acc:
            return n_acc > b_acc
    elif renewal == "loss":
        if n_loss != b_loss:
            return n_loss < b_loss
    elif renewal == "mse":
        if n_mse != b_mse:
            return n_mse < b_mse
    else:
        raise ValueError(f"Unknown renewal metric: {renewal}")

    if n_loss != b_loss:
        return n_loss < b_loss
    if n_acc != b_acc:
        return n_acc > b_acc
    if n_mse != b_mse:
        return n_mse < b_mse

    return False


def _is_at_least_delta(delta: float, min_delta: float) -> bool:
    """
    Returns True if directional improvement delta meets min_delta, handling floating-point boundaries.
    """
    if min_delta > 0:
        return delta > min_delta or math.isclose(
            delta, min_delta, rel_tol=1e-12, abs_tol=1e-15
        )
    return delta > 0


def _validate_score(
    score: Sequence[Any], particle_idx: int, iteration: int
) -> tuple[float, float, float]:
    """
    Validates that a score sequence contains 3 finite numbers.
    Raises FloatingPointError naming particle index and iteration if non-finite.
    """
    try:
        parsed = (float(score[0]), float(score[1]), float(score[2]))
    except (IndexError, TypeError, ValueError) as e:
        raise FloatingPointError(
            f"Invalid score format {score} for particle {particle_idx} at iteration {iteration}"
        ) from e

    if not all(math.isfinite(x) for x in parsed):
        raise FloatingPointError(
            f"Non-finite score {parsed} encountered for particle {particle_idx} at iteration {iteration}"
        )
    return parsed


class _RandomSource:
    """
    Private CPU/device random generator wrapper for drawing stochastic tensors and events.
    Transfers generated tensors to reference device and dtype if needed.
    """

    def __init__(self, seed: int | None = None, device: torch.device | str = "cpu"):
        self.cpu_generator = torch.Generator(device="cpu")
        dev = resolve_device(device) if not isinstance(device, torch.device) else device
        self.search_generator = (
            self.cpu_generator
            if dev.type == "cpu"
            else torch.Generator(device=dev)
        )

        if seed is not None:
            self.cpu_generator.manual_seed(seed)
            if self.search_generator is not self.cpu_generator:
                self.search_generator.manual_seed(seed)
        else:
            self.cpu_generator.seed()
            if self.search_generator is not self.cpu_generator:
                self.search_generator.seed()

    def uniform(
        self,
        shape: torch.Size | tuple[int, ...],
        low: float = 0.0,
        high: float = 1.0,
        device: torch.device | str | None = None,
        dtype: torch.dtype = torch.float32,
    ) -> torch.Tensor:
        if device is None:
            target_dev = self.search_generator.device
        else:
            target_dev = (
                resolve_device(device) if not isinstance(device, torch.device) else device
            )

        if self.search_generator.device.type == target_dev.type:
            r = torch.rand(
                shape, generator=self.search_generator, device=target_dev, dtype=dtype
            )
        else:
            r = torch.rand(
                shape, generator=self.cpu_generator, device="cpu", dtype=dtype
            ).to(device=target_dev, dtype=dtype)
        return low + (high - low) * r
    def bernoulli_event(self, p: float) -> bool:
        r = torch.rand((1,), generator=self.cpu_generator, device="cpu").item()
        return r < p

    def permutation(self, n: int) -> torch.Tensor:
        return torch.randperm(n, generator=self.cpu_generator, device="cpu")

    def choice(self, n: int, size: int) -> torch.Tensor:
        perm = torch.randperm(n, generator=self.cpu_generator, device="cpu")
        return perm[:size]

    def randint(
        self,
        low: int,
        high: int,
        size: tuple[int, ...] | int | None = None,
        device: torch.device | str | None = None,
    ) -> torch.Tensor:
        if size is None:
            shape = (1,)
        elif isinstance(size, int):
            shape = (size,)
        else:
            shape = size
        r = torch.randint(
            low=low,
            high=high,
            size=shape,
            generator=self.cpu_generator,
            device="cpu",
        )
        if device is not None and str(device) != "cpu":
            return r.to(device)
        return r

class Optimizer:
    """
    Particle Swarm Optimizer for PyTorch nn.Module models with stage-plugin architecture.
    """

    def __init__(
        self,
        model: nn.Module,
        loss: nn.Module,
        *,
        task: Literal["binary", "multiclass", "regression"],
        method: str | MovementPlugin = "original",
        initialization: str | InitializationPlugin = "model_noise",
        evaluation: str | EvaluationPlugin = "full",
        convergence: str | ConvergencePlugin = "none",
        refinement: str | RefinementPlugin = "none",
        method_options: dict[str, Any] | None = None,
        n_particles: int = 10,
        c0: float | None = None,
        c1: float | None = None,
        w_min: float | None = None,
        w_max: float | None = None,
        negative_swarm: float = 0.0,
        mutation_swarm: float = 0.0,
        particle_min: float | None = None,
        particle_max: float | None = None,
        velocity_limit_ratio: float | None = None,
        boundary_strategy: Literal["clip", "reflect"] = "clip",
        initial_position_noise: float = 0.05,
        seed: int | None = None,
        device: str | torch.device | None = None,
        fitness_size: int | None = None,
        convergence_patience: int = 10,
        convergence_min_delta: float = 0.0001,
        convergence_monitor: str = "loss",
        refinement_epochs: int = 0,
        refinement_lr: float = 0.001,
        moment_blend: float | None = None,
        moment_beta1: float | None = None,
        moment_beta2: float | None = None,
        moment_step_size: float | None = None,
        moment_epsilon: float | None = None,
    ):
        if model is None or not isinstance(model, nn.Module):
            raise ValueError("model must be an instance of torch.nn.Module")

        if loss is None or not isinstance(loss, nn.Module):
            raise ValueError("loss must be an instance of torch.nn.Module")

        if task not in ("binary", "multiclass", "regression"):
            raise ValueError(
                "task must be one of 'binary', 'multiclass', 'regression'"
            )

        if (
            isinstance(n_particles, bool)
            or not isinstance(n_particles, int)
            or n_particles < 1
        ):
            raise ValueError("n_particles must be an integer >= 1")

        for name, val in [
            ("c0", c0),
            ("c1", c1),
            ("w_min", w_min),
            ("w_max", w_max),
            ("moment_blend", moment_blend),
            ("moment_beta1", moment_beta1),
            ("moment_beta2", moment_beta2),
            ("moment_step_size", moment_step_size),
            ("moment_epsilon", moment_epsilon),
        ]:
            if val is not None and (
                isinstance(val, bool)
                or not isinstance(val, (int, float))
                or not math.isfinite(val)
            ):
                raise ValueError(f"{name} must be a finite number")

        if w_min is not None and w_max is not None and float(w_min) > float(w_max):
            raise ValueError("w_min must be <= w_max")

        for name, val in [
            ("negative_swarm", negative_swarm),
            ("mutation_swarm", mutation_swarm),
        ]:
            if (
                isinstance(val, bool)
                or not isinstance(val, (int, float))
                or not math.isfinite(val)
                or not (0.0 <= float(val) <= 1.0)
            ):
                raise ValueError(f"{name} must be a finite float in range [0, 1]")

        if (particle_min is None) != (particle_max is None):
            raise ValueError("particle_min and particle_max must be provided together")

        if particle_min is not None and particle_max is not None:
            if (
                isinstance(particle_min, bool)
                or not isinstance(particle_min, (int, float))
                or not math.isfinite(particle_min)
            ):
                raise ValueError("particle_min must be a finite number")
            if (
                isinstance(particle_max, bool)
                or not isinstance(particle_max, (int, float))
                or not math.isfinite(particle_max)
            ):
                raise ValueError("particle_max must be a finite number")
            if particle_min > particle_max:
                raise ValueError("particle_min must be <= particle_max")

        if velocity_limit_ratio is not None:
            if (
                isinstance(velocity_limit_ratio, bool)
                or not isinstance(velocity_limit_ratio, (int, float))
                or not math.isfinite(velocity_limit_ratio)
                or not (0.0 < float(velocity_limit_ratio) <= 1.0)
            ):
                raise ValueError(
                    "velocity_limit_ratio must be a finite float in range (0, 1]"
                )
            if particle_min is None or particle_max is None:
                raise ValueError(
                    "velocity_limit_ratio requires paired particle_min and particle_max bounds"
                )
            if particle_min >= particle_max:
                raise ValueError(
                    "velocity_limit_ratio requires particle_min < particle_max"
                )

        if boundary_strategy not in ("clip", "reflect"):
            raise ValueError("boundary_strategy must be one of 'clip', 'reflect'")

        if boundary_strategy == "reflect":
            if particle_min is None or particle_max is None:
                raise ValueError(
                    "boundary_strategy 'reflect' requires paired particle_min and particle_max bounds"
                )
            if particle_min >= particle_max:
                raise ValueError(
                    "boundary_strategy 'reflect' requires particle_min < particle_max"
                )

        if (
            isinstance(initial_position_noise, bool)
            or not isinstance(initial_position_noise, (int, float))
            or not math.isfinite(initial_position_noise)
            or initial_position_noise < 0.0
        ):
            raise ValueError(
                "initial_position_noise must be a finite nonnegative number"
            )

        if seed is not None:
            if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0:
                raise ValueError("seed must be a non-negative integer")

        if (
            isinstance(convergence_patience, bool)
            or not isinstance(convergence_patience, int)
            or convergence_patience <= 0
        ):
            raise ValueError("convergence_patience must be a positive integer")

        if (
            isinstance(convergence_min_delta, bool)
            or not isinstance(convergence_min_delta, (int, float))
            or not math.isfinite(convergence_min_delta)
            or convergence_min_delta < 0
        ):
            raise ValueError(
                "convergence_min_delta must be a finite nonnegative number"
            )

        if convergence_monitor not in ("loss", "acc", "accuracy", "mse"):
            raise ValueError(
                "convergence_monitor must be one of 'loss', 'acc', 'accuracy', 'mse'"
            )

        if (
            isinstance(refinement_epochs, bool)
            or not isinstance(refinement_epochs, int)
            or refinement_epochs < 0
        ):
            raise ValueError("refinement_epochs must be a non-negative integer")

        if (
            isinstance(refinement_lr, bool)
            or not isinstance(refinement_lr, (int, float))
            or not math.isfinite(refinement_lr)
            or float(refinement_lr) <= 0.0
        ):
            raise ValueError("refinement_lr must be a positive finite float")

        self.device = resolve_device(device)
        self.task = task

        self.model = copy.deepcopy(model)
        self.eval_model = copy.deepcopy(model).to(self.device)
        self.eval_loss = copy.deepcopy(loss).to(self.device)

        if self.device.type == "mps":
            self.eval_model = self.eval_model.to(dtype=torch.float32)
            if hasattr(self.eval_loss, "to"):
                self.eval_loss = self.eval_loss.to(dtype=torch.float32)

        self.eval_model.eval()
        if hasattr(self.eval_loss, "eval"):
            self.eval_loss.eval()

        self.codec = ParameterCodec(self.eval_model)
        self._base_vector = self.codec.encode(self.eval_model).clone().detach()
        base_vector = self._base_vector.clone()

        self.n_particles = n_particles
        self.negative_swarm = float(negative_swarm)
        self.mutation_swarm = float(mutation_swarm)
        self.particle_min = float(particle_min) if particle_min is not None else None
        self.particle_max = float(particle_max) if particle_max is not None else None
        self.velocity_limit_ratio = (
            float(velocity_limit_ratio) if velocity_limit_ratio is not None else None
        )
        self.boundary_strategy = boundary_strategy
        self.initial_position_noise = float(initial_position_noise)

        if (
            self.velocity_limit_ratio is not None
            and self.particle_min is not None
            and self.particle_max is not None
        ):
            self.velocity_limit = float(
                self.velocity_limit_ratio * (self.particle_max - self.particle_min)
            )
        else:
            self.velocity_limit = None

        self.seed = seed
        self.renewal = "acc"
        self.fitness_size = fitness_size
        self.convergence_patience = convergence_patience
        self.convergence_min_delta = float(convergence_min_delta)
        self.convergence_monitor = convergence_monitor
        self.refinement_epochs = refinement_epochs
        self.refinement_lr = float(refinement_lr)

        self._method_selector = method if isinstance(method, str) else method.metadata.title
        self._initialization_selector = initialization if isinstance(initialization, str) else initialization.metadata.title
        self._evaluation_selector = evaluation if isinstance(evaluation, str) else evaluation.metadata.title
        self._convergence_selector = convergence if isinstance(convergence, str) else convergence.metadata.title
        self._refinement_selector = refinement if isinstance(refinement, str) else refinement.metadata.title

        # Split stage options cleanly to prevent cross-stage consumption
        movement_opts = dict(method_options or {})
        if c0 is not None:
            movement_opts["c0"] = c0
        if c1 is not None:
            movement_opts["c1"] = c1
        if w_min is not None:
            movement_opts["w_min"] = w_min
        if w_max is not None:
            movement_opts["w_max"] = w_max
        if moment_blend is not None:
            movement_opts["moment_blend"] = moment_blend
        if moment_beta1 is not None:
            movement_opts["moment_beta1"] = moment_beta1
        if moment_beta2 is not None:
            movement_opts["moment_beta2"] = moment_beta2
        if moment_step_size is not None:
            movement_opts["moment_step_size"] = moment_step_size
        if moment_epsilon is not None:
            movement_opts["moment_epsilon"] = moment_epsilon

        init_opts = {}
        if self._initialization_selector not in ("uniform", "Uniform Bounded Space Initialization"):
            if initial_position_noise != 0.05:
                init_opts["noise"] = initial_position_noise

        eval_opts = {}
        if self._evaluation_selector not in ("full", "Full Dataset Evaluation"):
            if fitness_size is not None:
                eval_opts["fitness_size"] = fitness_size

        conv_opts = {}
        if self._convergence_selector not in ("none", "No Convergence Action"):
            conv_opts = {
                "patience": convergence_patience,
                "min_delta": convergence_min_delta,
                "monitor": convergence_monitor,
            }

        refine_opts = {}
        if self._refinement_selector not in ("none", "No Refinement"):
            refine_opts = {
                "epochs": refinement_epochs,
                "lr": refinement_lr,
            }
        # Handle custom movement plugin conflict validation
        if isinstance(method, MovementPlugin):
            for k in ("c0", "c1", "w_min", "w_max"):
                if getattr(method, k, None) is not None and movement_opts.get(k) is not None:
                    if float(getattr(method, k)) != float(movement_opts[k]):
                        raise ValueError(
                            f"Conflicting parameter '{k}' provided for preconfigured custom movement instance"
                        )

        self.movement_plugin: MovementPlugin = get_plugin("movement", method, movement_opts)  # type: ignore
        self.initialization_plugin: InitializationPlugin = get_plugin(
            "initialization", initialization, init_opts
        )  # type: ignore
        self.evaluation_plugin: EvaluationPlugin = get_plugin(
            "evaluation", evaluation, eval_opts
        )  # type: ignore
        self.convergence_plugin: ConvergencePlugin = get_plugin(
            "convergence", convergence, conv_opts
        )  # type: ignore
        self.refinement_plugin: RefinementPlugin = get_plugin(
            "refinement", refinement, refine_opts
        )  # type: ignore

        # Resolve scalar parameter attributes from movement plugin (or None if irrelevant)
        self.c0 = (
            float(getattr(self.movement_plugin, "c0"))
            if hasattr(self.movement_plugin, "c0") and getattr(self.movement_plugin, "c0", None) is not None
            else None
        )
        self.c1 = (
            float(getattr(self.movement_plugin, "c1"))
            if hasattr(self.movement_plugin, "c1") and getattr(self.movement_plugin, "c1", None) is not None
            else None
        )
        self.w_min = (
            float(getattr(self.movement_plugin, "w_min"))
            if hasattr(self.movement_plugin, "w_min") and getattr(self.movement_plugin, "w_min", None) is not None
            else None
        )
        self.w_max = (
            float(getattr(self.movement_plugin, "w_max"))
            if hasattr(self.movement_plugin, "w_max") and getattr(self.movement_plugin, "w_max", None) is not None
            else None
        )

        self.moment_blend = (
            float(getattr(self.movement_plugin, "moment_blend"))
            if hasattr(self.movement_plugin, "moment_blend")
            else (float(moment_blend) if moment_blend is not None else 0.0)
        )
        self.moment_beta1 = (
            float(getattr(self.movement_plugin, "moment_beta1"))
            if hasattr(self.movement_plugin, "moment_beta1")
            else (float(moment_beta1) if moment_beta1 is not None else 0.9)
        )
        self.moment_beta2 = (
            float(getattr(self.movement_plugin, "moment_beta2"))
            if hasattr(self.movement_plugin, "moment_beta2")
            else (float(moment_beta2) if moment_beta2 is not None else 0.999)
        )
        self.moment_step_size = (
            float(getattr(self.movement_plugin, "moment_step_size"))
            if hasattr(self.movement_plugin, "moment_step_size")
            else (float(moment_step_size) if moment_step_size is not None else 1.0)
        )
        self.moment_epsilon = (
            float(getattr(self.movement_plugin, "moment_epsilon"))
            if hasattr(self.movement_plugin, "moment_epsilon")
            else (float(moment_epsilon) if moment_epsilon is not None else 1e-8)
        )

        # Validate stage/option combinations
        eval_title = self.evaluation_plugin.metadata.title
        if self.fitness_size is not None and eval_title != "Fixed Subset Evaluation":
            raise ValueError("fitness_size is only valid with evaluation='fixed_subset'")

        refine_title = self.refinement_plugin.metadata.title
        if self.refinement_epochs > 0 and refine_title == "No Refinement":
            raise ValueError("refinement_epochs > 0 is valid only with refinement='adam'")

        mv_title = self.movement_plugin.metadata.title
        if self.negative_swarm != 0.0 and mv_title in (
            "Fully Informed Particle Swarm (FIPS)",
            "Comprehensive Learning PSO (CLPSO)",
            "Bare Bones PSO",
            "Quantum PSO",
        ):
            raise ValueError(
                f"negative_swarm is unsupported for movement method '{mv_title}'"
            )

        if self.mutation_swarm != 0.0 and mv_title in (
            "Bare Bones PSO",
            "Quantum PSO",
        ):
            raise ValueError(f"mutation_swarm is unsupported for {mv_title}")

        if self.velocity_limit is not None and mv_title in (
            "Bare Bones PSO",
            "Quantum PSO",
        ):
            raise ValueError(f"velocity_limit is unsupported for {mv_title}")

        self._random_source = _RandomSource(seed=self.seed, device=self.device)
        self.generator = self._random_source.cpu_generator

        self._global_best_score: tuple[float, float, float] | None = None
        self._global_best_weights: torch.Tensor | None = None
        self.particles: list[Particle] = []

    def get_best_model(self) -> nn.Module | None:
        """
        Returns a fresh deepcopied eval-mode nn.Module on selected device with best parameters,
        or None if optimization has not been run.
        """
        if self._global_best_weights is None:
            return None
        best_model = copy.deepcopy(self.model).to(self.device)
        if self.device.type == "mps":
            best_model = best_model.to(dtype=torch.float32)
        self.codec.apply_vector(self._global_best_weights, best_model)
        best_model.eval()
        return best_model

    def get_best_score(self) -> tuple[float, float, float] | None:
        """
        Returns the best score as an immutable 3-float tuple (loss, acc, mse),
        or None if optimization has not been run.
        """
        if self._global_best_score is None:
            return None
        return (
            float(self._global_best_score[0]),
            float(self._global_best_score[1]),
            float(self._global_best_score[2]),
        )

    def get_best_state_dict(self) -> collections.OrderedDict[str, torch.Tensor] | None:
        """
        Returns a defensive CPU-cloned state dict of the best model,
        or None if optimization has not been run.
        """
        if self._global_best_weights is None:
            return None
        return self.codec.to_state_dict(self._global_best_weights, self.eval_model)
    def evaluate(
        self,
        x: torch.Tensor,
        y: torch.Tensor,
        *,
        batch_size: int | None = None,
    ) -> tuple[float, float, float]:
        """
        Evaluates the current best model weights on input data (x, y) with aggregate metric semantics.
        Validates input tensor shapes and types. Mutates no state and produces no artifacts.
        """
        if self._global_best_weights is None:
            raise RuntimeError("Optimization has not been run or best weights are unavailable")

        if not isinstance(x, torch.Tensor) or not isinstance(y, torch.Tensor):
            raise TypeError("x and y must be torch.Tensor instances")

        if x.ndim == 0 or y.ndim == 0:
            raise ValueError("x and y must have a leading dimension (ndim >= 1)")

        len_x = x.shape[0]
        len_y = y.shape[0]
        if len_x != len_y:
            raise ValueError(f"x and y leading dimensions must match: {len_x} != {len_y}")
        if len_x == 0:
            raise ValueError("x and y leading dimensions must be nonzero")

        if batch_size is not None:
            if (
                isinstance(batch_size, bool)
                or not isinstance(batch_size, int)
                or batch_size <= 0
            ):
                raise ValueError("batch_size must be a positive integer")

        dtype_model = next(self.eval_model.parameters()).dtype
        x_dev = x.to(device=self.device, dtype=dtype_model)

        if self.task == "multiclass":
            if y.ndim > 1 and y.shape[-1] > 1:
                y_dev = y.to(device=self.device, dtype=dtype_model)
            else:
                y_dev = y.to(device=self.device, dtype=torch.int64)
        elif self.task in ("binary", "regression"):
            y_dev = y.to(device=self.device, dtype=dtype_model)

        raw_score = self._evaluate_aggregate_score(
            self._global_best_weights, x_dev, y_dev, batch_size=batch_size
        )
        return _validate_score(raw_score, particle_idx=-1, iteration=-1)

    def _compute_batch_loss(
        self, out: torch.Tensor, y_batch: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Normalizes targets and computes (raw_loss, loss_tensor).
        """
        batch_size = out.shape[0]
        if self.task in ("binary", "regression"):
            if y_batch.numel() != out.numel():
                raise ValueError(
                    f"Target element count ({y_batch.numel()}) does not match model output element count ({out.numel()}) for task '{self.task}'."
                )
            target_norm = y_batch.to(device=out.device, dtype=out.dtype).reshape_as(out)
            raw_loss = self.eval_loss(out, target_norm)
        elif self.task == "multiclass":
            if out.ndim < 2:
                raise ValueError(
                    f"Multiclass model output must have rank >= 2 (logits of shape [batch_size, num_classes]), got shape {tuple(out.shape)}."
                )
            if tuple(y_batch.shape) == tuple(out.shape):
                target_norm = y_batch.to(device=out.device, dtype=out.dtype)
            elif y_batch.numel() == batch_size:
                target_norm = y_batch.to(
                    device=out.device, dtype=torch.int64
                ).reshape(-1)
            else:
                raise ValueError(
                    f"Multiclass target shape {tuple(y_batch.shape)} (numel={y_batch.numel()}) is incompatible with model output shape {tuple(out.shape)} (batch_size={batch_size})."
                )
            raw_loss = self.eval_loss(out, target_norm)
        else:
            raise ValueError(f"Unknown task: {self.task}")

        if not isinstance(raw_loss, torch.Tensor):
            raise TypeError(
                f"Loss function must return a torch.Tensor, got {type(raw_loss).__name__}"
            )

        loss_tensor = (
            raw_loss.mean() if raw_loss.numel() > 1 else raw_loss.reshape(())
        )
        return raw_loss, loss_tensor

    def _compute_batch_metrics(
        self, out: torch.Tensor, y_batch: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Computes (raw_loss, loss_tensor, acc_tensor, mse_tensor) for model output and target batch.
        """
        batch_size = out.shape[0]
        raw_loss, loss_tensor = self._compute_batch_loss(out, y_batch)

        if self.task in ("binary", "regression"):
            target_norm = y_batch.to(device=out.device, dtype=out.dtype).reshape_as(out)
            if self.task == "binary":
                if isinstance(self.eval_loss, nn.BCEWithLogitsLoss):
                    probs = torch.sigmoid(out)
                else:
                    probs = out
                preds = (probs >= 0.5).to(out.dtype)
                acc_tensor = (preds == target_norm).to(dtype=out.dtype).mean()
                mse_tensor = torch.mean((probs - target_norm) ** 2)
            else:
                acc_tensor = torch.tensor(0.0, device=out.device, dtype=out.dtype)
                mse_tensor = torch.mean((out - target_norm) ** 2)

        elif self.task == "multiclass":
            if tuple(y_batch.shape) == tuple(out.shape):
                target_norm = y_batch.to(device=out.device, dtype=out.dtype)
                target_classes = torch.argmax(target_norm, dim=-1)
                y_one_hot = target_norm
            else:
                target_norm = y_batch.to(
                    device=out.device, dtype=torch.int64
                ).reshape(-1)
                target_classes = target_norm
                num_classes = out.shape[-1]
                y_one_hot = torch.nn.functional.one_hot(
                    target_classes, num_classes=num_classes
                ).to(dtype=out.dtype)

            probs = torch.softmax(out, dim=-1)
            preds = torch.argmax(out, dim=-1)
            acc_tensor = (preds == target_classes).to(dtype=out.dtype).mean()
            mse_tensor = torch.mean((probs - y_one_hot) ** 2)
        else:
            raise ValueError(f"Unknown task: {self.task}")

        return raw_loss, loss_tensor, acc_tensor, mse_tensor

    def _evaluate_batch_tensors(
        self, x_batch: torch.Tensor, y_batch: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Evaluates batch metrics on current installed model parameters.
        Returns 0D device tensors (loss, acc, mse).
        """
        out = self.eval_model(x_batch)
        if not isinstance(out, torch.Tensor) or out.ndim < 1:
            raise ValueError("Model output must be a torch.Tensor with rank >= 1.")

        batch_size = x_batch.shape[0]
        if out.shape[0] != batch_size:
            raise ValueError(
                f"Model output leading dimension ({out.shape[0]}) does not match batch size ({batch_size})."
            )

        _, loss_t, acc_t, mse_t = self._compute_batch_metrics(out, y_batch)
        return loss_t, acc_t, mse_t

    def _evaluate_batch(
        self, position: torch.Tensor, x_batch: torch.Tensor, y_batch: torch.Tensor
    ) -> tuple[float, float, float]:
        """
        Scalar tuple evaluation wrapper.
        """
        self.codec.apply_vector(position, self.eval_model)
        with torch.inference_mode():
            t_loss, t_acc, t_mse = self._evaluate_batch_tensors(x_batch, y_batch)
        return (t_loss.item(), t_acc.item(), t_mse.item())

    def _evaluate_aggregate_score(
        self,
        position: torch.Tensor,
        x_data: torch.Tensor,
        y_data: torch.Tensor,
        batch_size: int | None = None,
    ) -> tuple[float, float, float]:
        """
        Evaluates aggregate score across all batches of x_data, y_data for a given position.
        """
        self.codec.apply_vector(position, self.eval_model)
        n_samples = x_data.shape[0]
        effective_batch_size = (
            n_samples
            if batch_size is None or batch_size >= n_samples
            else batch_size
        )

        acc_dtype = self.codec.dtype
        batch_loss_sum = torch.tensor(0.0, device=self.device, dtype=acc_dtype)
        batch_acc_sum = torch.tensor(0.0, device=self.device, dtype=acc_dtype)
        batch_mse_sum = torch.tensor(0.0, device=self.device, dtype=acc_dtype)
        total_samples = 0

        with torch.inference_mode():
            for start in range(0, n_samples, effective_batch_size):
                end = min(start + effective_batch_size, n_samples)
                batch_len = end - start
                x_b = x_data[start:end]
                y_b = y_data[start:end]
                t_loss, t_acc, t_mse = self._evaluate_batch_tensors(x_b, y_b)
                batch_loss_sum += t_loss * batch_len
                batch_acc_sum += t_acc * batch_len
                batch_mse_sum += t_mse * batch_len
                total_samples += batch_len

        agg_tensor = torch.stack([
            batch_loss_sum / total_samples,
            batch_acc_sum / total_samples,
            batch_mse_sum / total_samples,
        ])
        agg_cpu = agg_tensor.detach().cpu().tolist()
        return (float(agg_cpu[0]), float(agg_cpu[1]), float(agg_cpu[2]))

    def _optimize(
        self,
        x_fitness: torch.Tensor,
        y_fitness: torch.Tensor,
        fit_context: FitContext,
        *,
        epochs: int = 10,
        batch_size: int | None = None,
        renewal: str = "acc",
        checkpoint_interval: int | None = None,
    ) -> tuple[
        tuple[float, float, float],
        list[dict[str, Any]],
        dict[int, torch.Tensor],
    ]:
        n_samples = x_fitness.shape[0]
        effective_batch_size = (
            n_samples
            if batch_size is None or batch_size >= n_samples
            else batch_size
        )

        history: list[dict[str, Any]] = []
        checkpoint_snapshots: dict[int, torch.Tensor] = {}

        for epoch in range(epochs):
            epoch_num = epoch + 1
            if self.w_min is not None and self.w_max is not None:
                if epochs <= 2:
                    w = self.w_max
                else:
                    w_raw = self.w_max - (self.w_max - self.w_min) * (epoch / (epochs - 2))
                    w = max(self.w_min, min(self.w_max, w_raw))
            else:
                w = 0.0

            epoch_scores = torch.zeros(
                (self.n_particles, 3), device=self.device, dtype=self.codec.dtype
            )

            for i, p in enumerate(self.particles):
                self.codec.apply_vector(p.position, self.eval_model)

                batch_loss_sum = torch.tensor(
                    0.0, device=self.device, dtype=self.codec.dtype
                )
                batch_acc_sum = torch.tensor(
                    0.0, device=self.device, dtype=self.codec.dtype
                )
                batch_mse_sum = torch.tensor(
                    0.0, device=self.device, dtype=self.codec.dtype
                )
                total_samples = 0

                with torch.inference_mode():
                    for start in range(0, n_samples, effective_batch_size):
                        end = min(start + effective_batch_size, n_samples)
                        batch_len = end - start

                        x_batch = x_fitness[start:end]
                        y_batch = y_fitness[start:end]

                        t_loss, t_acc, t_mse = self._evaluate_batch_tensors(
                            x_batch, y_batch
                        )
                        batch_loss_sum += t_loss * batch_len
                        batch_acc_sum += t_acc * batch_len
                        batch_mse_sum += t_mse * batch_len
                        total_samples += batch_len

                epoch_scores[i, 0] = batch_loss_sum / total_samples
                epoch_scores[i, 1] = batch_acc_sum / total_samples
                epoch_scores[i, 2] = batch_mse_sum / total_samples

            cpu_scores = epoch_scores.detach().cpu().tolist()
            pbest_improved = [False] * self.n_particles
            pending_resets = [False] * self.n_particles

            for i, p in enumerate(self.particles):
                score = _validate_score(
                    cpu_scores[i],
                    particle_idx=i,
                    iteration=epoch_num,
                )
                if p.personal_best_score is None or _is_better_score(
                    score, p.personal_best_score, renewal
                ):
                    p.personal_best_score = score
                    p.personal_best_weights = p.position.clone()
                    pbest_improved[i] = True

                if _is_better_score(score, self._global_best_score, renewal):
                    self._global_best_score = score
                    self._global_best_weights = p.position.clone()

                iter_ctx = IterationContext(
                    epoch=epoch_num,
                    total_epochs=epochs,
                    w=w,
                    particle_idx=i,
                    is_negative=p.negative,
                    rng=self._random_source,
                    optimizer=self,
                )
                should_reset = self.convergence_plugin.on_particle_evaluated(
                    i, score, pbest_improved[i], iter_ctx
                )
                if should_reset:
                    pending_resets[i] = True

            if self._global_best_weights is None:
                raise RuntimeError(
                    "Global best weights not set before velocity calculation"
                )

            gbest_improved = _is_better_score(
                self._global_best_score, getattr(self, "_prev_gbest_score", None), renewal
            )
            self._prev_gbest_score = self._global_best_score

            epoch_iter_ctx = IterationContext(
                epoch=epoch_num,
                total_epochs=epochs,
                w=w,
                particle_idx=-1,
                is_negative=False,
                rng=self._random_source,
                optimizer=self,
            )
            stop_early = self.convergence_plugin.on_epoch_end(
                self._global_best_score, gbest_improved, epoch_iter_ctx
            )

            if epoch < epochs - 1 and not stop_early:
                # Build SwarmState snapshot BEFORE movement calculation
                swarm_positions = tuple(p.position for p in self.particles)
                swarm_velocities = tuple(p.velocity for p in self.particles)
                swarm_pbests = tuple(
                    p.personal_best_weights if p.personal_best_weights is not None else p.position
                    for p in self.particles
                )
                pbest_scores_tuple = tuple(
                    p.personal_best_score
                    if p.personal_best_score is not None
                    else (float("inf"), float("-inf"), float("inf"))
                    for p in self.particles
                )
                swarm_state = SwarmState(
                    positions=swarm_positions,
                    velocities=swarm_velocities,
                    pbest_positions=swarm_pbests,
                    pbest_scores=pbest_scores_tuple,
                    gbest_position=self._global_best_weights.detach(),
                    gbest_score=self._global_best_score,
                    pbest_improved=tuple(pbest_improved),
                )

                self.movement_plugin.on_epoch_end(swarm_state, epoch_iter_ctx)

                for i, p in enumerate(self.particles):
                    p_iter_ctx = IterationContext(
                        epoch=epoch_num,
                        total_epochs=epochs,
                        w=w,
                        particle_idx=i,
                        is_negative=p.negative,
                        rng=self._random_source,
                        optimizer=self,
                    )
                    pos_override, vel_override = self.movement_plugin.propose(
                        i, swarm_state, p_iter_ctx
                    )
                    if pos_override is not None:
                        p.velocity = torch.zeros_like(p.velocity)
                        p.position = pos_override
                    elif vel_override is not None:
                        proposed_v = vel_override
                        if (
                            self.mutation_swarm > 0.0
                            and self._random_source.bernoulli_event(
                                self.mutation_swarm
                            )
                        ):
                            proposed_v = self._random_source.uniform(
                                p.position.shape,
                                -0.2,
                                0.2,
                                device=self.device,
                                dtype=self.codec.dtype,
                            )
                            self.movement_plugin.reset_particle_state(i)
                        if self.velocity_limit is not None:
                            proposed_v = torch.clamp(
                                proposed_v, -self.velocity_limit, self.velocity_limit
                            )
                        p.velocity = proposed_v
                        p.position = p.position + p.velocity

                    p.apply_boundary_strategy(
                        self.particle_min, self.particle_max, self.boundary_strategy
                    )

                for i, p in enumerate(self.particles):
                    if pending_resets[i]:
                        base_vec = fit_context.base_vector
                        p.reset(base_vec, fit_context, self.initialization_plugin)
                        self.movement_plugin.reset_particle_state(i)
                        self.convergence_plugin.reset_particle(i)
            best = self.get_best_score()
            if best is None:
                raise RuntimeError(
                    "Optimization epoch completed without recording a best score"
                )

            history.append(
                {
                    "epoch": epoch_num,
                    "loss": float(best[0]),
                    "accuracy": float(best[1]),
                    "mse": float(best[2]),
                }
            )

            if (
                checkpoint_interval is not None
                and epoch_num % checkpoint_interval == 0
            ):
                if self._global_best_weights is not None:
                    checkpoint_snapshots[epoch_num] = (
                        self._global_best_weights.detach().cpu().clone()
                    )

            if stop_early:
                break

        best = self.get_best_score()
        if best is None:
            raise RuntimeError(
                "Optimization completed without recording a best score"
            )

        return best, history, checkpoint_snapshots

    def _refine(
        self,
        x_fitness: torch.Tensor,
        y_fitness: torch.Tensor,
        *,
        refinement_epochs: int,
        refinement_lr: float,
        batch_size: int | None,
        renewal: str,
    ) -> None:
        """
        Adam local search on fitness tensors/batching.
        """
        if refinement_epochs <= 0 or self._global_best_weights is None:
            return

        candidate_weights = self._global_best_weights.clone()
        self.codec.apply_vector(candidate_weights, self.eval_model)
        self.eval_model.eval()

        optimizer = torch.optim.Adam(self.eval_model.parameters(), lr=refinement_lr)

        n_samples = x_fitness.shape[0]
        effective_batch_size = (
            n_samples
            if batch_size is None or batch_size >= n_samples
            else batch_size
        )

        for epoch in range(refinement_epochs):
            for start in range(0, n_samples, effective_batch_size):
                end = min(start + effective_batch_size, n_samples)
                x_b = x_fitness[start:end]
                y_b = y_fitness[start:end]

                optimizer.zero_grad()
                out = self.eval_model(x_b)
                if not isinstance(out, torch.Tensor) or out.ndim < 1:
                    raise ValueError("Model output must be a torch.Tensor with rank >= 1.")
                if out.shape[0] != x_b.shape[0]:
                    raise ValueError(
                        f"Model output leading dimension ({out.shape[0]}) does not match batch size ({x_b.shape[0]})."
                    )

                _, batch_loss = self._compute_batch_loss(out, y_b)
                batch_loss.backward()
                optimizer.step()

                if self.particle_min is not None and self.particle_max is not None:
                    with torch.no_grad():
                        for p in self.eval_model.parameters():
                            p.clamp_(self.particle_min, self.particle_max)

            candidate_w = self.codec.encode(self.eval_model)
            candidate_score = self._evaluate_aggregate_score(
                candidate_w, x_fitness, y_fitness, batch_size
            )
            validated_score = _validate_score(
                candidate_score, particle_idx=-1, iteration=epoch + 1
            )
            if _is_better_score(validated_score, self._global_best_score, renewal):
                self._global_best_score = validated_score
                self._global_best_weights = candidate_w.clone()

    def _save_artifacts(
        self,
        *,
        output_dir: str | os.PathLike,
        epochs: int,
        batch_size: int | None,
        fitness_size: int | None,
        renewal: str,
        refinement_epochs: int = 0,
        refinement_lr: float = 0.001,
        validation_split: float | None,
        val_source: str | None,
        log_format: str,
        checkpoint_interval: int | None,
        save_info: bool,
        best_score: tuple[float, float, float],
        val_score: tuple[float, float, float] | None,
        val_sample_count: int | None,
        history: list[dict[str, Any]],
        checkpoint_snapshots: dict[int, torch.Tensor],
    ) -> None:
        os.makedirs(output_dir, exist_ok=True)

        best_state_dict = self.get_best_state_dict()
        if best_state_dict is None:
            raise RuntimeError(
                "Cannot save best model because optimization state is missing"
            )

        best_checkpoint = {
            "model_state_dict": best_state_dict,
            "score": best_score,
            "task": self.task,
            "device": self.device.type,
            "version": __version__,
        }
        torch.save(best_checkpoint, os.path.join(output_dir, "best_model.pt"))

        if checkpoint_snapshots:
            checkpoints_dir = os.path.join(output_dir, "checkpoints")
            os.makedirs(checkpoints_dir, exist_ok=True)
            for epoch_num, weights_vector in sorted(checkpoint_snapshots.items()):
                ckpt_path = os.path.join(checkpoints_dir, f"epoch-{epoch_num}.pt")
                ckpt_state_dict = self.codec.to_state_dict(
                    weights_vector, self.eval_model
                )
                ckpt_payload = {
                    "epoch": epoch_num,
                    "model_state_dict": ckpt_state_dict,
                    "score": (
                        history[epoch_num - 1]["loss"],
                        history[epoch_num - 1]["accuracy"],
                        history[epoch_num - 1]["mse"],
                    ),
                    "task": self.task,
                    "device": self.device.type,
                    "version": __version__,
                }
                torch.save(ckpt_payload, ckpt_path)

        if log_format == "csv":
            csv_path = os.path.join(output_dir, "history.csv")
            with open(csv_path, "w", newline="", encoding="utf-8") as f:
                writer = csv.DictWriter(
                    f, fieldnames=["epoch", "loss", "accuracy", "mse"]
                )
                writer.writeheader()
                writer.writerows(history)
        elif log_format == "tensorboard":
            from torch.utils.tensorboard import SummaryWriter

            tb_dir = os.path.join(output_dir, "tensorboard")
            writer = SummaryWriter(log_dir=tb_dir)
            for row in history:
                ep = row["epoch"]
                writer.add_scalar("loss/train", row["loss"], ep)
                writer.add_scalar("accuracy/train", row["accuracy"], ep)
                writer.add_scalar("mse/train", row["mse"], ep)
            writer.close()

        if save_info:
            loss_name = getattr(
                self.eval_loss, "__class__", type(self.eval_loss)
            ).__name__
            run_info = {
                "version": __version__,
                "task": self.task,
                "device": self.device.type,
                "loss_function": loss_name,
                "config": {
                    "method": self._method_selector,
                    "initialization": self._initialization_selector,
                    "evaluation": self._evaluation_selector,
                    "convergence": self._convergence_selector,
                    "refinement": self._refinement_selector,
                    "plugins": {
                        "movement": {
                            "title": self.movement_plugin.metadata.title,
                            "source": self.movement_plugin.metadata.source,
                            "gradient_required": self.movement_plugin.metadata.gradient_required,
                            "fidelity": self.movement_plugin.metadata.fidelity,
                            "options": self.movement_plugin.get_options(),
                        },
                        "initialization": {
                            "title": self.initialization_plugin.metadata.title,
                            "source": self.initialization_plugin.metadata.source,
                            "gradient_required": self.initialization_plugin.metadata.gradient_required,
                            "fidelity": self.initialization_plugin.metadata.fidelity,
                            "options": self.initialization_plugin.get_options(),
                        },
                        "evaluation": {
                            "title": self.evaluation_plugin.metadata.title,
                            "source": self.evaluation_plugin.metadata.source,
                            "gradient_required": self.evaluation_plugin.metadata.gradient_required,
                            "fidelity": self.evaluation_plugin.metadata.fidelity,
                            "options": self.evaluation_plugin.get_options(),
                        },
                        "convergence": {
                            "title": self.convergence_plugin.metadata.title,
                            "source": self.convergence_plugin.metadata.source,
                            "gradient_required": self.convergence_plugin.metadata.gradient_required,
                            "fidelity": self.convergence_plugin.metadata.fidelity,
                            "options": self.convergence_plugin.get_options(),
                        },
                        "refinement": {
                            "title": self.refinement_plugin.metadata.title,
                            "source": self.refinement_plugin.metadata.source,
                            "gradient_required": self.refinement_plugin.metadata.gradient_required,
                            "fidelity": self.refinement_plugin.metadata.fidelity,
                            "options": self.refinement_plugin.get_options(),
                        },
                    },
                    "n_particles": self.n_particles,
                    "c0": self.c0,
                    "c1": self.c1,
                    "w_min": self.w_min,
                    "w_max": self.w_max,
                    "negative_swarm": self.negative_swarm,
                    "mutation_swarm": self.mutation_swarm,
                    "particle_min": self.particle_min,
                    "particle_max": self.particle_max,
                    "velocity_limit_ratio": self.velocity_limit_ratio,
                    "boundary_strategy": self.boundary_strategy,
                    "initial_position_noise": self.initial_position_noise,
                    "seed": self.seed,
                    "fitness_size": fitness_size,
                    "convergence_patience": self.convergence_patience,
                    "convergence_min_delta": self.convergence_min_delta,
                    "convergence_monitor": self.convergence_monitor,
                    "moment_blend": self.moment_blend,
                    "moment_beta1": self.moment_beta1,
                    "moment_beta2": self.moment_beta2,
                    "moment_step_size": self.moment_step_size,
                    "moment_epsilon": self.moment_epsilon,
                    "epochs": epochs,
                    "batch_size": batch_size,
                    "renewal": renewal,
                    "refinement_epochs": refinement_epochs,
                    "refinement_lr": refinement_lr,
                    "validation_source": val_source,
                    "validation_split": validation_split,
                    "output_dir": str(output_dir),
                    "log_format": log_format,
                    "checkpoint_interval": checkpoint_interval,
                    "save_info": save_info,
                },
                "best_training_score": [float(x) for x in best_score],
                "validation_score": (
                    [float(x) for x in val_score] if val_score is not None else None
                ),
                "validation_source": val_source,
                "validation_sample_count": val_sample_count,
            }
            with open(os.path.join(output_dir, "run.json"), "w", encoding="utf-8") as f:
                json.dump(run_info, f, indent=2)

    def fit(
        self,
        x: torch.Tensor,
        y: torch.Tensor,
        *,
        epochs: int = 10,
        batch_size: int | None = None,
        fitness_size: int | None = None,
        renewal: str = "acc",
        refinement_epochs: int = 0,
        refinement_lr: float = 0.001,
        validation_data: tuple[torch.Tensor, torch.Tensor] | None = None,
        validation_split: float | None = None,
        output_dir: str | os.PathLike | None = None,
        log_format: Literal["none", "csv", "tensorboard"] = "none",
        checkpoint_interval: int | None = None,
        save_info: bool = False,
    ) -> tuple[float, float, float]:
        if not isinstance(x, torch.Tensor) or not isinstance(y, torch.Tensor):
            raise TypeError("x and y must be torch.Tensor instances")

        if x.ndim == 0 or y.ndim == 0:
            raise ValueError("x and y must have a leading dimension (ndim >= 1)")

        len_x = x.shape[0]
        len_y = y.shape[0]
        if len_x != len_y:
            raise ValueError(f"x and y leading dimensions must match: {len_x} != {len_y}")
        if len_x == 0:
            raise ValueError("x and y leading dimensions must be nonzero")

        if isinstance(epochs, bool) or not isinstance(epochs, int) or epochs <= 0:
            raise ValueError("epochs must be a positive integer")

        if batch_size is not None:
            if (
                isinstance(batch_size, bool)
                or not isinstance(batch_size, int)
                or batch_size <= 0
            ):
                raise ValueError("batch_size must be a positive integer")

        if renewal not in ("acc", "loss", "mse"):
            raise ValueError("renewal must be one of 'acc', 'loss', 'mse'")
        self.renewal = renewal

        if (
            isinstance(refinement_epochs, bool)
            or not isinstance(refinement_epochs, int)
            or refinement_epochs < 0
        ):
            raise ValueError("refinement_epochs must be a non-negative integer")

        if (
            isinstance(refinement_lr, bool)
            or not isinstance(refinement_lr, (int, float))
            or not math.isfinite(refinement_lr)
            or float(refinement_lr) <= 0.0
        ):
            raise ValueError("refinement_lr must be a positive finite float")

        if log_format not in ("none", "csv", "tensorboard"):
            raise ValueError("log_format must be one of 'none', 'csv', 'tensorboard'")

        if checkpoint_interval is not None:
            if (
                isinstance(checkpoint_interval, bool)
                or not isinstance(checkpoint_interval, int)
                or checkpoint_interval <= 0
            ):
                raise ValueError("checkpoint_interval must be a positive integer")

        if not isinstance(save_info, bool):
            raise ValueError("save_info must be a boolean")

        if validation_data is not None and validation_split is not None:
            raise ValueError(
                "validation_data and validation_split are mutually exclusive"
            )

        eval_title = self.evaluation_plugin.metadata.title
        if fitness_size is not None and eval_title != "Fixed Subset Evaluation":
            raise ValueError("fitness_size is only valid with evaluation='fixed_subset'")

        refine_title = self.refinement_plugin.metadata.title
        if refinement_epochs > 0 and refine_title == "No Refinement":
            raise ValueError("refinement_epochs > 0 is valid only with refinement='adam'")

        # Restore eval model parameters from constructor-time base vector and reset run state
        self.codec.apply_vector(self._base_vector.clone(), self.eval_model)
        self._global_best_score = None
        self._global_best_weights = None
        if hasattr(self, "_prev_gbest_score"):
            del self._prev_gbest_score

        val_x, val_y = None, None
        val_source = None

        if validation_data is not None:
            if not isinstance(validation_data, tuple) or len(validation_data) != 2:
                raise ValueError("validation_data must be a tuple of (val_x, val_y)")
            v_x, v_y = validation_data
            if not isinstance(v_x, torch.Tensor) or not isinstance(v_y, torch.Tensor):
                raise TypeError(
                    "validation_data elements must be torch.Tensor instances"
                )
            if v_x.ndim == 0 or v_y.ndim == 0:
                raise ValueError(
                    "validation_data elements must have a leading dimension"
                )
            len_val_x = v_x.shape[0]
            len_val_y = v_y.shape[0]
            if len_val_x != len_val_y:
                raise ValueError(
                    f"validation_data leading dimensions must match: {len_val_x} != {len_val_y}"
                )
            if len_val_x == 0:
                raise ValueError("validation_data leading dimensions must be nonzero")
            val_x, val_y = v_x, v_y
            val_source = "validation_data"

        if validation_split is not None:
            if (
                isinstance(validation_split, bool)
                or not isinstance(validation_split, (int, float))
                or not math.isfinite(validation_split)
                or not (0.0 < float(validation_split) < 1.0)
            ):
                raise ValueError(
                    "validation_split must be a finite numeric strictly between 0 and 1"
                )
            val_source = "validation_split"

        if output_dir is not None:
            if isinstance(output_dir, bool) or not isinstance(
                output_dir, (str, os.PathLike)
            ):
                raise ValueError("output_dir must be a valid path-like string or Path")

        if output_dir is None and (
            log_format != "none" or checkpoint_interval is not None or save_info
        ):
            raise ValueError(
                "output_dir is required when log_format != 'none', checkpoint_interval is set, or save_info is True"
            )

        if validation_split is not None:
            n_samples = len_x
            n_val = int(math.floor(n_samples * float(validation_split)))
            if n_val == 0 or n_val >= n_samples:
                raise ValueError(
                    "validation_split results in an empty training or validation set"
                )
            perm = self._random_source.permutation(n_samples)
            val_indices = perm[:n_val]
            train_indices = perm[n_val:]

            x_train = x[train_indices]
            val_x = x[val_indices]
            y_train = y[train_indices]
            val_y = y[val_indices]
        else:
            x_train = x
            y_train = y

        dtype_model = next(self.eval_model.parameters()).dtype
        x_train_dev = x_train.to(device=self.device, dtype=dtype_model)

        if self.task == "multiclass":
            if y_train.ndim > 1 and y_train.shape[-1] > 1:
                y_train_dev = y_train.to(device=self.device, dtype=dtype_model)
            else:
                y_train_dev = y_train.to(device=self.device, dtype=torch.int64)
        elif self.task in ("binary", "regression"):
            y_train_dev = y_train.to(device=self.device, dtype=dtype_model)

        if val_x is not None and val_y is not None:
            val_x = val_x.to(device=self.device, dtype=dtype_model)
            if self.task == "multiclass":
                if val_y.ndim > 1 and val_y.shape[-1] > 1:
                    val_y = val_y.to(device=self.device, dtype=dtype_model)
                else:
                    val_y = val_y.to(device=self.device, dtype=torch.int64)
            elif self.task in ("binary", "regression"):
                val_y = val_y.to(device=self.device, dtype=dtype_model)

        base_vector = self.codec.encode(self.eval_model)
        effective_fitness_size = fitness_size if fitness_size is not None else self.fitness_size
        effective_refine_epochs = refinement_epochs if refinement_epochs > 0 else self.refinement_epochs
        effective_refine_lr = refinement_lr if refinement_lr != 0.001 else self.refinement_lr

        fit_ctx = FitContext(
            optimizer=self,
            model=self.model,
            eval_model=self.eval_model,
            codec=self.codec,
            base_vector=base_vector,
            n_particles=self.n_particles,
            particle_min=self.particle_min,
            particle_max=self.particle_max,
            velocity_limit=self.velocity_limit,
            boundary_strategy=self.boundary_strategy,
            initial_position_noise=self.initial_position_noise,
            seed=self.seed,
            device=self.device,
            rng=self._random_source,
            task=self.task,
            x_train=x_train_dev,
            y_train=y_train_dev,
            batch_size=batch_size,
            fitness_size=effective_fitness_size,
            renewal=renewal,
            epochs=epochs,
            refinement_epochs=effective_refine_epochs,
            refinement_lr=effective_refine_lr,
            c0=self.c0,
            c1=self.c1,
            w_min=self.w_min,
            w_max=self.w_max,
            negative_swarm=self.negative_swarm,
            mutation_swarm=self.mutation_swarm,
        )

        # Prepare fit for all 5 stage plugins
        self.initialization_plugin.prepare_fit(fit_ctx)
        self.evaluation_plugin.prepare_fit(fit_ctx)
        self.movement_plugin.prepare_fit(fit_ctx)
        self.convergence_plugin.prepare_fit(fit_ctx)
        self.refinement_plugin.prepare_fit(fit_ctx)

        # Create fresh swarm for each fit
        num_negative = int(round(self.negative_swarm * self.n_particles))
        self.particles = []
        for i in range(self.n_particles):
            p = Particle(
                index=i,
                base_vector=base_vector,
                context=fit_ctx,
                init_plugin=self.initialization_plugin,
                negative=(i < num_negative),
            )
            self.particles.append(p)

        x_fitness, y_fitness = self.evaluation_plugin.get_fitness_data(
            x_train_dev, y_train_dev, fit_ctx
        )

        best_score, history, checkpoint_snapshots = self._optimize(
            x_fitness,
            y_fitness,
            fit_ctx,
            epochs=epochs,
            batch_size=batch_size,
            renewal=renewal,
            checkpoint_interval=checkpoint_interval,
        )

        if effective_refine_epochs > 0:
            refined_w, refined_s = self.refinement_plugin.refine(
                self._global_best_weights
                if self._global_best_weights is not None
                else base_vector,
                best_score,
                self._evaluate_batch,
                fit_ctx,
            )
            best_score = refined_s

        val_score = None
        val_sample_count = None
        if val_x is not None and val_y is not None:
            if self._global_best_weights is None:
                raise RuntimeError("Best model not available after optimization")
            raw_val = self._evaluate_aggregate_score(
                self._global_best_weights, val_x, val_y, batch_size=batch_size
            )
            val_score = _validate_score(raw_val, particle_idx=-1, iteration=-1)
            val_sample_count = val_x.shape[0]

        if output_dir is not None:
            self._save_artifacts(
                output_dir=output_dir,
                epochs=epochs,
                batch_size=batch_size,
                fitness_size=effective_fitness_size,
                renewal=renewal,
                refinement_epochs=effective_refine_epochs,
                refinement_lr=effective_refine_lr,
                validation_split=validation_split,
                val_source=val_source,
                log_format=log_format,
                checkpoint_interval=checkpoint_interval,
                save_info=save_info,
                best_score=best_score,
                val_score=val_score,
                val_sample_count=val_sample_count,
                history=history,
                checkpoint_snapshots=checkpoint_snapshots,
            )

        return best_score