import math

import torch
import torch.nn as nn

from model.normalization import select_norm_layer
from model import MODEL


class GANImageBuffer(object):
    """This class implements an image buffer that stores previously
    generated images.
    This buffer allows us to update the discriminator using a history of
    generated images rather than the ones produced by the latest generator
    to reduce model oscillation.
    Args:
        buffer_size (int): The size of image buffer. If buffer_size = 0,
            no buffer will be created.
        buffer_ratio (float): The chance / possibility  to use the images
            previously stored in the buffer.
    """

    def __init__(self, buffer_size, buffer_ratio=0.5):
        self.buffer_size = buffer_size
        # create an empty buffer
        if self.buffer_size > 0:
            self.img_num = 0
            self.image_buffer = []
        self.buffer_ratio = buffer_ratio

    def query(self, images):
        """Query current image batch using a history of generated images.
        Args:
            images (Tensor): Current image batch without history information.
        """
        if self.buffer_size == 0:  # if the buffer size is 0, do nothing
            return images
        return_images = []
        for image in images:
            image = torch.unsqueeze(image.data, 0)
            # if the buffer is not full, keep inserting current images
            if self.img_num < self.buffer_size:
                self.img_num = self.img_num + 1
                self.image_buffer.append(image)
                return_images.append(image)
            else:
                use_buffer = torch.rand(1) < self.buffer_ratio
                # by self.buffer_ratio, the buffer will return a previously
                # stored image, and insert the current image into the buffer
                if use_buffer:
                    random_id = torch.randint(0, self.buffer_size, (1,)).item()
                    image_tmp = self.image_buffer[random_id].clone()
                    self.image_buffer[random_id] = image
                    return_images.append(image_tmp)
                # by (1 - self.buffer_ratio), the buffer will return the
                # current image
                else:
                    return_images.append(image)
        # collect all the images and return
        return_images = torch.cat(return_images, 0)
        return return_images


# based SPADE or pix2pixHD Discriminator
@MODEL.register_module("PatchDiscriminator")
class PatchDiscriminator(nn.Module):
    def __init__(self, in_channels, base_channels, num_conv=4, use_spectral=False, norm_type="IN",
                 need_intermediate_feature=False):
        super().__init__()
        self.need_intermediate_feature = need_intermediate_feature

        kernel_size = 4
        padding = math.ceil((kernel_size - 1.0) / 2)
        norm_layer = select_norm_layer(norm_type)
        use_bias = norm_type == "IN"
        padding_mode = "zeros"

        sequence = [nn.Sequential(
            nn.Conv2d(in_channels, base_channels, kernel_size, stride=2, padding=padding),
            nn.LeakyReLU(0.2, False)
        )]
        multiple_now = 1
        for i in range(1, num_conv):
            multiple_prev = multiple_now
            multiple_now = min(2 ** i, 2 ** 3)
            stride = 1 if i == num_conv - 1 else 2
            sequence.append(nn.Sequential(
                self.build_conv2d(use_spectral, base_channels * multiple_prev, base_channels * multiple_now,
                                  kernel_size, stride, padding, bias=use_bias, padding_mode=padding_mode),
                norm_layer(base_channels * multiple_now),
                nn.LeakyReLU(0.2, inplace=False),
            ))
        multiple_now = min(2 ** num_conv, 8)
        sequence.append(nn.Conv2d(base_channels * multiple_now, 1, kernel_size, stride=1, padding=padding,
                                  padding_mode=padding_mode))
        self.conv_blocks = nn.ModuleList(sequence)

    @staticmethod
    def build_conv2d(use_spectral, in_channels: int, out_channels: int, kernel_size, stride, padding,
                     bias=True, padding_mode: str = 'zeros'):
        conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias, padding_mode=padding_mode)
        if not use_spectral:
            return conv
        return nn.utils.spectral_norm(conv)

    def forward(self, x):
        if self.need_intermediate_feature:
            intermediate_feature = []
            for layer in self.conv_blocks:
                x = layer(x)
                intermediate_feature.append(x)
            return tuple(intermediate_feature)
        else:
            for layer in self.conv_blocks:
                x = layer(x)
            return x


@MODEL.register_module()
class ResidualBlock(nn.Module):
    def __init__(self, num_channels, padding_mode='reflect', norm_type="IN", use_bias=None):
        super(ResidualBlock, self).__init__()
        if use_bias is None:
            # Only for IN, use bias since it does not have affine parameters.
            use_bias = norm_type == "IN"
        norm_layer = select_norm_layer(norm_type)
        self.conv1 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode,
                               bias=use_bias)
        self.norm1 = norm_layer(num_channels)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode,
                               bias=use_bias)
        self.norm2 = norm_layer(num_channels)

    def forward(self, x):
        res = x
        x = self.relu1(self.norm1(self.conv1(x)))
        x = self.norm2(self.conv2(x))
        return x + res