"""Module for collecting and checking FSDP modules in a PyTorch model hierarchy."""
import torch
from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel


def collect_fsdp_modules(
    module_name: str,
    module: torch.nn.Module,
    maps: tuple[dict[str, FullyShardedDataParallel], dict[str, torch.nn.Module]],
) -> tuple[dict[str, FullyShardedDataParallel], dict[str, torch.nn.Module]]:
    """Recursively collect FSDP and non-FSDP modules from a PyTorch model hierarchy.

    This function traverses the model hierarchy and populates two dictionaries:
    one for FSDP modules and another for non-FSDP modules. The dictionaries map
    module names to their corresponding module instances.

    Parameters
    ----------
    module_name : str
        The name of the current module being processed.
    module : torch.nn.Module
        The current module instance being processed.
    maps : tuple[dict[str, FullyShardedDataParallel], dict[str, torch.nn.Module]]
        A tuple containing two dictionaries to be populated:
        - The first dictionary maps module names to FSDP module instances
        - The second dictionary maps module names to non-FSDP module instances

    Returns
    -------
    tuple[dict[str, FullyShardedDataParallel], dict[str, torch.nn.Module]]
        The updated tuple of dictionaries containing the collected modules.

    """
    fsdp_modules, non_fsdp_modules = maps
    # Add the current module if it's an FSDP module
    if isinstance(module, FullyShardedDataParallel):
        fsdp_modules[module_name] = module
    if isinstance(module, torch.nn.Module):
        non_fsdp_modules[module_name] = module
    # Recursively explore all children
    for child_name, child_module in module.named_children():
        full_child_name = f"{module_name}.{child_name}" if module_name else child_name
        collect_fsdp_modules(
            full_child_name,
            child_module,
            (fsdp_modules, non_fsdp_modules),
        )
    return (fsdp_modules, non_fsdp_modules)


def is_leaf_fsdp_module(
    module_name: str,
    maps: tuple[dict[str, FullyShardedDataParallel], dict[str, torch.nn.Module]],
) -> bool:
    """Check if the module is a leaf FSDP module.

    A leaf FSDP module is one that doesn't contain any child FSDP modules.
    This function helps identify modules that can be processed directly.

    Parameters
    ----------
    module_name : str
        The name of the module to check.
    maps : tuple[dict[str, FullyShardedDataParallel], dict[str, torch.nn.Module]]
        A tuple containing two dictionaries:
        - The first maps module names to FSDP module instances
        - The second maps module names to non-FSDP module instances

    Returns
    -------
    bool
        True if the module is a leaf FSDP module, False otherwise.

    """
    fsdp_modules, non_fsdp_modules = maps
    # Get the module object
    if module_name in fsdp_modules:
        module = fsdp_modules[module_name]
    else:
        module = non_fsdp_modules[module_name]

    # Check if any child is an FSDP module
    for child_name, child_module in module.named_children():
        full_child_name = f"{module_name}.{child_name}" if module_name else child_name
        if isinstance(child_module, FullyShardedDataParallel):
            return False
        return is_leaf_fsdp_module(full_child_name, maps)
    return True
