import collections


def collect_hierarchical_parameter_name(base_name, model, parameter_to_name=None, with_grad_only=False):
    """
        Create a meaningful name of the module's parameters based on the module hierarchy

        Args:
            base_name: the base name
            model: the model
            parameter_to_name: where to store the module to name conversion
            with_grad_only: only the parameters requiring gradient are collected

        Returns:
            a dictionary with mapping nn.Parameter to string
        """
    if parameter_to_name is None:
        parameter_to_name = collections.OrderedDict()

    for child_id, child in enumerate(model.children()):
        child_name = base_name + '/' + type(child).__name__ + f'_{child_id}'
        for name, parameter in child.named_parameters(recurse=False):
            if with_grad_only and not parameter.requires_grad:
                # discard if not gradient
                continue
            parameter_name = child_name + '/' + name
            parameter_to_name[parameter] = parameter_name

        collect_hierarchical_parameter_name(
            child_name,
            child,
            parameter_to_name=parameter_to_name,
            with_grad_only=with_grad_only)

    return parameter_to_name