"""Custom optimizer builder for mt_dao models.

Allowing
models to specify their own parameter groups. It integrates with the llmfoundry
"""

from collections.abc import Iterable
from typing import Any, Protocol, runtime_checkable

import torch
from llmfoundry import registry
from llmfoundry.utils.builders import _extract_param_groups  # noqa: PLC2701
from llmfoundry.utils.builders import build_optimizer as llmfoundry_build_optimizer
from llmfoundry.utils.registry_utils import construct_from_registry
from torch.optim.optimizer import Optimizer


@runtime_checkable
class HasGetParamGroups(Protocol):
    """Protocol for models that can provide their optimizer parameter groups.

    This is useful for models that have custom parameter groups or need to
    modify the optimizer configuration based on their parameters. Such as models
    using MuP or CompleteP.
    """

    # NOTE: This is a protocol, so it can be used with isinstance checks.
    # NOTE: The method signature must match the one in the model.

    def get_optimizer_param_groups(  # <- exact signature
        self,
        optimizer_config: dict[str, Any],
    ) -> tuple[Iterable[torch.Tensor] | Iterable[dict[str, Any]], dict[str, Any]]:
        """Get the optimizer parameter groups and any modified optimizer config.

        Parameters
        ----------
        optimizer_config: (dict[str, Any])
            The optimizer configuration.

        Returns
        -------
            tuple[Iterable[torch.Tensor] | Iterable[dict[str, Any]], dict[str,
            Any]]: A tuple containing the parameter groups and any modified
            optimizer configuration.

        """
        ...


def build_optimizer(
    model: torch.nn.Module,
    name: str,
    optimizer_config: dict[str, Any],
) -> Optimizer:
    """Build an llmfoundry optimizer.

    This allows the model to specify its own parameter groups.

    Parameters
    ----------
    model: torch.nn.Module
        The model for which to build the optimizer.
    name: str
        The name of the optimizer to build.
    optimizer_config: dict[str, Any]
        The configuration for the optimizer, which may include parameters like
    learning rate, weight decay, etc.

    Returns
    -------
    Optimizer
        An instance of the optimizer built with the specified configuration and
        the model's parameters.

    Raises
    ------
    ValueError
        If the `params` key is present in the `optimizer_config`, as it will be
        automatically extracted from the model and optimizer config.

    """
    if not isinstance(model, HasGetParamGroups):
        # If the model does not implement get_optimizer_param_groups, use the
        # default build_optimizer function.
        return llmfoundry_build_optimizer(
            model=model,
            name=name,
            optimizer_config=optimizer_config,
        )

    # NOTE: Modified from Composer's build_optimizer to support models that
    # have custom parameter groups or need to modify the optimizer configuration
    # based on their parameters. Such as models using MuP or CompleteP.

    # NOTE: call to make sure disable_grad is correctly applied
    # as well as any other transformations
    # NOTE: incompatible with manually specifying
    # param groups in optimizer_config
    _extract_param_groups(model, optimizer_config)

    params, new_optim_config = model.get_optimizer_param_groups(
        optimizer_config,
    )
    kwargs = {
        **new_optim_config,
    }

    if "params" in kwargs:
        msg = (
            "The `params` will be automatically extracted from the model and "
            "optimizer config. Please remove it from the optimizer config kwargs."
        )
        raise ValueError(
            msg,
        )

    kwargs["params"] = params
    return construct_from_registry(
        name=name,
        registry=registry.optimizers,
        partial_function=True,
        pre_validation_function=Optimizer,
        post_validation_function=None,
        kwargs=kwargs,
    )
