import torch
from torch.nn.modules.conv import _ConvNd
from typing import Tuple


def count_convNd(m: _ConvNd, layer_input: Tuple[torch.Tensor, ],
                 layer_output: torch.Tensor):
    layer_input = layer_input[0]

    kernel_ops = torch.zeros(m.weight.size()[2:]).numel()  # Kw x Kh
    bias_ops = 1 if m.bias is not None else 0

    # N x Cout x H x W x  (Cin x Kw x Kh + bias)
    total_ops = layer_output.nelement() * (
        m.in_channels // m.groups * kernel_ops + bias_ops)

    m.total_ops += torch.DoubleTensor([int(total_ops)])


def count_convNd_ver2(m: _ConvNd, layer_input: Tuple[torch.Tensor, ],
                      layer_output: torch.Tensor):
    layer_input = layer_input[0]

    # N x H x W (exclude Cout)
    output_size = torch.zeros(
        (layer_output.size()[:1] + layer_output.size()[2:])).numel()
    # Cout x Cin x Kw x Kh
    kernel_ops = m.weight.nelement()
    if m.bias is not None:
        # Cout x 1
        kernel_ops += +m.bias.nelement()
    # x N x H x W x Cout x (Cin x Kw x Kh + bias)
    m.total_ops += torch.DoubleTensor([int(output_size * kernel_ops)])
