import torch
from typing import Union
from typing import Type
from typing import Tuple


def count_module_instances(module: torch.nn.Module, module_class: Union[Type[torch.nn.Module],
                                                                        Tuple[Type[torch.nn.Module], ...]]) -> int:
    """Counts the number of instances of ``module_class`` in ``module``, recursively.

    .. rubric:: Example

    .. testsetup::

        from composer.utils.module_surgery import count_module_instances

    .. doctest::

        >>> from torch import nn
        >>> module = nn.Sequential(nn.Linear(16, 32), nn.Linear(32, 64), nn.ReLU())
        >>> count_module_instances(module, nn.Linear)
        2
        >>> count_module_instances(module, (nn.Linear, nn.ReLU))
        3

    Args:
        module (torch.nn.Module): The source module.
        module_class (Type[torch.nn.Module] | Tuple[Type[torch.nn.Module], ...]):
            The module type (or tuple of module types) to count.

    Returns:
        int: The number of instances of ``module_class`` in ``module``
    """
    count = 0
    for _, child in module.named_children():
        if isinstance(child, module_class):
            count += 1
        count += count_module_instances(child, module_class)

    return count