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


class ConvTranspose2d(nn.ConvTranspose2d):
    # only support symetric quantization for activation and weight
    def __init__(self,  *args, is_relu=1, padding_value=None, 
                weight_bitwidth=4,
                weight_half_range=False,
                input_bitwidth=8,
                input_half_range=False,
                output_bitwidth=32,
                bias_bit=32,**kwargs) -> None:
        super().__init__(*args, **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_bit = bias_bit
        self.padding_value = nn.Parameter(padding_value.to(torch.float32).round().to(torch.int32), requires_grad=False)


    def _add_padding_value(self, input):
        b, c, h, w = input.shape
        pad_num   = self.kernel_size[0] // 2
        tmp_input = torch.full([b, 3*c, h, w], self.padding_value.data.item()).to(input.device)
        pad_input = torch.full([b, c, 2 * h - 1 + 2 * pad_num, 2 * w - 1 + 2 * pad_num], self.padding_value.data.item()).to(input.device).to(input.dtype)
        new_input = torch.cat([input, tmp_input], dim=1)
        tmp_input = new_input.reshape(b, 4, -1, h, w).permute(0, 2, 1, 3, 4).\
            reshape(b, c, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c, 2*h, 2*w)
        pad_input[..., pad_num:2 * h - 1 + pad_num, pad_num:2 * w - 1 + pad_num] = tmp_input[..., :(2*h - 1), :(2*w - 1)]
        return pad_input


    def forward(self, input):
        if is_tracing_state():
            padding_value = self.padding_value.cpu().detach().numpy()[0]
            output = Conv2dTransposeFixpoint.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:
            if self.kernel_size[0] == 2 and self.stride[0] == 2:
                input = self._add_padding_value(input)
                # flip weight
                rot_weight = torch.rot90(self.weight.permute(1, 0, 2, 3), k=2, dims=[2, 3])
                output = F.conv2d(input, rot_weight, self.bias, stride=1, padding=0)
                output = (output + 0.5).floor()

            else:
                output = super().forward(input)
                output = (output + 0.5).floor()

        return output

class Conv2dTransposeFixpoint(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias, 
                    dilations,
                    group,
                    kernel_shapes,
                    pads,
                    strides,
                    is_relu,
                    padding_value,
                    weight_bitwidth):
        return F.conv_transpose2d(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("::Conv2dTransposeFixpoint", 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 ConvTransposeReLU2d(_FusedModule):
    r"""This is a sequential container which calls the ConvTranspose2d and ReLU modules.
    During quantization this will be replaced with the corresponding fused module."""
    def __init__(self, deconv, relu):
        super().__init__(deconv, relu)

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