# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

# pyre-strict

from collections import defaultdict
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple

import tabulate
import torch
from torch import nn

from .activation_count import ActivationCountAnalysis
from .flop_count import FlopCountAnalysis
from .parameter_count import parameter_count


### Pre-processing functions ###


def _format_size(x: int, sig_figs: int = 3, hide_zero: bool = False) -> str:
    """
    Formats an integer for printing in a table or model representation.
    Expresses the number in terms of 'kilo', 'mega', etc., using
    'K', 'M', etc. as a suffix.

    Args:
        x (int) : The integer to format.
        sig_figs (int) : The number of significant figures to keep
        hide_zero (bool) : If True, x=0 is replaced with an empty string
            instead of '0'.

    Returns:
        str : The formatted string.
    """
    if hide_zero and x == 0:
        return str("")

    def fmt(x: float) -> str:
        # use fixed point to avoid scientific notation
        return "{{:.{}f}}".format(sig_figs).format(x).rstrip("0").rstrip(".")

    if abs(x) > 1e14:
        return fmt(x / 1e15) + "P"
    if abs(x) > 1e11:
        return fmt(x / 1e12) + "T"
    if abs(x) > 1e8:
        return fmt(x / 1e9) + "G"
    if abs(x) > 1e5:
        return fmt(x / 1e6) + "M"
    if abs(x) > 1e2:
        return fmt(x / 1e3) + "K"
    return str(x)


def _pretty_statistics(
    statistics: Dict[str, Dict[str, int]], sig_figs: int = 3, hide_zero: bool = False
) -> Dict[str, Dict[str, str]]:
    """
    Converts numeric statistics to strings with kilo/mega/giga/etc.
    labels.

    Args:
        statistics (dict(str, dict(str, int))) : the statistics to
            format. Organized as a dictionary over modules, which are
            each a dictionary over statistic types.
        sig_figs (int) : the number of significant figures for each stat
        hide_zero (bool) : if True, statistics that are zero will be
            written as an empty string. Defaults to False.

    Return:
        dict(str, dict(str, str)) : the input statistics as pretty strings
    """
    out_stats = {}
    for mod, stats in statistics.items():
        out_stats[mod] = {
            s: _format_size(val, sig_figs, hide_zero) for s, val in stats.items()
        }
    return out_stats


def _group_by_module(
    statistics: Dict[str, Dict[str, Any]]
) -> Dict[str, Dict[str, Any]]:
    """
    Converts statistics organized first by statistic type and then by module
    to statistics organized first by module and then by statistic type.

    Args:
        statistics (dict(str, dict(str, any))) : the statistics to convert

    Returns:
        dict(str, dict(str, any)) : the reorganized statistics
    """
    out_stats = defaultdict(dict)
    for stat_name, stat in statistics.items():
        for mod, val in stat.items():
            out_stats[mod][stat_name] = val
    return dict(out_stats)


def _indicate_uncalled_modules(
    statistics: Dict[str, Dict[str, str]],
    stat_name: str,
    uncalled_modules: Set[str],
    uncalled_indicator: str = "N/A",
) -> Dict[str, Dict[str, str]]:
    """
    If a module is in the set of uncalled modules, replace its statistics
    with the specified indicator, instead of using the existing string.
    Assumes the statistic is already formatting in string form.

    Args:
        statistics (dict(str, dict(str, str))) : the statistics to
            format. Organized as a dictionary over modules, which are
            each a dictionary over statistic types. Expects statistics
            have already been converted to strings.
        stat_name (str) : the name of the statistic being modified
        uncalled_modules set(str) : a set of names of uncalled modules.
        indicator (str) : the string that will be used to indicate
            unused modules. Defaults to 'N/A'.

    Returns:
        dict(str, dict(str, str)) : the modified statistics
    """

    stats_out = {mod: stats.copy() for mod, stats in statistics.items()}
    for mod in uncalled_modules:
        if mod not in stats_out:
            stats_out[mod] = {}
        stats_out[mod][stat_name] = uncalled_indicator
    return stats_out


