"""
"""

# Import standard libraries.
import copy

# Import 3rd-party packages.
import torch

# Import thoplw scripts.
from .counts import *
from .linfo  import LayersInfo

# Map from layer classes to hook functions.
COUNT_FUNCTIONS = {

    # Fully connected layers.
    torch.nn.Linear: count_linear,

    # Convolution layers.
    torch.nn.Conv1d: count_convNd,
    torch.nn.Conv2d: count_convNd,
    torch.nn.Conv3d: count_convNd,

    # Transposed convolution layers.
    torch.nn.ConvTranspose1d: count_convNd,
    torch.nn.ConvTranspose2d: count_convNd,
    torch.nn.ConvTranspose3d: count_convNd,

    # Normalization layers.
    torch.nn.BatchNorm1d   : count_batchnorm,
    torch.nn.BatchNorm2d   : count_batchnorm,
    torch.nn.BatchNorm3d   : count_batchnorm,
    torch.nn.LayerNorm     : count_layernorm,
    torch.nn.InstanceNorm1d: count_instancenorm,
    torch.nn.InstanceNorm2d: count_instancenorm,
    torch.nn.InstanceNorm3d: count_instancenorm,

    # Activation layers.
    torch.nn.ReLU     : count_zero,
    torch.nn.ReLU6    : count_zero,
    torch.nn.LeakyReLU: count_lrelu,
    torch.nn.PReLU    : count_prelu,
    torch.nn.Softmax  : count_softmax,

    # Pooling layers.
    torch.nn.MaxPool1d        : count_zero,
    torch.nn.MaxPool2d        : count_zero,
    torch.nn.MaxPool3d        : count_zero,
    torch.nn.AdaptiveMaxPool1d: count_zero,
    torch.nn.AdaptiveMaxPool2d: count_zero,
    torch.nn.AdaptiveMaxPool3d: count_zero,
    torch.nn.AvgPool1d        : count_avgpool,
    torch.nn.AvgPool2d        : count_avgpool,
    torch.nn.AvgPool3d        : count_avgpool,
    torch.nn.AdaptiveAvgPool1d: count_avgpool,
    torch.nn.AdaptiveAvgPool2d: count_avgpool,
    torch.nn.AdaptiveAvgPool3d: count_avgpool,

    # Dropout layers.
    torch.nn.Dropout: count_zero,

    # Upsampling layers.
    torch.nn.Upsample            : count_upsample,
    torch.nn.UpsamplingBilinear2d: count_upsample,
    torch.nn.UpsamplingNearest2d : count_upsample,

    # # RNN layers.
    # torch.nn.RNNCell : count_rnn_cell,
    # torch.nn.GRUCell : count_gru_cell,
    # torch.nn.LSTMCell: count_lstm_cell,
    # torch.nn.RNN     : count_rnn,
    # torch.nn.GRU     : count_gru,
    # torch.nn.LSTM    : count_lstm,

    # Other layers.
    torch.nn.Flatten     : count_zero,
    torch.nn.PixelShuffle: count_zero,
    torch.nn.ZeroPad2d   : count_zero,

    # The following layers should NOT be registered.
    torch.nn.Sequential: count_none,
}


def profile(model, tensor, custom_ops={}, verbose=True):
    """
    Compute MACs and number of parameters.

    Args:
        model      (torch.nn.Module): Target NN model.
        tensor     (torch.Tensor)   : Input tensor.
        custom_ops (dict)           : Custom count functions.
        verbose    (bool)           : Dump verbose output if True.

    Returns:
        (tuple): A tuple of (macs, number of params, layer details).
    """
    unknown_layers = set()

    count_functions = copy.copy(COUNT_FUNCTIONS)
    count_functions.update(custom_ops)

    for name, module in model.named_modules():

        # Get the corresponding hook function.
        count_func = count_functions.get(module.__class__, count_none)

        # Register the hook function.
        module.register_forward_hook(get_hook_func(count_func))

        # Memorize the layer class name if unknown.
        if (module.__class__ not in count_functions) and (module != model):
            unknown_layers.add(str(module.__class__))

    # Change to the evaluation mode.
    model.eval()

    # Forward inference with hook functions.
    with torch.no_grad():
        model(tensor)

    # Instanciate the LayerInfo class that stores the details of layer information.
    layers_info = LayersInfo(model.__class__.__name__, tensor.shape)

    # Collect log information generated by the forward inference.
    for name, module in model.named_modules():
        if hasattr(module, "thoplw_log"):
            layers_info.append(name, module.thoplw_log)

    # Compute total MACs and params.
    total_macs, total_params = layers_info.total()

    # Print verbose info.
    if verbose:

        # Print unknown layers.
        if len(unknown_layers) > 0:

            print("The following layers will be unexpectedly ignored:")

            for class_name in sorted(unknown_layers):
                print("  - " + class_name)

    return (total_macs, total_params, layers_info)


# vim: expandtab tabstop=4 shiftwidth=4 fdm=marker
