# Copyright 2022 MosaicML Composer authors
# SPDX-License-Identifier: Apache-2.0

"""Monitor vs during training."""

from composer.core import Callback, State
from composer.loggers import Logger
from llmfoundry.callbacks import callbacks

__all__ = ["QuasiHyperbolicParameterMonitor"]


@callbacks.register(name="quasi_hyperbolic_parameter_monitor")
class QuasiHyperbolicParameterMonitor(Callback):
    """Logs hyperbolic v1 parameter.

    This callback iterates over all optimizers and their parameter groups to log vs
    under the ``v{i}-{OPTIMIZER_NAME}/group{GROUP_NUMBER}`` key.

    Example:
        .. doctest::

            >>> from composer import Trainer
            >>> from composer.callbacks import LRMonitor
            >>> # constructing trainer object with this callback
            >>> trainer = Trainer(
            ...     model=model,
            ...     train_dataloader=train_dataloader,
            ...     eval_dataloader=eval_dataloader,
            ...     optimizers=optimizer,
            ...     max_duration="1ep",
            ...     callbacks=[LRMonitor()],
            ... )

    The vs are logged by the :class:`.Logger` to the following key as described
    below.

    +---------------------------------------------+---------------------------------------+
    | Key                                         | Logged data                      |
    +=============================================+=======================================+
    |                                             | vs for each optimizer and  |
    | ``v_{idx}-{OPTIMIZER_NAME}/group{GROUP_NUMBER}`` | parameter group is|
    |                                             | logged to a separate key.         |
    +---------------------------------------------+---------------------------------------+

    """

    def batch_end(self, state: State, logger: Logger) -> None:  # noqa: PLR6301
        """Log the vs at the end of the batch.

        In case we are using schedulers over the vs.

        Parameters
        ----------
        state : State
            The current state of training.
        logger : Logger
            The logger to log the vs.

        """
        assert state.optimizers is not None, "optimizers must be defined"
        step = state.timestamp.batch.value
        for optimizer in state.optimizers:
            vs: list[tuple[float, ...]] = [
                group["vs"] for group in optimizer.param_groups
            ]
            for idx, inner_vs in enumerate(vs):
                for j, v in enumerate(inner_vs):
                    logger.log_metrics({f"v{j}-group{idx}": v}, step)
