#!/usr/bin/env python
# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
from torch.nn import Conv2d
from torch.nn import ConvTranspose2d

def conv_with_kaiming_uniform(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1):
    conv = Conv2d(
        in_channels, 
        out_channels, 
        kernel_size=kernel_size, 
        stride=stride, 
        padding=padding, 
        dilation=dilation, 
        bias= True
        )
        # Caffe2 implementation uses XavierFill, which in fact
        # corresponds to kaiming_uniform_ in PyTorch
    nn.init.kaiming_uniform_(conv.weight, a=1)
    nn.init.constant_(conv.bias, 0)
    return conv


def convtrans_with_kaiming_uniform(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1):
    conv = ConvTranspose2d(
        in_channels, 
        out_channels, 
        kernel_size=kernel_size, 
        stride=stride, 
        padding=padding, 
        dilation=dilation, 
        bias= True
        )
        # Caffe2 implementation uses XavierFill, which in fact
        # corresponds to kaiming_uniform_ in PyTorch
    nn.init.kaiming_uniform_(conv.weight, a=1)
    nn.init.constant_(conv.bias, 0)
    return conv


class UpsampleCropBlock(nn.Module):
    def __init__(self, in_channels, out_channels, up_kernel_size, up_stride, up_padding, pixelshuffle=False):
        """
        Combines Conv2d, ConvTransposed2d and Cropping. Simulates the caffe2 crop layer in the forward function.
        Used for DRIU and HED. 
        
        Attributes
        ----------
            in_channels : number of channels of intermediate layer
            out_channels : number of output channels
            up_kernel_size : kernel size for transposed convolution
            up_stride : stride for transposed convolution
            up_padding : padding for transposed convolution
        """
        super().__init__()
        # NOTE: Kaiming init, replace with nn.Conv2d and nn.ConvTranspose2d to get original DRIU impl.
        self.conv = conv_with_kaiming_uniform(in_channels, out_channels, 3, 1, 1)
        if pixelshuffle:
            self.upconv = PixelShuffle_ICNR( out_channels, out_channels, scale = up_stride)
        else:
            self.upconv = convtrans_with_kaiming_uniform(out_channels, out_channels, up_kernel_size, up_stride, up_padding)        
        
        
    def forward(self, x, input_res):
        """
        Forward pass of UpsampleBlock. Upsampled feature maps are cropped to the resolution of the input image.
        Attributes
        ----------
        x : input channels
        input_res : tuple (h,w)    
            Resolution of the input image
        """
        img_h = input_res[0]
        img_w = input_res[1]
        x = self.conv(x)
        x = self.upconv(x)
        # determine center crop
        # height
        up_h = x.shape[2]
        h_crop = up_h - img_h
        h_s = h_crop//2
        h_e = up_h - (h_crop - h_s)
        # width
        up_w = x.shape[3]
        w_crop = up_w-img_w
        w_s = w_crop//2
        w_e = up_w - (w_crop - w_s)
        # perform crop 
        # needs explicit ranges for onnx export 
        x = x[:,:,h_s:h_e,w_s:w_e] # crop to input size 
        
        return x



def ifnone(a, b):
    "`a` if `a` is not None, otherwise `b`."
    return b if a is None else a

def icnr(x, scale=2, init=nn.init.kaiming_normal_):
    """
    https://docs.fast.ai/layers.html#PixelShuffle_ICNR
    ICNR init of `x`, with `scale` and `init` function.
    """
    ni,nf,h,w = x.shape
    ni2 = int(ni/(scale**2))
    k = init(torch.zeros([ni2,nf,h,w])).transpose(0, 1)
    k = k.contiguous().view(ni2, nf, -1)
    k = k.repeat(1, 1, scale**2)
    k = k.contiguous().view([nf,ni,h,w]).transpose(0, 1)
    x.data.copy_(k)

class PixelShuffle_ICNR(nn.Module):
    """
    https://docs.fast.ai/layers.html#PixelShuffle_ICNR 
    Upsample by `scale` from `ni` filters to `nf` (default `ni`), using `nn.PixelShuffle`, `icnr` init, and `weight_norm`.
    """
    def __init__(self, ni:int, nf:int=None, scale:int=2):
        super().__init__()
        nf = ifnone(nf, ni)
        self.conv = conv_with_kaiming_uniform(ni, nf*(scale**2), 1)
        icnr(self.conv.weight)
        self.shuf = nn.PixelShuffle(scale)
        # Blurring over (h*w) kernel
        # "Super-Resolution using Convolutional Neural Networks without Any Checkerboard Artifacts"
        # - https://arxiv.org/abs/1806.02658
        self.pad = nn.ReplicationPad2d((1,0,1,0))
        self.blur = nn.AvgPool2d(2, stride=1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self,x):
        x = self.shuf(self.relu(self.conv(x)))
        x = self.blur(self.pad(x))
        return x

class UnetBlock(nn.Module):
    def __init__(self, up_in_c, x_in_c, pixel_shuffle=False, middle_block=False):
        super().__init__()

        # middle block for VGG based U-Net
        if middle_block:
            up_out_c =  up_in_c
        else:
            up_out_c =  up_in_c // 2
        cat_channels = x_in_c + up_out_c
        inner_channels = cat_channels // 2
        
        if pixel_shuffle:
            self.upsample = PixelShuffle_ICNR( up_in_c, up_out_c )
        else:
            self.upsample = convtrans_with_kaiming_uniform( up_in_c, up_out_c, 2, 2)
        self.convtrans1 = convtrans_with_kaiming_uniform( cat_channels, inner_channels, 3, 1, 1)
        self.convtrans2 = convtrans_with_kaiming_uniform( inner_channels, inner_channels, 3, 1, 1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, up_in, x_in):
        up_out = self.upsample(up_in)
        cat_x = torch.cat([up_out, x_in] , dim=1)
        x = self.relu(self.convtrans1(cat_x))
        x = self.relu(self.convtrans2(x))
        return x