"""Utilities for retrieving optimizer states in PyTorch models."""
import os
from typing import cast

import torch
from torch.distributed.checkpoint.state_dict import (
    OptimizerStateType,
    StateDictOptions,
    get_optimizer_state_dict,
)
from torch.distributed.fsdp import FullyShardedDataParallel

from mt_dao.constants import RANK


def get_optimizer_states_simplified(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    fsdp_state_dict_type: str | None = None,
    *,
    fsdp_enabled: bool = False,
) -> dict[str, OptimizerStateType] | dict[str, dict[str, dict[str, torch.Tensor]]]:
    """Retrieve the optimizer states for a given model and optimizer.

    This function handles both standard and FSDP (Fully Sharded Data Parallel) models.
    If the model is wrapped in FSDP, it retrieves the optimizer state using
    `get_optimizer_state_dict` with appropriate options. If the model is not wrapped in
    FSDP, it retrieves the optimizer state using `fsdp_get_optim_state_dict` or
    `optimizer.state_dict()` based on the `fsdp_state_dict_type`.

    Parameters
    ----------
    model : torch.nn.Module
        The model for which the optimizer states are to be retrieved.
    optimizer : torch.optim.Optimizer
        The optimizer whose states are to be retrieved.
    fsdp_state_dict_type : str | None, optional
        The type of FSDP state dictionary to use. If None, the optimizer state will
        be retrieved without FSDP considerations. If "full" or "sharded", it will
        retrieve the full or sharded state dictionary, respectively.
    fsdp_enabled : bool, optional
        Whether the model is wrapped in FSDP. Defaults to False.

    Returns
    -------
    dict[str, OptimizerStateType] | dict[str, dict[str, dict[str, torch.Tensor]]]
        A dictionary containing the optimizer states. If FSDP is enabled and the
        `fsdp_state_dict_type` is set to "full" or "sharded", it returns a more complex
        structure with state information. Otherwise, it returns a simpler structure.

    """
    optim_state_dict = get_optimizer_state_dict(
        model=model,
        optimizers=optimizer,
        submodules=None,
        options=StateDictOptions(
            full_state_dict=fsdp_state_dict_type == "full",
            cpu_offload=fsdp_enabled,
        ),
    )
    return {type(optimizer).__qualname__: optim_state_dict}


def get_optimizer_states_base(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    fsdp_state_dict_type: str | None = None,
    *,
    fsdp_enabled: bool = False,
) -> dict[str, dict[str, torch.Tensor]] | None:
    """Retrieve the optimizer states for a given model and optimizer.

    This function handles both standard and FSDP (Fully Sharded Data Parallel) models.
    If the model is wrapped in FSDP, it retrieves the optimizer state using
    `get_optimizer_state_dict` with appropriate options. If the model is not wrapped in
    FSDP, it retrieves the optimizer state using `get_optimizer_states_simplified`.

    Parameters
    ----------
    model : torch.nn.Module
        The model for which the optimizer states are to be retrieved.
    optimizer : torch.optim.Optimizer
        The optimizer whose states are to be retrieved.
    fsdp_state_dict_type : str | None, optional
        The type of FSDP state dictionary to use. If None, the optimizer state will
        be retrieved without FSDP considerations. If "full" or "sharded", it will
        retrieve the full or sharded state dictionary, respectively.
    fsdp_enabled : bool, optional
        Whether the model is wrapped in FSDP. Defaults to False.

    Returns
    -------
    dict[str, dict[str, torch.Tensor]] | None
        A dictionary containing the optimizer states, where the keys are the optimizer
        state names (e.g., 'exp_avg', 'exp_avg_sq') and the values are the corresponding
        tensors. If the model is wrapped in FSDP and the rank is not 0, it returns None.

    """
    # Check if the model is wrapped in FSDP
    is_fsdp = False
    if hasattr(model, "model") and isinstance(
        model,
        FullyShardedDataParallel,
    ):
        is_fsdp = True

    # Get the optimizer state dict based on whether FSDP is used
    if is_fsdp:
        optim_state_dict = get_optimizer_state_dict(
            model=model,
            optimizers=optimizer,
            options=StateDictOptions(
                full_state_dict=True,
                cpu_offload=True,
            ),
        )
    else:
        optim_state_dict = get_optimizer_states_simplified(
            model=model,
            optimizer=optimizer,
            fsdp_state_dict_type=fsdp_state_dict_type,
            fsdp_enabled=fsdp_enabled,
        )[type(optimizer).__qualname__]

    return (
        cast("dict[str, dict[str, torch.Tensor]]", optim_state_dict["state"])
        if int(os.environ.get(RANK, "0")) == 0
        else None
    )
