import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.intrinsic as nni
from torch.nn.intrinsic.modules import _FusedModule
from mqbench.utils.utils import is_tracing_state


class Conv2d(nn.Conv2d):
    r"""
    A Conv2d module used for fix-point forward.

    We adopt the same interface as `torch.nn.Conv2d`, please see
    https://pytorch.org/docs/stable/nn.html?highlight=conv2d#torch.nn.Conv2d
    for documentation.

    Similar to `torch.nn.Conv2d`, with DeQuantize modules initialized to
    default.

    """

    def __init__(self, 
                in_channels,
                out_channels,
                kernel_size,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=True,
                padding_mode='zeros',
                padding_value=None, 
                is_relu=0,
                weight_bitwidth=4,
                weight_half_range=False,
                input_bitwidth=8,
                input_half_range=False,
                output_bitwidth=32,
                bias_bitwidth=32,
                **kwargs) -> None:
        super().__init__(in_channels,
                        out_channels,
                        kernel_size=kernel_size,
                        stride=stride,
                        padding=padding,
                        dilation=dilation,
                        groups=groups,
                        bias=bias,
                        padding_mode=padding_mode,
                        **kwargs)
        self.is_relu = is_relu
        self.weight_bitwidth = weight_bitwidth
        self.weight_half_range = weight_half_range
        self.input_bitwidth = input_bitwidth
        self.input_half_range = input_half_range
        self.output_bitwidth = output_bitwidth
        self.bias_bitwidth = bias_bitwidth
        self.padding_value = nn.Parameter(padding_value.to(torch.float32).round().to(torch.int32), requires_grad=False)

    def _forward_conv(self, input):
        input = self._add_padding_value(input)
        output = super().forward(input)
        output = self._remove_padding_value(output)
        return output

    def _add_padding_value(self, input):
        if self.padding_value is not None:
            if self.padding_value.shape[0] == 1:
                input = F.pad(input, (self.kernel_size[0]*self.stride[0],  \
                                        self.kernel_size[0]*self.stride[0], \
                                            self.kernel_size[0]*self.stride[0], \
                                                self.kernel_size[0]*self.stride[0], 0, 0), \
                            mode='constant', value=self.padding_value.item())
            else:
                input = F.pad(input, (self.kernel_size[0]*self.stride[0],  \
                                        self.kernel_size[0]*self.stride[0], \
                                            self.kernel_size[0]*self.stride[0], \
                                                self.kernel_size[0]*self.stride[0], 0, 0), \
                            mode='constant', value=0.)
                for i in range(self.padding_value.shape[0]):
                    input[:, i, :self.kernel_size[0]*self.stride[0], :].fill_(self.padding_value[i].item())
                    input[:, i, -self.kernel_size[0]*self.stride[0]:, :].fill_(self.padding_value[i].item())
                    input[:, i, :, :self.kernel_size[0]*self.stride[0]].fill_(self.padding_value[i].item())
                    input[:, i, :, -self.kernel_size[0]*self.stride[0]:].fill_(self.padding_value[i].item())
        return input

    def _remove_padding_value(self, output):
        if self.padding_value is not None:
            output = output[:, :, self.kernel_size[0]:-self.kernel_size[0], self.kernel_size[0]:-self.kernel_size[0]]
        return output

    @torch.jit.export
    def extra_repr(self):
        return 'dilation={}, groups={}, kernel_size={}, ' \
               'padding={}, stride={}, is_relu={}, padding_value={}, weight_bitwidth={}'.format(
                   self.dilation, self.groups, self.kernel_size,
                   self.padding, self.stride,
                   self.is_relu, self.padding_value, self.weight_bitwidth)

    def forward(self, input):
        if is_tracing_state():
            padding_value = self.padding_value.cpu().detach().numpy()[0]
            output = Conv2dFixpoint.apply(input, self.weight, self.bias, 
                                        self.dilation, 
                                        self.groups,
                                        self.kernel_size,
                                        self.padding,
                                        self.stride,
                                        self.is_relu,
                                        padding_value,
                                        self.weight_bitwidth)
        else:
            output = self._forward_conv(input)
            output = (output + 0.5).floor()
        return output

class Conv2dFixpoint(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias, 
                    dilations,
                    group,
                    kernel_shapes,
                    pads,
                    strides,
                    is_relu,
                    padding_value,
                    weight_bitwidth):
        return F.conv2d(x, weight, bias, stride=strides, padding=pads, dilation=dilations, groups=group)

    @staticmethod
    def symbolic(g, x, weight, bias, 
                    dilations,
                    group,
                    kernel_shapes,
                    pads,
                    strides,
                    is_relu,
                    padding_value,
                    weight_bitwidth):
        return g.op("::Conv2dFixpoint", x, weight, bias, 
                                                    dilations_i=dilations,
                                                    group_i=group,
                                                    kernel_shapes_i=kernel_shapes,
                                                    pads_i=pads,
                                                    strides_i=strides,
                                                    is_relu_i=is_relu,
                                                    padding_value_i=padding_value,
                                                    weight_bitwidth_i=weight_bitwidth)


class ConvReLU2d(_FusedModule):
    r"""This is a sequential container which calls the Conv2d and ReLU modules.
    During quantization this will be replaced with the corresponding fused module."""
    def __init__(self, conv, relu):
        super().__init__(conv, relu)

    def forward(self, input):
        output = super().forward(input)
        output = output.round()
        return output