def _remove_zero_statistics(
    statistics: Dict[str, Dict[str, int]],
    force_keep: Optional[Set[str]] = None,
    require_trivial_children: bool = False,
) -> Dict[str, Dict[str, int]]:
    """
    Any module that has zero for all available statistics is removed from the
    set of statistics. This can help declutter the reporting of statistics
    if many submodules have zero statistics. Assumes the statistics have
    a model hierarchy starting with a root that has name ''.

    Args:
        statistics (dict(str, dict(str, int))) : the statistics to
            remove zeros from. Organized as a dictionary over modules,
            which are each a dictionary over statistic types.
        force_keep (set(str) or None) : a set of modules to always keep, even
            if they are all zero.
        require_trivial_children (bool) : If True, a statistic will only
            be deleted if all its children are also deleted. Defaults to
            False.

    Returns:
        dict(str, dict(str, int)) : the input statistics dictionary,
            with submodules removed if they have zero for all statistics.
    """
    out_stats: Dict[str, Dict[str, int]] = {}
    _force_keep: Set[str] = force_keep if force_keep else set() | {""}

    def keep_stat(name: str) -> None:
        prefix = name + ("." if name else "")
        trivial_children = True
        for mod in statistics:
            # 'if mod' excludes root = '', which is never a child
            if mod and mod.count(".") == prefix.count(".") and mod.startswith(prefix):
                keep_stat(mod)
                trivial_children &= mod not in out_stats

        if (
            (not all(val == 0 for val in statistics[name].values()))
            or (name in _force_keep)
            or (require_trivial_children and not trivial_children)
        ):
            out_stats[name] = statistics[name].copy()

    keep_stat("")
    return out_stats


def _fill_missing_statistics(
    model: nn.Module, statistics: Dict[str, Dict[str, int]]
) -> Dict[str, Dict[str, int]]:
    """
    If, for a given submodule name in the model, a statistic is missing
    from statistics, fills it in with zero. This visually uniformizes
    the reporting of statistics.

    Args:
        model (nn.Module) : the model whose submodule names will be
            used to fill statistics
        statistics (dict(str, dict(str, int))) : the statistics to
            fill in missing values for. Organized as a dictionary
            over statistics, which are each a dictionary over submodules'
            names. The statistics are assumed to be formatted already
            to the desired string format for printing.

    Returns:
        dict(str, dict(str, int)) : the input statistics with missing
            values filled with zero.
    """
    out_stats = {name: stat.copy() for name, stat in statistics.items()}
    for mod_name, _ in model.named_modules():
        for stat in out_stats.values():
            if mod_name not in stat:
                stat[mod_name] = 0
    return out_stats


### Model String Printing ###


