"""
Learning rate schedulers for gradient ascent optimization.

Provides various LR scheduling strategies for reward-guided gradient ascent,
including cosine annealing, linear decay, and custom schedules.
"""

import math
from typing import Optional, Literal


class LRScheduler:
    """Base class for learning rate schedulers."""
    
    def __init__(self, initial_lr: float, num_steps: int):
        """
        Initialize LR scheduler.
        
        Args:
            initial_lr: Initial learning rate
            num_steps: Total number of optimization steps
        """
        self.initial_lr = initial_lr
        self.num_steps = num_steps
        self.current_step = 0
    
    def get_lr(self) -> float:
        """Get current learning rate."""
        raise NotImplementedError
    
    def step(self):
        """Update scheduler state after a step."""
        self.current_step += 1
    
    def reset(self):
        """Reset scheduler state."""
        self.current_step = 0


class ConstantLR(LRScheduler):
    """Constant learning rate (no scheduling)."""
    
    def get_lr(self) -> float:
        return self.initial_lr


class LinearLR(LRScheduler):
    """Linear learning rate decay."""
    
    def __init__(
        self,
        initial_lr: float,
        num_steps: int,
        end_lr: float = 0.0,
        start_step: int = 0,
    ):
        """
        Initialize linear LR scheduler.
        
        Args:
            initial_lr: Starting learning rate
            num_steps: Total number of steps
            end_lr: Ending learning rate (default: 0.0)
            start_step: Step to begin decay (default: 0)
        """
        super().__init__(initial_lr, num_steps)
        self.end_lr = end_lr
        self.start_step = start_step
    
    def get_lr(self) -> float:
        if self.current_step < self.start_step:
            return self.initial_lr
        
        progress = (self.current_step - self.start_step) / (self.num_steps - self.start_step)
        progress = min(1.0, progress)
        
        return self.initial_lr + (self.end_lr - self.initial_lr) * progress


class CosineLR(LRScheduler):
    """Cosine annealing learning rate schedule."""
    
    def __init__(
        self,
        initial_lr: float,
        num_steps: int,
        min_lr: float = 0.0,
        warmup_steps: int = 0,
    ):
        """
        Initialize cosine LR scheduler.
        
        Args:
            initial_lr: Maximum learning rate
            num_steps: Total number of steps
            min_lr: Minimum learning rate (default: 0.0)
            warmup_steps: Number of linear warmup steps (default: 0)
        """
        super().__init__(initial_lr, num_steps)
        self.min_lr = min_lr
        self.warmup_steps = warmup_steps
    
    def get_lr(self) -> float:
        if self.current_step < self.warmup_steps:
            # Linear warmup
            return self.initial_lr * (self.current_step / self.warmup_steps)
        
        # Cosine annealing
        progress = (self.current_step - self.warmup_steps) / (self.num_steps - self.warmup_steps)
        progress = min(1.0, progress)
        
        cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
        return self.min_lr + (self.initial_lr - self.min_lr) * cosine_decay


class ExponentialLR(LRScheduler):
    """Exponential learning rate decay."""
    
    def __init__(
        self,
        initial_lr: float,
        num_steps: int,
        gamma: float = 0.95,
    ):
        """
        Initialize exponential LR scheduler.
        
        Args:
            initial_lr: Starting learning rate
            num_steps: Total number of steps
            gamma: Multiplicative decay factor per step
        """
        super().__init__(initial_lr, num_steps)
        self.gamma = gamma
    
    def get_lr(self) -> float:
        return self.initial_lr * (self.gamma ** self.current_step)


class StepLR(LRScheduler):
    """Step-wise learning rate decay."""
    
    def __init__(
        self,
        initial_lr: float,
        num_steps: int,
        step_size: int,
        gamma: float = 0.1,
    ):
        """
        Initialize step LR scheduler.
        
        Args:
            initial_lr: Starting learning rate
            num_steps: Total number of steps
            step_size: Number of steps between each decay
            gamma: Multiplicative decay factor
        """
        super().__init__(initial_lr, num_steps)
        self.step_size = step_size
        self.gamma = gamma
    
    def get_lr(self) -> float:
        num_decays = self.current_step // self.step_size
        return self.initial_lr * (self.gamma ** num_decays)


def create_lr_scheduler(
    scheduler_type: Literal["constant", "linear", "cosine", "exponential", "step"],
    initial_lr: float,
    num_steps: int,
    **kwargs
) -> LRScheduler:
    """
    Factory function to create learning rate schedulers.
    
    Args:
        scheduler_type: Type of scheduler ("constant", "linear", "cosine", "exponential", "step")
        initial_lr: Initial learning rate
        num_steps: Total number of optimization steps
        **kwargs: Additional scheduler-specific arguments
            For linear: end_lr, start_step
            For cosine: min_lr, warmup_steps
            For exponential: gamma
            For step: step_size, gamma
    
    Returns:
        LRScheduler instance
    
    Examples:
        # Constant LR
        scheduler = create_lr_scheduler("constant", initial_lr=0.1, num_steps=100)
        
        # Linear decay
        scheduler = create_lr_scheduler("linear", initial_lr=0.1, num_steps=100, end_lr=0.01)
        
        # Cosine annealing with warmup
        scheduler = create_lr_scheduler("cosine", initial_lr=0.1, num_steps=100, 
                                       min_lr=0.001, warmup_steps=10)
    """
    if scheduler_type == "constant":
        return ConstantLR(initial_lr, num_steps)
    
    elif scheduler_type == "linear":
        return LinearLR(
            initial_lr, num_steps,
            end_lr=kwargs.get("end_lr", 0.0),
            start_step=kwargs.get("start_step", 0),
        )
    
    elif scheduler_type == "cosine":
        return CosineLR(
            initial_lr, num_steps,
            min_lr=kwargs.get("min_lr", 0.0),
            warmup_steps=kwargs.get("warmup_steps", 0),
        )
    
    elif scheduler_type == "exponential":
        return ExponentialLR(
            initial_lr, num_steps,
            gamma=kwargs.get("gamma", 0.95),
        )
    
    elif scheduler_type == "step":
        return StepLR(
            initial_lr, num_steps,
            step_size=kwargs.get("step_size", 10),
            gamma=kwargs.get("gamma", 0.1),
        )
    
    else:
        raise ValueError(f"Unknown scheduler type: {scheduler_type}. "
                        f"Choose from: constant, linear, cosine, exponential, step")