__all__ = [
    'LRScheduleCallback',
    'LRScaleCallback',
]

from typing import Any, Mapping, cast

import torch

from ...base import CallbackRegistry, Config, LrSchedulerRegistry
from ...utils import get_world_size
from ..runners import Trainer
from .base import BaseCallback
from .interval import IntervalMixin

Memo = dict[str, Any]


@CallbackRegistry.register()
class LRScheduleCallback(IntervalMixin, BaseCallback):

    def __init__(
        self,
        *args,
        lr_scheduler: Config,
        interval: int = 1,
        **kwargs,
    ) -> None:
        super().__init__(*args, interval=interval, **kwargs)
        assert isinstance(self._runner, Trainer)
        self._lr_scheduler_config = lr_scheduler

    def connect(self) -> None:
        super().connect()
        self._build_lr_scheduler()

    def _build_lr_scheduler(self) -> None:
        runner = cast(Trainer, self._runner)
        self._lr_scheduler: torch.optim.lr_scheduler.LRScheduler = \
            LrSchedulerRegistry.build(
                self._lr_scheduler_config,
                optimizer=runner.optimizer,
            )

    def after_run_iter(self, batch, memo: Memo) -> None:
        super().after_run_iter(batch, memo)
        if self._should_run_iter():
            self._lr_scheduler.step()
        if 'log' in memo:
            memo['log']['lr'] = [
                f'{lr:.3e}' for lr in self._lr_scheduler.get_last_lr()
            ]

    def after_run_epoch(self, epoch_memo: Memo, memo: Memo) -> None:
        super().after_run_epoch(epoch_memo, memo)
        if self._should_run_epoch():
            self._lr_scheduler.step()

    def load_state_dict(
        self,
        state_dict: Mapping[str, Any],
        *args,
        **kwargs,
    ) -> None:
        super().load_state_dict(state_dict, *args, **kwargs)
        self._lr_scheduler.load_state_dict(state_dict['lr_scheduler'])

    def state_dict(self, *args, **kwargs) -> dict[str, Any]:
        state_dict = super().state_dict(*args, **kwargs)
        state_dict['lr_scheduler'] = self._lr_scheduler.state_dict()
        return state_dict


@CallbackRegistry.register()
class LRScaleCallback(BaseCallback):

    def __init__(self, *args, lr_scaler: Config, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        assert isinstance(self._runner, Trainer)
        self._lr_scaler_config = lr_scaler

    def _scale_lr(self, config: Config) -> None:
        runner = cast(Trainer, self._runner)
        base_batch_size = config.base_batch_size
        batch_size = get_world_size() * runner.dataloader.batch_size
        lr_scaler = batch_size / base_batch_size
        if 'lr' in runner.optimizer.defaults:
            runner.optimizer.defaults['lr'] *= lr_scaler
        for param_group in runner.optimizer.param_groups:
            if 'lr' in param_group:
                param_group['lr'] *= lr_scaler
        runner.logger.info(
            f"{base_batch_size=} {batch_size=} {lr_scaler=:.3f}"
        )

    def connect(self) -> None:
        super().connect()
        self._scale_lr(self._lr_scaler_config)