def _model_stats_str(model: nn.Module, statistics: Dict[str, Dict[str, str]]) -> str:
    """
    This produces a representation of the model much like 'str(model)'
    would, except the provided statistics are written out as additional
    information for each submodule.

    Args:
        model (nn.Module) : the model to form a representation of.
        statistics (dict(str, dict(str, str))) : the statistics to
            include in the model representations. Organized as a dictionary
            over module names, which are each a dictionary over statistics.
            The statistics are assumed to be formatted already to the
            desired string format for printing.

    Returns:
        str : the string representation of the model with the statistics
            inserted.
    """

    # Copied from nn.Module._addindent
    def _addindent(s_: str, numSpaces: int) -> str:
        s = s_.split("\n")
        # don't do anything for single-line stuff
        if len(s) == 1:
            return s_
        first = s.pop(0)
        s = [(numSpaces * " ") + line for line in s]
        s = "\n".join(s)
        s = first + "\n" + s
        return s

    def print_statistics(name: str) -> str:
        if name not in statistics:
            return ""
        printed_stats = ["{}: {}".format(k, v) for k, v in statistics[name].items()]
        return ", ".join(printed_stats)

    # This comes directly from nn.Module.__repr__ with small changes
    # to include the statistics.
    def repr_with_statistics(module: nn.Module, name: str) -> str:
        # We treat the extra repr like the sub-module, one item per line
        extra_lines = []
        extra_repr = module.extra_repr()
        printed_stats = print_statistics(name)
        # empty string will be split into list ['']
        if extra_repr:
            extra_lines.extend(extra_repr.split("\n"))
        if printed_stats:
            extra_lines.extend(printed_stats.split("\n"))
        child_lines = []
        for key, submod in module._modules.items():
            submod_name = name + ("." if name else "") + key
            # pyre-fixme[6]: Expected `Module` for 1st param but got
            #  `Optional[nn.modules.module.Module]`.
            submod_str = repr_with_statistics(submod, submod_name)
            submod_str = _addindent(submod_str, 2)
            child_lines.append("(" + key + "): " + submod_str)
        lines = extra_lines + child_lines

        main_str = module._get_name() + "("
        if lines:
            # simple one-liner info, which most builtin Modules will use
            if len(extra_lines) == 1 and not child_lines:
                main_str += extra_lines[0]
            else:
                main_str += "\n  " + "\n  ".join(lines) + "\n"

        main_str += ")"
        return main_str

    return repr_with_statistics(model, "")


def _get_input_sizes(iterable: Iterable[Any]) -> List[Any]:  # pyre-ignore[2,3]
    """
    Gets the sizes of all torch tensors in an iterable. If an element
    of the iterable is a non-torch tensor iterable, it recurses into
    that iterable to continue calculating sizes. Any non-iterable is
    given a size of None. The output consists of nested lists with the
    same nesting structure as the input iterables.
    """
    out_list = []
    for i in iterable:
        if isinstance(i, torch.Tensor):
            out_list.append(list(i.size()))
        elif isinstance(i, Iterable):
            sublist_sizes = _get_input_sizes(i)
            if all(j is None for j in sublist_sizes):
                out_list.append(None)
            else:
                out_list.append(sublist_sizes)
        else:
            out_list.append(None)
    return out_list


