"""LR scheduler factory — build a scheduler from a :class:`TrainConfig`.

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

Known scheduler names (StepLR, CosineAnnealingLR, ExponentialLR) get
sensible defaults from ``TrainConfig`` fields.  Any name not in the
known set is looked up dynamically from ``torch.optim.lr_scheduler``,
passing only ``config.scheduler_kwargs``.

``config.scheduler_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:
    import torch.optim as optim
    from torch.optim.lr_scheduler import LRScheduler

    from matrice_models.config.base import TrainConfig

logger = logging.getLogger(__name__)


def build_scheduler(
    optimizer: optim.Optimizer,
    config: TrainConfig,
) -> LRScheduler | None:
    """Construct the LR scheduler specified in *config*.

    Returns ``None`` when ``config.scheduler`` is unset, giving a
    constant learning rate.

    For known names (StepLR, CosineAnnealingLR, ExponentialLR) a base
    kwargs dict is built from ``config`` fields.
    ``config.scheduler_kwargs`` is merged on top, so explicit user
    values always win.

    For any other name the class is looked up dynamically from
    ``torch.optim.lr_scheduler`` and only ``config.scheduler_kwargs``
    is forwarded.

    Args:
        optimizer: The optimizer whose LR will be scheduled.
        config: Training configuration.

    Returns:
        A configured scheduler, or ``None`` for constant LR.

    Raises:
        ValueError: If the scheduler name cannot be resolved.
    """
    name = (config.scheduler or "").strip()
    if not name:
        return None

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

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


def _base_kwargs_for(name: str, config: TrainConfig) -> dict[str, Any]:
    """Build default kwargs from well-known config fields."""
    if name == "steplr":
        return {"step_size": config.lr_step_size, "gamma": config.lr_gamma}
    if name == "cosineannealinglr":
        return {"T_max": config.epochs, "eta_min": config.lr_min}
    if name == "exponentiallr":
        return {"gamma": config.lr_gamma}
    return {}


def _resolve_class(name: str) -> type[LRScheduler]:
    """Resolve a scheduler class by name from ``torch.optim.lr_scheduler``."""
    from torch.optim import lr_scheduler as sched_module
    from torch.optim.lr_scheduler import LRScheduler as _Base

    name_lower = name.lower()
    for attr in dir(sched_module):
        if attr.lower() == name_lower:
            cls = getattr(sched_module, attr)
            if isinstance(cls, type) and issubclass(cls, _Base):
                return cls

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