# python3.7
"""Contains the implementation of discriminator described in StyleGAN2.

Compared to that of StyleGAN, the discriminator in StyleGAN2 mainly adds skip
connections, increases model size and disables progressive growth. This script
ONLY supports config F in the original paper.

Paper: https://arxiv.org/pdf/1912.04958.pdf

Official TensorFlow implementation: https://github.com/NVlabs/stylegan2
"""

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

__all__ = ['StyleGAN2Discriminator']

# Resolutions allowed.
_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]

# Initial resolution.
_INIT_RES = 4

# Architectures allowed.
_ARCHITECTURES_ALLOWED = ['resnet', 'skip', 'origin']

# Default gain factor for weight scaling.
_WSCALE_GAIN = 1.0


class StyleGAN2Discriminator(nn.Module):
    """Defines the discriminator network in StyleGAN2.

    NOTE: The discriminator takes images with `RGB` channel order and pixel
    range [-1, 1] as inputs.

    Settings for the network:

    (1) resolution: The resolution of the input image.
    (2) image_channels: Number of channels of the input image. (default: 3)
    (3) label_size: Size of the additional label for conditional generation.
        (default: 0)
    (4) architecture: Type of architecture. Support `origin`, `skip`, and
        `resnet`. (default: `resnet`)
    (5) use_wscale: Whether to use weight scaling. (default: True)
    (6) minibatch_std_group_size: Group size for the minibatch standard
        deviation layer. 0 means disable. (default: 4)
    (7) minibatch_std_channels: Number of new channels after the minibatch
        standard deviation layer. (default: 1)
    (8) fmaps_base: Factor to control number of feature maps for each layer.
        (default: 32 << 10)
    (9) fmaps_max: Maximum number of feature maps in each layer. (default: 512)
    """

    def __init__(self,
                 resolution,
                 image_channels=3,
                 label_size=0,
                 architecture='resnet',
                 use_wscale=True,
                 minibatch_std_group_size=4,
                 minibatch_std_channels=1,
                 fmaps_base=32 << 10,
                 fmaps_max=512):
        """Initializes with basic settings.

        Raises:
            ValueError: If the `resolution` is not supported, or `architecture`
                is not supported.
        """
        super().__init__()

        if resolution not in _RESOLUTIONS_ALLOWED:
            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')
        if architecture not in _ARCHITECTURES_ALLOWED:
            raise ValueError(f'Invalid architecture: `{architecture}`!\n'
                             f'Architectures allowed: '
                             f'{_ARCHITECTURES_ALLOWED}.')

        self.init_res = _INIT_RES
        self.init_res_log2 = int(np.log2(self.init_res))
        self.resolution = resolution
        self.final_res_log2 = int(np.log2(self.resolution))
        self.image_channels = image_channels
        self.label_size = label_size
        self.architecture = architecture
        self.use_wscale = use_wscale
        self.minibatch_std_group_size = minibatch_std_group_size
        self.minibatch_std_channels = minibatch_std_channels
        self.fmaps_base = fmaps_base
        self.fmaps_max = fmaps_max

        self.pth_to_tf_var_mapping = {}
        for res_log2 in range(self.final_res_log2, self.init_res_log2 - 1, -1):
            res = 2 ** res_log2
            block_idx = self.final_res_log2 - res_log2

            # Input convolution layer for each resolution (if needed).
            if res_log2 == self.final_res_log2 or self.architecture == 'skip':
                self.add_module(
                    f'input{block_idx}',
                    ConvBlock(in_channels=self.image_channels,
                              out_channels=self.get_nf(res),
                              kernel_size=1,
                              use_wscale=self.use_wscale))
                self.pth_to_tf_var_mapping[f'input{block_idx}.weight'] = (
                    f'{res}x{res}/FromRGB/weight')
                self.pth_to_tf_var_mapping[f'input{block_idx}.bias'] = (
                    f'{res}x{res}/FromRGB/bias')

            # Convolution block for each resolution (except the last one).
            if res != self.init_res:
                self.add_module(
                    f'layer{2 * block_idx}',
                    ConvBlock(in_channels=self.get_nf(res),
                              out_channels=self.get_nf(res),
                              use_wscale=self.use_wscale))
                tf_layer0_name = 'Conv0'
                self.add_module(
                    f'layer{2 * block_idx + 1}',
                    ConvBlock(in_channels=self.get_nf(res),
                              out_channels=self.get_nf(res // 2),
                              scale_factor=2,
                              use_wscale=self.use_wscale))
                tf_layer1_name = 'Conv1_down'

                if self.architecture == 'resnet':
                    layer_name = f'skip_layer{block_idx}'
                    self.add_module(
                        layer_name,
                        ConvBlock(in_channels=self.get_nf(res),
                                  out_channels=self.get_nf(res // 2),
                                  kernel_size=1,
                                  add_bias=False,
                                  scale_factor=2,
                                  use_wscale=self.use_wscale,
                                  activation_type='linear'))
                    self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
                        f'{res}x{res}/Skip/weight')

            # Convolution block for last resolution.
            else:
                self.add_module(
                    f'layer{2 * block_idx}',
                    ConvBlock(in_channels=self.get_nf(res),
                              out_channels=self.get_nf(res),
                              use_wscale=self.use_wscale,
                              minibatch_std_group_size=minibatch_std_group_size,
                              minibatch_std_channels=minibatch_std_channels))
                tf_layer0_name = 'Conv'
                self.add_module(
                    f'layer{2 * block_idx + 1}',
                    DenseBlock(in_channels=self.get_nf(res) * res * res,
                               out_channels=self.get_nf(res // 2),
                               use_wscale=self.use_wscale))
                tf_layer1_name = 'Dense0'

            self.pth_to_tf_var_mapping[f'layer{2 * block_idx}.weight'] = (
                f'{res}x{res}/{tf_layer0_name}/weight')
            self.pth_to_tf_var_mapping[f'layer{2 * block_idx}.bias'] = (
                f'{res}x{res}/{tf_layer0_name}/bias')
            self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 1}.weight'] = (
                f'{res}x{res}/{tf_layer1_name}/weight')
            self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 1}.bias'] = (
                f'{res}x{res}/{tf_layer1_name}/bias')

            # Final dense block.
            self.add_module(
                f'layer{2 * block_idx + 2}',
                DenseBlock(in_channels=self.get_nf(res // 2),
                           out_channels=max(self.label_size, 1),
                           use_wscale=self.use_wscale,
                           activation_type='linear'))
            self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 2}.weight'] = (
                f'Output/weight')
            self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 2}.bias'] = (
                f'Output/bias')

        if self.architecture == 'skip':
            self.downsample = DownsamplingLayer()

    def get_nf(self, res):
        """Gets number of feature maps according to current resolution."""
        return min(self.fmaps_base // res, self.fmaps_max)

    def forward(self, image, label=None, **_unused_kwargs):
        expected_shape = (self.image_channels, self.resolution, self.resolution)
        if image.ndim != 4 or image.shape[1:] != expected_shape:
            raise ValueError(f'The input tensor should be with shape '
                             f'[batch_size, channel, height, width], where '
                             f'`channel` equals to {self.image_channels}, '
                             f'`height`, `width` equal to {self.resolution}!\n'
                             f'But `{image.shape}` is received!')
        if self.label_size:
            if label is None:
                raise ValueError(f'Model requires an additional label '
                                 f'(with size {self.label_size}) as inputs, '
                                 f'but no label is received!')
            batch_size = image.shape[0]
            if label.ndim != 2 or label.shape != (batch_size, self.label_size):
                raise ValueError(f'Input label should be with shape '
                                 f'[batch_size, label_size], where '
                                 f'`batch_size` equals to that of '
                                 f'images ({image.shape[0]}) and '
                                 f'`label_size` equals to {self.label_size}!\n'
                                 f'But `{label.shape}` is received!')

        x = self.input0(image)
        for res_log2 in range(self.final_res_log2, self.init_res_log2 - 1, -1):
            block_idx = self.final_res_log2 - res_log2
            if self.architecture == 'skip' and block_idx > 0:
                image = self.downsample(image)
                x = x + self.__getattr__(f'input{block_idx}')(image)
            if self.architecture == 'resnet' and res_log2 != self.init_res_log2:
                residual = self.__getattr__(f'skip_layer{block_idx}')(x)
            x = self.__getattr__(f'layer{2 * block_idx}')(x)
            x = self.__getattr__(f'layer{2 * block_idx + 1}')(x)
            if self.architecture == 'resnet' and res_log2 != self.init_res_log2:
                x = (x + residual) / np.sqrt(2.0)
        x = self.__getattr__(f'layer{2 * block_idx + 2}')(x)

        if self.label_size:
            x = torch.sum(x * label, dim=1, keepdim=True)
        return x


class MiniBatchSTDLayer(nn.Module):
    """Implements the minibatch standard deviation layer."""

    def __init__(self, group_size=4, new_channels=1, epsilon=1e-8):
        super().__init__()
        self.group_size = group_size
        self.new_channels = new_channels
        self.epsilon = epsilon

    def forward(self, x):
        if self.group_size <= 1:
            return x
        ng = min(self.group_size, x.shape[0])
        nc = self.new_channels
        temp_c = x.shape[1] // nc                               # [NCHW]
        y = x.view(ng, -1, nc, temp_c, x.shape[2], x.shape[3])  # [GMncHW]
        y = y - torch.mean(y, dim=0, keepdim=True)              # [GMncHW]
        y = torch.mean(y ** 2, dim=0)                           # [MncHW]
        y = torch.sqrt(y + self.epsilon)                        # [MncHW]
        y = torch.mean(y, dim=[2, 3, 4], keepdim=True)          # [Mn111]
        y = torch.mean(y, dim=2)                                # [Mn11]
        y = y.repeat(ng, 1, x.shape[2], x.shape[3])             # [NnHW]
        return torch.cat([x, y], dim=1)


class DownsamplingLayer(nn.Module):
    """Implements the downsampling layer.

    This layer can also be used as filtering by setting `scale_factor` as 1.
    """

    def __init__(self, scale_factor=2, kernel=(1, 3, 3, 1), extra_padding=0):
        super().__init__()
        assert scale_factor >= 1
        self.scale_factor = scale_factor

        if extra_padding != 0:
            assert scale_factor == 1

        if kernel is None:
            kernel = np.ones((scale_factor), dtype=np.float32)
        else:
            kernel = np.array(kernel, dtype=np.float32)
        assert kernel.ndim == 1
        kernel = np.outer(kernel, kernel)
        kernel = kernel / np.sum(kernel)
        assert kernel.ndim == 2
        assert kernel.shape[0] == kernel.shape[1]
        kernel = kernel[np.newaxis, np.newaxis]
        self.register_buffer('kernel', torch.from_numpy(kernel))
        self.kernel = self.kernel.flip(0, 1)
        padding = kernel.shape[2] - scale_factor + extra_padding
        self.padding = ((padding + 1) // 2, padding // 2,
                        (padding + 1) // 2, padding // 2)

    def forward(self, x):
        assert x.ndim == 4
        channels = x.shape[1]
        x = x.view(-1, 1, x.shape[2], x.shape[3])
        x = F.pad(x, self.padding, mode='constant', value=0)
        x = F.conv2d(x, self.kernel, stride=self.scale_factor)
        x = x.view(-1, channels, x.shape[2], x.shape[3])
        return x


class ConvBlock(nn.Module):
    """Implements the convolutional block.

    Basically, this block executes minibatch standard deviation layer (if
    needed), filtering layer (if needed), convolutional layer, and activation
    layer in sequence.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 add_bias=True,
                 scale_factor=1,
                 filtering_kernel=(1, 3, 3, 1),
                 use_wscale=True,
                 wscale_gain=_WSCALE_GAIN,
                 lr_mul=1.0,
                 activation_type='lrelu',
                 minibatch_std_group_size=0,
                 minibatch_std_channels=1):
        """Initializes with block settings.

        Args:
            in_channels: Number of channels of the input tensor.
            out_channels: Number of channels of the output tensor.
            kernel_size: Size of the convolutional kernels. (default: 3)
            add_bias: Whether to add bias onto the convolutional result.
                (default: True)
            scale_factor: Scale factor for downsampling. `1` means skip
                downsampling. (default: 1)
            filtering_kernel: Kernel used for filtering before downsampling.
                (default: (1, 3, 3, 1))
            use_wscale: Whether to use weight scaling. (default: True)
            wscale_gain: Gain factor for weight scaling. (default: _WSCALE_GAIN)
            lr_mul: Learning multiplier for both weight and bias. (default: 1.0)
            activation_type: Type of activation. Support `linear` and `lrelu`.
                (default: `lrelu`)
            minibatch_std_group_size: Group size for the minibatch standard
                deviation layer. 0 means disable. (default: 0)
            minibatch_std_channels: Number of new channels after the minibatch
                standard deviation layer. (default: 1)

        Raises:
            NotImplementedError: If the `activation_type` is not supported.
        """
        super().__init__()

        if minibatch_std_group_size > 1:
            in_channels = in_channels + minibatch_std_channels
            self.mbstd = MiniBatchSTDLayer(group_size=minibatch_std_group_size,
                                           new_channels=minibatch_std_channels)
        else:
            self.mbstd = nn.Identity()

        if scale_factor > 1:
            extra_padding = kernel_size - scale_factor
            self.filter = DownsamplingLayer(scale_factor=1,
                                            kernel=filtering_kernel,
                                            extra_padding=extra_padding)
            self.stride = scale_factor
            self.padding = 0  # Padding is done in `DownsamplingLayer`.
        else:
            self.filter = nn.Identity()
            assert kernel_size % 2 == 1
            self.stride = 1
            self.padding = kernel_size // 2

        weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
        fan_in = kernel_size * kernel_size * in_channels
        wscale = wscale_gain / np.sqrt(fan_in)
        if use_wscale:
            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
            self.wscale = wscale * lr_mul
        else:
            self.weight = nn.Parameter(
                torch.randn(*weight_shape) * wscale / lr_mul)
            self.wscale = lr_mul

        if add_bias:
            self.bias = nn.Parameter(torch.zeros(out_channels))
        else:
            self.bias = None
        self.bscale = lr_mul

        if activation_type == 'linear':
            self.activate = nn.Identity()
            self.activate_scale = 1.0
        elif activation_type == 'lrelu':
            self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True)
            self.activate_scale = np.sqrt(2.0)
        else:
            raise NotImplementedError(f'Not implemented activation function: '
                                      f'`{activation_type}`!')

    def forward(self, x):
        x = self.mbstd(x)
        x = self.filter(x)
        weight = self.weight * self.wscale
        bias = self.bias * self.bscale if self.bias is not None else None
        x = F.conv2d(x,
                     weight=weight,
                     bias=bias,
                     stride=self.stride,
                     padding=self.padding)
        x = self.activate(x) * self.activate_scale
        return x


class DenseBlock(nn.Module):
    """Implements the dense block.

    Basically, this block executes fully-connected layer and activation layer.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 add_bias=True,
                 use_wscale=True,
                 wscale_gain=_WSCALE_GAIN,
                 lr_mul=1.0,
                 activation_type='lrelu'):
        """Initializes with block settings.

        Args:
            in_channels: Number of channels of the input tensor.
            out_channels: Number of channels of the output tensor.
            add_bias: Whether to add bias onto the fully-connected result.
                (default: True)
            use_wscale: Whether to use weight scaling. (default: True)
            wscale_gain: Gain factor for weight scaling. (default: _WSCALE_GAIN)
            lr_mul: Learning multiplier for both weight and bias. (default: 1.0)
            activation_type: Type of activation. Support `linear` and `lrelu`.
                (default: `lrelu`)

        Raises:
            NotImplementedError: If the `activation_type` is not supported.
        """
        super().__init__()
        weight_shape = (out_channels, in_channels)
        wscale = wscale_gain / np.sqrt(in_channels)
        if use_wscale:
            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
            self.wscale = wscale * lr_mul
        else:
            self.weight = nn.Parameter(
                torch.randn(*weight_shape) * wscale / lr_mul)
            self.wscale = lr_mul

        if add_bias:
            self.bias = nn.Parameter(torch.zeros(out_channels))
        else:
            self.bias = None
        self.bscale = lr_mul

        if activation_type == 'linear':
            self.activate = nn.Identity()
            self.activate_scale = 1.0
        elif activation_type == 'lrelu':
            self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True)
            self.activate_scale = np.sqrt(2.0)
        else:
            raise NotImplementedError(f'Not implemented activation function: '
                                      f'`{activation_type}`!')

    def forward(self, x):
        if x.ndim != 2:
            x = x.view(x.shape[0], -1)
        bias = self.bias * self.bscale if self.bias is not None else None
        x = F.linear(x, weight=self.weight * self.wscale, bias=bias)
        x = self.activate(x) * self.activate_scale
        return x