import math
import typing

from torch.optim.lr_scheduler import LambdaLR, LinearLR, SequentialLR, _LRScheduler

if typing.TYPE_CHECKING:
    from torch.optim import Optimizer


def get_warmup_scheduler(
    optimizer: 'Optimizer',
    warmup_steps: int,
    main_scheduler: '_LRScheduler',
    warmup_start_lr: float = 0.0,
    warmup_mode: typing.Literal['linear', 'cosine', 'quadratic'] = 'linear',
    last_epoch: int = -1,
) -> _LRScheduler | SequentialLR:
    """
    Create a warmup scheduler using PyTorch's native SequentialLR.

    This function creates a warmup phase followed by the main scheduler using
    PyTorch's built-in schedulers for better compatibility and maintainability.

    Parameters
    ----------
    - `optimizer` : `Optimizer`
        - The optimizer to schedule
    - `warmup_steps` : `int`
        - Number of steps for the warmup phase
    - `main_scheduler` : `_LRScheduler`
        - The main scheduler to use after warmup
    - `warmup_start_lr` : `float`, optional
        - Initial learning rate for warmup, by default 0.0
    - `warmup_mode` : `Literal['linear', 'cosine', 'quadratic']`, optional
        - Warmup schedule type, by default 'linear'
    - `last_epoch` : `int`, optional
        - The index of last epoch, by default -1

    Returns
    ----------
    - `SequentialLR`
        - Configured SequentialLR combining warmup and main scheduler

    Examples
    --------
    ```python
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    main_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
    scheduler = get_warmup_scheduler(
        optimizer,
        warmup_steps=10,
        main_scheduler=main_scheduler,
        warmup_mode='linear'
    )
    ```
    """
    if warmup_steps <= 0:
        return main_scheduler

    # Get base learning rates for all parameter groups
    base_lrs = [group['lr'] for group in optimizer.param_groups]

    # Create warmup scheduler based on mode
    if warmup_mode == 'linear':
        # Use LinearLR for linear warmup
        # start_factor = warmup_start_lr / base_lr, end_factor = 1.0
        start_factors = [
            warmup_start_lr / base_lr if base_lr > 0 else 0.0 for base_lr in base_lrs
        ]
        # Handle case where all parameter groups have the same ratio
        start_factor = start_factors[0] if len(set(start_factors)) == 1 else start_factors

        warmup_scheduler = LinearLR(
            optimizer,
            start_factor=typing.cast(float, start_factor),
            end_factor=1.0,
            total_iters=warmup_steps,
            last_epoch=last_epoch,
        )

    elif warmup_mode == 'cosine':
        # Use LambdaLR for cosine warmup
        def make_cosine_warmup_lambda(base_lr, warmup_start_lr):
            def cosine_warmup_lambda(epoch):
                if epoch >= warmup_steps:
                    return 1.0
                progress = epoch / warmup_steps
                if base_lr > 0:
                    start_factor = warmup_start_lr / base_lr
                    return (
                        start_factor
                        + (1.0 - start_factor) * (1 - math.cos(math.pi * progress)) / 2
                    )
                else:
                    return 0.0

            return cosine_warmup_lambda

        # Create lambda functions for each parameter group
        lambda_funcs = [
            make_cosine_warmup_lambda(base_lr, warmup_start_lr) for base_lr in base_lrs
        ]

        warmup_scheduler = LambdaLR(
            optimizer, lr_lambda=lambda_funcs, last_epoch=last_epoch
        )

    elif warmup_mode == 'quadratic':
        # Use LambdaLR for quadratic (quadratic) warmup
        def make_quadratic_warmup_lambda(base_lr, warmup_start_lr):
            def quadratic_warmup_lambda(epoch):
                if epoch >= warmup_steps:
                    return 1.0
                progress = epoch / warmup_steps
                if base_lr > 0:
                    start_factor = warmup_start_lr / base_lr
                    return start_factor + (1.0 - start_factor) * (progress**2)
                else:
                    return 0.0

            return quadratic_warmup_lambda

        # Create lambda functions for each parameter group
        lambda_funcs = [
            make_quadratic_warmup_lambda(base_lr, warmup_start_lr) for base_lr in base_lrs
        ]

        warmup_scheduler = LambdaLR(
            optimizer, lr_lambda=lambda_funcs, last_epoch=last_epoch
        )

    else:
        raise ValueError(f'Unknown warmup_mode: {warmup_mode}')

    # Create SequentialLR combining warmup and main scheduler
    return SequentialLR(
        optimizer,
        schedulers=[warmup_scheduler, main_scheduler],
        milestones=[warmup_steps],
        last_epoch=last_epoch,
    )


def get_learning_rate(optimizer: 'Optimizer') -> typing.List[float]:
    """Get current learning rates from the optimizer.

    Parameters
    ----------
    - `optimizer` : `Optimizer`
        - The optimizer to query

    Returns
    ----------
    - `List[float]`
        - List of current learning rates for each parameter group
    """
    return [group['lr'] for group in optimizer.param_groups]
