"""Quasi-hyperbolic scheduler that switches to new parameters at specified times."""

from collections import defaultdict
from logging import WARNING
from typing import Any, cast

import torch
from composer.core import Callback, State
from composer.loggers import Logger
from composer.optim.scheduler import _convert_time  # noqa: PLC2701
from flwr.common import log as flwr_log
from llmfoundry.callbacks import callbacks

from mt_dao.masks_utils import ModelStateNames
from mt_dao.pytorch_utils.model_parameters_getters import get_tensors_dict
from mt_dao.pytorch_utils.optimizer_states_setters import set_optimizers_states


@callbacks.register_class("quasi_hyper_switch_callback")
class QuasiHyperBolicSwitchCallback(Callback):
    r"""Functional quasi-hyperbolic scheduler.

    Switches to quasi-hyperbolic formulation at specified times.
    """

    def __init__(
        self,
        new_vs: tuple[float, ...] | None = None,
        new_betas: tuple[float, ...] | None = None,
        t_switch: str | Any = "1dur",  # noqa: ANN401
        momenta_to_reset: tuple[str, ...] = ("EXP_AVG",),
    ) -> None:
        """Initialize the quasi-hyperbolic scheduler.

        Parameters
        ----------
        new_vs : tuple of float
            The new v values to switch to.
        new_betas : tuple of float
            The new beta values to switch to.
        t_switch : str or Any, optional
            The time to switch to the new parameters, by default "1dur".
        momenta_to_reset : tuple of str, optional
            The momenta to reset, by default ("EXP_AVG",).

        Raises
        ------
        ValueError
            If parameters are included in momenta_to_reset.

        """
        self.t_switch = t_switch
        self.new_vs = new_vs
        self.new_betas = new_betas
        self.momenta_to_reset = tuple(ModelStateNames[v] for v in momenta_to_reset)
        if ModelStateNames.PARAMETERS in self.momenta_to_reset:
            msg = "Cannot reset parameters."
            raise ValueError(msg)
        flwr_log(WARNING, "Instantiated qh scheduler")

    def batch_end(self, state: State, logger: Logger) -> None:  # noqa: ARG002, C901
        """Update the quasi-hyperbolic parameters based on the current time.

        Parameters
        ----------
        state : State
            The current state of training.
        logger : Logger
            The logger to use for logging.

        """
        # Extract current base values directly from the optimizer parameter groups.
        optimizer = state.optimizers[0]
        groups = optimizer.param_groups

        # Extract the current step from the optimizer

        # Convert durations using Composer's _convert_time.
        converted_t_switch = _convert_time(self.t_switch, state)
        converted_t_switch_value = float(converted_t_switch.value)

        # Get current time from state.timestamp using the unit of t_v1.
        current_time = state.timestamp.get(converted_t_switch.unit)
        t = float(current_time.value)

        tensors_dict = get_tensors_dict(
            state.model,
            only_requires_grad=True,
            sort_dict=True,
        )

        if t == converted_t_switch_value:
            flwr_log(
                WARNING,
                f"Scheduler time {t} , target time{converted_t_switch_value}",
            )
            opt_states_dict: dict[str, dict[str, torch.Tensor]] = cast(
                "dict[str, dict[str, torch.Tensor]]",
                defaultdict(lambda: defaultdict(dict)),
            )

            for name, val in tensors_dict.items():
                for momenta in self.momenta_to_reset:
                    opt_states_dict[momenta][
                        name.replace("model.", "").replace(
                            "module.",
                            "",
                        )
                    ] = torch.zeros_like(val)

            set_optimizers_states(
                model=state.model,
                optimizer=optimizer,
                step=None,
                opt_state_dict=opt_states_dict,
            )

            # Set constants
            for group in groups:
                if self.new_vs is not None:
                    if isinstance(group["vs"], tuple):
                        group["vs"] = self.new_vs
                    elif isinstance(group["vs"], torch.Tensor):
                        group["vs"].data.copy_(
                            torch.tensor(
                                self.new_vs,
                                device=group["vs"].device,
                                dtype=group["vs"].dtype,
                            ),
                        )
                    else:
                        group["vs"] = self.new_vs
                if self.new_betas is not None:
                    if isinstance(group["betas"], tuple):
                        group["betas"] = self.new_betas
                    elif isinstance(group["betas"], torch.Tensor):
                        group["betas"].data.copy_(
                            torch.tensor(
                                self.new_betas,
                                device=group["betas"].device,
                                dtype=group["betas"].dtype,
                            ),
                        )
                    else:
                        group["betas"] = self.new_betas
