import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.spectral_norm as SpectralNorm
from torch.autograd import Function


class BlurFunctionBackward(Function):

    @staticmethod
    def forward(ctx, grad_output, kernel, kernel_flip):
        ctx.save_for_backward(kernel, kernel_flip)
        grad_input = F.conv2d(grad_output, kernel_flip, padding=1, groups=grad_output.shape[1])
        return grad_input

    @staticmethod
    def backward(ctx, gradgrad_output):
        kernel, kernel_flip = ctx.saved_tensors
        grad_input = F.conv2d(gradgrad_output, kernel, padding=1, groups=gradgrad_output.shape[1])
        return grad_input, None, None


class BlurFunction(Function):

    @staticmethod
    def forward(ctx, x, kernel, kernel_flip):
        ctx.save_for_backward(kernel, kernel_flip)
        output = F.conv2d(x, kernel, padding=1, groups=x.shape[1])
        return output

    @staticmethod
    def backward(ctx, grad_output):
        kernel, kernel_flip = ctx.saved_tensors
        grad_input = BlurFunctionBackward.apply(grad_output, kernel, kernel_flip)
        return grad_input, None, None


blur = BlurFunction.apply


class Blur(nn.Module):

    def __init__(self, channel):
        super().__init__()
        kernel = torch.tensor([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=torch.float32)
        kernel = kernel.view(1, 1, 3, 3)
        kernel = kernel / kernel.sum()
        kernel_flip = torch.flip(kernel, [2, 3])

        self.kernel = kernel.repeat(channel, 1, 1, 1)
        self.kernel_flip = kernel_flip.repeat(channel, 1, 1, 1)

    def forward(self, x):
        return blur(x, self.kernel.type_as(x), self.kernel_flip.type_as(x))


def calc_mean_std(feat, eps=1e-5):
    """Calculate mean and std for adaptive_instance_normalization.

    Args:
        feat (Tensor): 4D tensor.
        eps (float): A small value added to the variance to avoid
            divide-by-zero. Default: 1e-5.
    """
    size = feat.size()
    assert len(size) == 4, 'The input feature should be 4D tensor.'
    n, c = size[:2]
    feat_var = feat.view(n, c, -1).var(dim=2) + eps
    feat_std = feat_var.sqrt().view(n, c, 1, 1)
    feat_mean = feat.view(n, c, -1).mean(dim=2).view(n, c, 1, 1)
    return feat_mean, feat_std


def adaptive_instance_normalization(content_feat, style_feat):
    """Adaptive instance normalization.

    Adjust the reference features to have the similar color and illuminations
    as those in the degradate features.

    Args:
        content_feat (Tensor): The reference feature.
        style_feat (Tensor): The degradate features.
    """
    size = content_feat.size()
    style_mean, style_std = calc_mean_std(style_feat)
    content_mean, content_std = calc_mean_std(content_feat)
    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
    return normalized_feat * style_std.expand(size) + style_mean.expand(size)


def AttentionBlock(in_channel):
    return nn.Sequential(
        SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1)), nn.LeakyReLU(0.2, True),
        SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1)))


def conv_block(in_channels, out_channels, kernel_size=3, stride=1, dilation=1, bias=True):
    """Conv block used in MSDilationBlock."""

    return nn.Sequential(
        SpectralNorm(
            nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size=kernel_size,
                stride=stride,
                dilation=dilation,
                padding=((kernel_size - 1) // 2) * dilation,
                bias=bias)),
        nn.LeakyReLU(0.2),
        SpectralNorm(
            nn.Conv2d(
                out_channels,
                out_channels,
                kernel_size=kernel_size,
                stride=stride,
                dilation=dilation,
                padding=((kernel_size - 1) // 2) * dilation,
                bias=bias)),
    )


class MSDilationBlock(nn.Module):
    """Multi-scale dilation block."""

    def __init__(self, in_channels, kernel_size=3, dilation=(1, 1, 1, 1), bias=True):
        super(MSDilationBlock, self).__init__()

        self.conv_blocks = nn.ModuleList()
        for i in range(4):
            self.conv_blocks.append(conv_block(in_channels, in_channels, kernel_size, dilation=dilation[i], bias=bias))
        self.conv_fusion = SpectralNorm(
            nn.Conv2d(
                in_channels * 4,
                in_channels,
                kernel_size=kernel_size,
                stride=1,
                padding=(kernel_size - 1) // 2,
                bias=bias))

    def forward(self, x):
        out = []
        for i in range(4):
            out.append(self.conv_blocks[i](x))
        out = torch.cat(out, 1)
        out = self.conv_fusion(out) + x
        return out


class UpResBlock(nn.Module):

    def __init__(self, in_channel):
        super(UpResBlock, self).__init__()
        self.body = nn.Sequential(
            nn.Conv2d(in_channel, in_channel, 3, 1, 1),
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(in_channel, in_channel, 3, 1, 1),
        )

    def forward(self, x):
        out = x + self.body(x)
        return out