"""Optimizer factory — build a ``torch.optim.Optimizer`` from a :class:`TrainConfig`.

Extracted from the duplicated ``setup_optimizer()`` functions in
``ml_pytorch_vision_classification/train.py`` and ``r2plus1d/train.py``.

Known optimizer names (SGD, Adam, AdamW, Adamax, RMSprop) get sensible
defaults from ``TrainConfig`` fields.  Any name not in the known set is
looked up dynamically from ``torch.optim``, passing only
``config.optimizer_kwargs``.

``config.optimizer_kwargs`` is always merged last, so user-provided
values override the config-derived defaults.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any


if TYPE_CHECKING:
    from collections.abc import Iterator

    import torch.nn as nn
    import torch.optim as optim

    from matrice_models.config.base import TrainConfig

logger = logging.getLogger(__name__)

_MOMENTUM_OPTIMIZERS = frozenset({"sgd", "rmsprop"})


def build_optimizer(
    model: nn.Module,
    config: TrainConfig,
    *,
    parameter_groups: Iterator[Any] | None = None,
) -> optim.Optimizer:
    """Construct the optimizer specified in *config*.

    For known names (SGD, Adam, AdamW, Adamax, RMSprop) a base kwargs
    dict is built from ``config`` fields (``learning_rate``,
    ``weight_decay``, ``momentum``).  ``config.optimizer_kwargs`` is
    merged on top, so explicit user values always win.

    For any other name the class is looked up dynamically from
    ``torch.optim`` and only ``config.optimizer_kwargs`` (plus ``lr``)
    is forwarded.

    Args:
        model: The model whose parameters will be optimized.
        config: Training configuration.
        parameter_groups: Custom parameter iterator.  When ``None``
            (the default), ``model.parameters()`` is used.

    Returns:
        A configured :class:`torch.optim.Optimizer`.

    Raises:
        ValueError: If the optimizer name cannot be resolved.
    """
    params = parameter_groups if parameter_groups is not None else model.parameters()
    name = config.optimizer.lower()
    extra = config.optimizer_kwargs

    base_kwargs = _base_kwargs_for(name, config)
    merged = {**base_kwargs, **extra}

    cls = _resolve_class(config.optimizer)
    logger.debug("Building optimizer %s with kwargs %s", cls.__name__, merged)
    return cls(params, **merged)


def _base_kwargs_for(name: str, config: TrainConfig) -> dict[str, Any]:
    """Build default kwargs from well-known config fields."""
    kwargs: dict[str, Any] = {
        "lr": config.learning_rate,
        "weight_decay": config.weight_decay,
    }
    if name in _MOMENTUM_OPTIMIZERS or name.startswith("sgd"):
        kwargs["momentum"] = config.momentum
    if name.startswith("sgd"):
        kwargs["nesterov"] = "nesterov" in name
    return kwargs


def _resolve_class(name: str) -> type[optim.Optimizer]:
    """Resolve an optimizer class by name from ``torch.optim``."""
    import torch.optim as optim_module

    # Case-insensitive lookup against torch.optim attributes
    name_lower = name.lower()
    for attr in dir(optim_module):
        if attr.lower() == name_lower:
            cls = getattr(optim_module, attr)
            if isinstance(cls, type) and issubclass(cls, optim_module.Optimizer):
                return cls

    msg = (
        f"Optimizer '{name}' not found in torch.optim. "
        f"Available: {[c for c in dir(optim_module) if not c.startswith('_')]}"
    )
    raise ValueError(msg)
