import torch.nn as nn
import functools
from .registry import MODEL


def _select_norm_layer(norm_type):
    if norm_type == "BN":
        return functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
    elif norm_type == "IN":
        return functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
    elif norm_type == "NONE":
        return lambda x: nn.Identity()
    else:
        raise NotImplemented(f'normalization layer {norm_type} is not found')


@MODEL.register_module()
class ResidualBlock(nn.Module):
    def __init__(self, num_channels, padding_mode='reflect', norm_type="IN", use_dropout=False):
        super(ResidualBlock, self).__init__()

        # Only for IN, use bias since it does not have affine parameters.
        use_bias = norm_type == "IN"
        norm_layer = _select_norm_layer(norm_type)
        models = [nn.Sequential(
            nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias),
            norm_layer(num_channels),
            nn.ReLU(inplace=True),
        )]
        if use_dropout:
            models.append(nn.Dropout(0.5))
        models.append(nn.Sequential(
            nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias),
            norm_layer(num_channels),
        ))
        self.block = nn.Sequential(*models)

    def forward(self, x):
        return x + self.block(x)


@MODEL.register_module()
class ResGenerator(nn.Module):
    def __init__(self, in_channels, out_channels, base_channels=64, num_blocks=9, padding_mode='reflect',
                 norm_type="IN", use_dropout=False):
        super(ResGenerator, self).__init__()
        assert num_blocks >= 0, f'Number of residual blocks must be non-negative, but got {num_blocks}.'
        norm_layer = _select_norm_layer(norm_type)
        use_bias = norm_type == "IN"

        self.start_conv = nn.Sequential(
            nn.Conv2d(in_channels, base_channels, kernel_size=7, stride=1, padding_mode=padding_mode, padding=3,
                      bias=use_bias),
            norm_layer(num_features=base_channels),
            nn.ReLU(inplace=True)
        )

        # down sampling
        submodules = []
        num_down_sampling = 2
        for i in range(num_down_sampling):
            multiple = 2 ** i
            submodules += [
                nn.Conv2d(in_channels=base_channels * multiple, out_channels=base_channels * multiple * 2,
                          kernel_size=3,
                          stride=2, padding=1, bias=use_bias),
                norm_layer(num_features=base_channels * multiple * 2),
                nn.ReLU(inplace=True)
            ]
        self.encoder = nn.Sequential(*submodules)

        res_block_channels = num_down_sampling ** 2 * base_channels
        self.res_blocks = nn.ModuleList(
            [ResidualBlock(res_block_channels, padding_mode, norm_type, use_dropout=use_dropout) for _ in
             range(num_blocks)])

        # up sampling
        submodules = []
        for i in range(num_down_sampling):
            multiple = 2 ** (num_down_sampling - i)
            submodules += [
                nn.ConvTranspose2d(base_channels * multiple, base_channels * multiple // 2, kernel_size=3, stride=2,
                                   padding=1, output_padding=1, bias=use_bias),
                norm_layer(num_features=base_channels * multiple // 2),
                nn.ReLU(inplace=True),
            ]
        self.decoder = nn.Sequential(*submodules)

        self.end_conv = nn.Sequential(
            nn.Conv2d(base_channels, out_channels, kernel_size=7, padding=3, padding_mode=padding_mode),
            nn.Tanh()
        )

    def forward(self, x):
        x = self.encoder(self.start_conv(x))
        for rb in self.res_blocks:
            x = rb(x)
        return self.end_conv(self.decoder(x))


@MODEL.register_module()
class PatchDiscriminator(nn.Module):
    def __init__(self, in_channels, base_channels=64, num_conv=3, norm_type="BN"):
        super(PatchDiscriminator, self).__init__()
        assert num_conv >= 0, f'Number of conv blocks must be non-negative, but got {num_conv}.'
        norm_layer = _select_norm_layer(norm_type)
        use_bias = norm_type == "IN"

        kernel_size = 4
        padding = 1
        sequence = [
            nn.Conv2d(in_channels, base_channels, kernel_size=kernel_size, stride=2, padding=padding),
            nn.LeakyReLU(0.2, inplace=True),
        ]

        # stacked intermediate layers,
        # gradually increasing the number of filters
        multiple_now = 1
        for n in range(1, num_conv):
            multiple_prev = multiple_now
            multiple_now = min(2 ** n, 8)
            sequence += [
                nn.Conv2d(base_channels * multiple_prev, base_channels * multiple_now, kernel_size=kernel_size,
                          padding=padding, stride=2, bias=use_bias),
                norm_layer(base_channels * multiple_now),
                nn.LeakyReLU(0.2, inplace=True)
            ]
        multiple_prev = multiple_now
        multiple_now = min(2 ** num_conv, 8)
        sequence += [
            nn.Conv2d(base_channels * multiple_prev, base_channels * multiple_now, kernel_size, stride=1,
                      padding=padding, bias=use_bias),
            norm_layer(base_channels * multiple_now),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(base_channels * multiple_now, 1, kernel_size, stride=1, padding=padding)
        ]
        self.model = nn.Sequential(*sequence)

    def forward(self, x):
        return self.model(x)