from torch.optim import Optimizer
from transformers.trainer import Trainer as HfTrainer
from typing import TYPE_CHECKING

try:
    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
except ImportError:
    from torch.optim.lr_scheduler import LRScheduler

if TYPE_CHECKING:
    from swift.trainers import Trainer, TrainingArguments


class OptimizerCallback:
    """
    Callback for creating and managing optimizer and learning rate scheduler.

    This callback provides hooks for customizing the creation of optimizers and
    learning rate schedulers during the training process. It delegates to the
    trainer's methods by default but can be subclassed to implement custom
    optimization strategies.

    Args:
        args (TrainingArguments): The training arguments containing hyperparameters
            and configuration settings.
        trainer (Trainer): The trainer instance that will use this callback.
    """

    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
        self.args = args
        self.trainer = trainer

    def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
        """
        Create both optimizer and learning rate scheduler for training.

        This method initializes the optimizer and scheduler by calling their
        respective creation methods and assigns them to the trainer instance.

        Args:
            num_training_steps (int): The total number of training steps, used
                for scheduler configuration (e.g., warmup steps, decay schedule).

        Returns:
            None: The optimizer and scheduler are set directly on the trainer.
        """
        trainer = self.trainer
        trainer.optimizer = self.create_optimizer()
        trainer.scheduler = self.create_scheduler(num_training_steps, trainer.optimizer)

    def create_optimizer(self, model=None) -> Optimizer:
        kwargs = {} if model is None else {'model': model}
        return HfTrainer.create_optimizer(self.trainer, **kwargs)

    def create_scheduler(self, num_training_steps: int, optimizer: Optimizer) -> LRScheduler:
        return HfTrainer.create_scheduler(self.trainer, num_training_steps, optimizer)