def flop_count_str(
    flops: FlopCountAnalysis, activations: Optional[ActivationCountAnalysis] = None
) -> str:
    """
    Calculates the parameters and flops of the model with the given inputs
    and returns a string representation of the model that includes the
    parameters and flops of every submodule. The string is structured
    to be similar that given by str(model), though it is not guaranteed to
    be identical in form if the default string representation of a module has
    been overridden. If a module has zero parameters and flops, statistics
    will not be reported for succinctness.

    The trace can only register the scope of a module if it is called
    directly, which means flops (and activations) arising from explicit
    calls to .forward() or to other python functions of the module will
    not be attributed to that module. Modules that are never called will
    have 'N/A' listed for their flops; this means they are either unused
    or their statistics are missing for this reason. Any such flops are still
    counted towards the parent

    Example:

    >>> import torch
    >>> import torch.nn as nn

    >>> class InnerNet(nn.Module):
    ...     def __init__(self):
    ...         super().__init__()
    ...         self.fc1 = nn.Linear(10,10)
    ...         self.fc2 = nn.Linear(10,10)
    ...     def forward(self, x):
    ...         return self.fc1(self.fc2(x))

    >>> class TestNet(nn.Module):
    ...     def __init__(self):
    ...         super().__init__()
    ...         self.fc1 = nn.Linear(10,10)
    ...         self.fc2 = nn.Linear(10,10)
    ...         self.inner = InnerNet()
    ...     def forward(self, x):
    ...         return self.fc1(self.fc2(self.inner(x)))

    >>> inputs = torch.randn((1,10))
    >>> print(flop_count_str(FlopCountAnalysis(model, inputs)))
    TestNet(
      #params: 0.44K, #flops: 0.4K
      (fc1): Linear(
        in_features=10, out_features=10, bias=True
        #params: 0.11K, #flops: 100
      )
      (fc2): Linear(
        in_features=10, out_features=10, bias=True
        #params: 0.11K, #flops: 100
      )
      (inner): InnerNet(
        #params: 0.22K, #flops: 0.2K
        (fc1): Linear(
          in_features=10, out_features=10, bias=True
          #params: 0.11K, #flops: 100
        )
        (fc2): Linear(
          in_features=10, out_features=10, bias=True
          #params: 0.11K, #flops: 100
        )
      )
    )


    Args:
        flops (FlopCountAnalysis): the flop counting object
        activations (bool) : If given, the activations of each layer will
            also be calculated and included in the representation.

    Returns:
        str:
            a string representation of the model with the number of
            parameters and flops included.
    """
    # cast to dict since pyre doesn't like the implicit defaultdict->dict
    model = flops._model
    params = dict(parameter_count(model))

    flops.unsupported_ops_warnings(False)
    flops.uncalled_modules_warnings(False)
    flops.tracer_warnings("none")
    stats = {"#params": params, "#flops": flops.by_module()}

    if activations is not None:
        activations.unsupported_ops_warnings(False)
        activations.uncalled_modules_warnings(False)
        activations.tracer_warnings("none")
        stats["#acts"] = activations.by_module()

    all_uncalled = flops.uncalled_modules() | (
        activations.uncalled_modules() if activations is not None else set()
    )
    stats = _fill_missing_statistics(model, stats)
    stats = _group_by_module(stats)
    stats = _remove_zero_statistics(stats, force_keep=all_uncalled)
    stats = _pretty_statistics(stats, sig_figs=2)
    stats = _indicate_uncalled_modules(stats, "#flops", flops.uncalled_modules())
    if activations is not None:
        stats = _indicate_uncalled_modules(
            stats, "#acts", activations.uncalled_modules()
        )

    model_string = ""
    if all_uncalled:
        model_string += (
            "N/A indicates a possibly missing statistic due to how "
            "the module was called. Missing values are still included "
            "in the parent's total.\n"
        )
    model_string += _model_stats_str(model, stats)
    return model_string


### Table Printing ###


def _get_single_child(
    name: str, statistics: Dict[str, Dict[str, str]]
) -> Optional[str]:
    """
    If the given module has only a single child in statistics, return it.
    Otherwise, return None.
    """
    prefix = name + ("." if name else "")
    child = None
    for mod in statistics:
        # 'if mod' excludes root = '', which is never a child
        if mod and mod.count(".") == prefix.count(".") and mod.startswith(prefix):
            if child is None:
                child = mod
            else:
                return None  # We found a second child, so return None
    return child


def _try_combine(
    stats1: Dict[str, str], stats2: Dict[str, str]
) -> Optional[Dict[str, str]]:
    """
    Try combine two statistics dict to display in one row. If they conflict,
    returns None.
    """
    ret = {}
    if set(stats1.keys()) != set(stats2.keys()):
        return None
    for k, v1 in stats1.items():
        v2 = stats2[k]
        if v1 != v2 and len(v1) and len(v2):
            return None
        ret[k] = v1 if len(v1) else v2
    return ret


def _fastforward(
    name: str, statistics: Dict[str, Dict[str, str]]
) -> Tuple[str, Dict[str, str]]:
    """
    If the given module has only a single child and matches statistics
    with that child, merge statistics and their names into one row.
    Then repeat until the condition isn't met.

    Returns:
        str: the new name
        dict: the combined statistics of this row
    """
    single_child = _get_single_child(name, statistics)
    if single_child is None:
        return name, statistics[name]
    combined = _try_combine(statistics[name], statistics[single_child])
    if combined is None:
        return name, statistics[name]
    statistics[single_child] = combined
    return _fastforward(single_child, statistics)


