"""
Learning Rate Schedulers
========================

Provides a comprehensive set of learning rate schedulers with a factory
function for easy selection via CLI arguments.

Supported Schedulers:
    - plateau: ReduceLROnPlateau (default, adaptive)
    - cosine: CosineAnnealingLR
    - cosine_restarts: CosineAnnealingWarmRestarts
    - onecycle: OneCycleLR
    - step: StepLR
    - multistep: MultiStepLR
    - exponential: ExponentialLR
    - linear_warmup: LinearLR (warmup phase)

Author: Ductho Le (ductho.le@outlook.com)
Version: 1.0.0
"""

import torch.optim as optim
from torch.optim.lr_scheduler import (
    CosineAnnealingLR,
    CosineAnnealingWarmRestarts,
    ExponentialLR,
    LinearLR,
    LRScheduler,
    MultiStepLR,
    OneCycleLR,
    ReduceLROnPlateau,
    SequentialLR,
    StepLR,
)


# ==============================================================================
# SCHEDULER REGISTRY
# ==============================================================================
def list_schedulers() -> list[str]:
    """
    Return list of available scheduler names.

    Returns:
        List of registered scheduler names
    """
    return [
        "plateau",
        "cosine",
        "cosine_restarts",
        "onecycle",
        "step",
        "multistep",
        "exponential",
        "linear_warmup",
    ]


def get_scheduler(
    name: str,
    optimizer: optim.Optimizer,
    # Common parameters
    epochs: int = 100,
    steps_per_epoch: int | None = None,
    min_lr: float = 1e-6,
    # ReduceLROnPlateau parameters
    patience: int = 10,
    factor: float = 0.5,
    # Cosine parameters
    T_max: int | None = None,
    T_0: int = 10,
    T_mult: int = 2,
    # OneCycleLR parameters
    max_lr: float | None = None,
    pct_start: float = 0.3,
    # Step/MultiStep parameters
    step_size: int = 30,
    milestones: list[int] | None = None,
    gamma: float = 0.1,
    # Linear warmup parameters
    warmup_epochs: int = 5,
    start_factor: float = 0.1,
    **kwargs,
) -> LRScheduler:
    """
    Factory function to create learning rate scheduler by name.

    Args:
        name: Scheduler name (see list_schedulers())
        optimizer: Optimizer instance to schedule
        epochs: Total training epochs (for cosine, onecycle)
        steps_per_epoch: Steps per epoch (required for onecycle)
        min_lr: Minimum learning rate (eta_min for cosine)
        patience: Patience for ReduceLROnPlateau
        factor: Reduction factor for plateau/step
        T_max: Period for CosineAnnealingLR (default: epochs)
        T_0: Initial period for CosineAnnealingWarmRestarts
        T_mult: Period multiplier for warm restarts
        max_lr: Maximum LR for OneCycleLR (default: optimizer's initial LR)
        pct_start: Percentage of cycle spent increasing LR (OneCycleLR)
        step_size: Period for StepLR
        milestones: Epochs to decay LR for MultiStepLR
        gamma: Decay factor for step/multistep/exponential
        warmup_epochs: Number of warmup epochs for linear_warmup
        start_factor: Starting LR factor for warmup (LR * start_factor)
        **kwargs: Additional arguments passed to scheduler

    Returns:
        Instantiated learning rate scheduler

    Raises:
        ValueError: If scheduler name is not recognized

    Example:
        >>> scheduler = get_scheduler("plateau", optimizer, patience=15)
        >>> scheduler = get_scheduler("cosine", optimizer, epochs=100)
        >>> scheduler = get_scheduler(
        ...     "onecycle", optimizer, epochs=100, steps_per_epoch=1000, max_lr=1e-3
        ... )
    """
    name_lower = name.lower().replace("-", "_")

    # Get initial LR from optimizer
    base_lr = optimizer.param_groups[0]["lr"]

    if name_lower == "plateau":
        return ReduceLROnPlateau(
            optimizer,
            mode="min",
            factor=factor,
            patience=patience,
            min_lr=min_lr,
            **kwargs,
        )

    elif name_lower == "cosine":
        return CosineAnnealingLR(
            optimizer,
            T_max=T_max if T_max is not None else epochs,
            eta_min=min_lr,
            **kwargs,
        )

    elif name_lower == "cosine_restarts":
        return CosineAnnealingWarmRestarts(
            optimizer, T_0=T_0, T_mult=T_mult, eta_min=min_lr, **kwargs
        )

    elif name_lower == "onecycle":
        if steps_per_epoch is None:
            raise ValueError(
                "OneCycleLR requires 'steps_per_epoch'. "
                "Pass len(train_dataloader) as steps_per_epoch."
            )
        return OneCycleLR(
            optimizer,
            max_lr=max_lr if max_lr is not None else base_lr,
            epochs=epochs,
            steps_per_epoch=steps_per_epoch,
            pct_start=pct_start,
            **kwargs,
        )

    elif name_lower == "step":
        return StepLR(optimizer, step_size=step_size, gamma=gamma, **kwargs)

    elif name_lower == "multistep":
        if milestones is None:
            # Default milestones at 30%, 60%, 90% of epochs
            milestones = [int(epochs * 0.3), int(epochs * 0.6), int(epochs * 0.9)]
        return MultiStepLR(optimizer, milestones=milestones, gamma=gamma, **kwargs)

    elif name_lower == "exponential":
        return ExponentialLR(optimizer, gamma=gamma, **kwargs)

    elif name_lower == "linear_warmup":
        return LinearLR(
            optimizer,
            start_factor=start_factor,
            end_factor=1.0,
            total_iters=warmup_epochs,
            **kwargs,
        )

    else:
        available = ", ".join(list_schedulers())
        raise ValueError(f"Unknown scheduler: '{name}'. Available options: {available}")


def get_scheduler_with_warmup(
    name: str,
    optimizer: optim.Optimizer,
    warmup_epochs: int = 5,
    start_factor: float = 0.1,
    **kwargs,
) -> LRScheduler:
    """
    Create a scheduler with linear warmup phase.

    Combines LinearLR warmup with any other scheduler using SequentialLR.

    Args:
        name: Main scheduler name (after warmup)
        optimizer: Optimizer instance
        warmup_epochs: Number of warmup epochs
        start_factor: Starting LR factor for warmup
        **kwargs: Arguments for main scheduler (see get_scheduler)

    Returns:
        SequentialLR combining warmup and main scheduler

    Example:
        >>> scheduler = get_scheduler_with_warmup(
        ...     "cosine", optimizer, warmup_epochs=5, epochs=100
        ... )
    """
    # Create warmup scheduler
    warmup_scheduler = LinearLR(
        optimizer,
        start_factor=start_factor,
        end_factor=1.0,
        total_iters=warmup_epochs,
    )

    # Create main scheduler
    main_scheduler = get_scheduler(name, optimizer, **kwargs)

    # Combine with SequentialLR
    return SequentialLR(
        optimizer,
        schedulers=[warmup_scheduler, main_scheduler],
        milestones=[warmup_epochs],
    )


def is_epoch_based(name: str) -> bool:
    """
    Check if scheduler should be stepped per epoch (True) or per batch (False).

    Args:
        name: Scheduler name

    Returns:
        True if scheduler should step per epoch, False for per batch
    """
    name_lower = name.lower().replace("-", "_")

    # OneCycleLR steps per batch, all others step per epoch
    per_batch_schedulers = {"onecycle"}

    return name_lower not in per_batch_schedulers