def _model_stats_table(
    statistics: Dict[str, Dict[str, str]],
    max_depth: int = 3,
    stat_columns: Optional[List[str]] = None,
) -> str:
    """
    Formats the statistics obtained from a model in a nice table.

    Args:
        statistics (dict(str, dict(str, str))) : The statistics to print.
            Organized as a dictionary over modules, then as a dictionary
            over statistics in the model. The statistics are assumed to
            already be formatted for printing.
        max_depth (int) : The maximum submodule depth to recurse to.
        stat_columns (list(str)) : Specify the order of the columns to print.
            If None, columns are found automatically from the provided
            statistics.

    Return:
        str : The formatted table.
    """
    if stat_columns is None:
        stat_columns = set()
        for stats in statistics.values():
            stat_columns.update(stats.keys())
        stat_columns = list(stat_columns)

    headers = ["module"] + stat_columns
    table: List[List[str]] = []

    def build_row(name: str, stats: Dict[str, str], indent_lvl: int) -> List[str]:
        indent = " " * indent_lvl
        row = [indent + name]
        for stat_name in stat_columns:  # pyre-ignore[16] Is not None at this point
            row_str = (indent + stats[stat_name]) if stat_name in stats else ""
            row.append(row_str)
        return row

    def fill(indent_lvl: int, prefix: str) -> None:
        if indent_lvl > max_depth:
            return
        for mod_name in statistics:
            # 'if mod' excludes root = '', which is never a child
            if (
                mod_name
                and mod_name.count(".") == prefix.count(".")
                and mod_name.startswith(prefix)
            ):
                mod_name, curr_stats = _fastforward(mod_name, statistics)
                if root_prefix and mod_name.startswith(root_prefix):
                    # Skip the root_prefix shared by all submodules as it carries 0 information
                    pretty_mod_name = mod_name[len(root_prefix) :]
                else:
                    pretty_mod_name = mod_name
                row = build_row(pretty_mod_name, curr_stats, indent_lvl)
                table.append(row)
                fill(indent_lvl + 1, mod_name + ".")

    root_name, curr_stats = _fastforward("", statistics)
    row = build_row(root_name or "model", curr_stats, indent_lvl=0)
    table.append(row)
    root_prefix = root_name + ("." if root_name else "")
    fill(indent_lvl=1, prefix=root_prefix)

    old_ws = tabulate.PRESERVE_WHITESPACE
    tabulate.PRESERVE_WHITESPACE = True
    tab = tabulate.tabulate(table, headers=headers, tablefmt="pipe")
    tabulate.PRESERVE_WHITESPACE = old_ws
    return tab


def flop_count_table(
    flops: FlopCountAnalysis,
    max_depth: int = 3,
    activations: Optional[ActivationCountAnalysis] = None,
    show_param_shapes: bool = True,
) -> str:
    """
    Format the per-module parameters and flops of a model in a table.
    It looks like this:
    ::
        | model                            | #parameters or shape   | #flops    |
        |:---------------------------------|:-----------------------|:----------|
        | model                            | 34.6M                  | 65.7G     |
        |  s1                              |  15.4K                 |  4.32G    |
        |   s1.pathway0_stem               |   9.54K                |   1.23G   |
        |    s1.pathway0_stem.conv         |    9.41K               |    1.23G  |
        |    s1.pathway0_stem.bn           |    0.128K              |           |
        |   s1.pathway1_stem               |   5.9K                 |   3.08G   |
        |    s1.pathway1_stem.conv         |    5.88K               |    3.08G  |
        |    s1.pathway1_stem.bn           |    16                  |           |
        |  s1_fuse                         |  0.928K                |  29.4M    |
        |   s1_fuse.conv_f2s               |   0.896K               |   29.4M   |
        |    s1_fuse.conv_f2s.weight       |    (16, 8, 7, 1, 1)    |           |
        |   s1_fuse.bn                     |   32                   |           |
        |    s1_fuse.bn.weight             |    (16,)               |           |
        |    s1_fuse.bn.bias               |    (16,)               |           |
        |  s2                              |  0.226M                |  7.73G    |
        |   s2.pathway0_res0               |   80.1K                |   2.58G   |
        |    s2.pathway0_res0.branch1      |    20.5K               |    0.671G |
        |    s2.pathway0_res0.branch1_bn   |    0.512K              |           |
        |    s2.pathway0_res0.branch2      |    59.1K               |    1.91G  |
        |   s2.pathway0_res1.branch2       |   70.4K                |   2.28G   |
        |    s2.pathway0_res1.branch2.a    |    16.4K               |    0.537G |
        |    s2.pathway0_res1.branch2.a_bn |    0.128K              |           |
        |    s2.pathway0_res1.branch2.b    |    36.9K               |    1.21G  |
        |    s2.pathway0_res1.branch2.b_bn |    0.128K              |           |
        |    s2.pathway0_res1.branch2.c    |    16.4K               |    0.537G |
        |    s2.pathway0_res1.branch2.c_bn |    0.512K              |           |
        |   s2.pathway0_res2.branch2       |   70.4K                |   2.28G   |
        |    s2.pathway0_res2.branch2.a    |    16.4K               |    0.537G |
        |    s2.pathway0_res2.branch2.a_bn |    0.128K              |           |
        |    s2.pathway0_res2.branch2.b    |    36.9K               |    1.21G  |
        |    s2.pathway0_res2.branch2.b_bn |    0.128K              |           |
        |    s2.pathway0_res2.branch2.c    |    16.4K               |    0.537G |
        |    s2.pathway0_res2.branch2.c_bn |    0.512K              |           |
        |    ............................. |    ......              |    ...... |

    Args:
        flops (FlopCountAnalysis): the flop counting object
        max_depth (int) : The max depth of submodules to include in the
            table. Defaults to 3.
        activations (ActivationCountAnalysis or None) : If given, include
            activation counts as an additional column in the table.
        show_param_shapes (bool) : If true, shapes for parameters will be
            included in the table. Defaults to True.

    Returns:
        str : The formatted table.

    Examples:
    ::
        print(flop_count_table(FlopCountAnalysis(model, inputs)))
    """
    params_header = "#parameters" + (" or shape" if show_param_shapes else "")
    flops_header, acts_header = "#flops", "#activations"

    model = flops._model
    # cast to dict since pyre doesn't like the implicit defaultdict->dict
    params = dict(parameter_count(model))

    flops.unsupported_ops_warnings(False)
    flops.uncalled_modules_warnings(False)
    flops.tracer_warnings("none")

    stats = {params_header: params, flops_header: flops.by_module()}
    stat_columns = [params_header, flops_header]

    if activations is not None:
        activations.unsupported_ops_warnings(False)
        activations.uncalled_modules_warnings(False)
        activations.tracer_warnings("none")
        stats[acts_header] = activations.by_module()
        stat_columns += [acts_header]

    stats = _group_by_module(stats)
    stats = _remove_zero_statistics(stats, require_trivial_children=True)
    stats = _pretty_statistics(stats, hide_zero=False)
    stats = _indicate_uncalled_modules(
        stats,
        flops_header,
        flops.uncalled_modules() & stats.keys(),
        uncalled_indicator="",
    )
    if activations:
        stats = _indicate_uncalled_modules(
            stats,
            acts_header,
            activations.uncalled_modules() & stats.keys(),
            uncalled_indicator="",
        )

    # Swap in shapes for parameters or delete shapes from dict
    param_shapes: Dict[str, Tuple[int, ...]] = {
        k: tuple(v.shape) for k, v in model.named_parameters()
    }
    to_delete = []
    for mod in stats:
        if mod in param_shapes:
            if show_param_shapes:
                stats[mod][params_header] = str(param_shapes[mod])
            else:
                to_delete.append(mod)
    for mod in to_delete:
        del stats[mod]

    return _model_stats_table(
        statistics=stats,
        max_depth=max_depth,
        stat_columns=stat_columns,
    )
