# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from collections.abc import Sequence

import torch
import torch.nn as nn

from monai.networks.blocks import Convolution
from monai.networks.layers import Act
from monai.networks.utils import normal_init


class MultiScalePatchDiscriminator(nn.Sequential):
    """
    Multi-scale Patch-GAN discriminator based on Pix2PixHD:
    High-Resolution Image Synthesis and Semantic Manipulation with Conditional GANs (https://arxiv.org/abs/1711.11585)

    The Multi-scale discriminator made up of several PatchGAN discriminators, that process the images
    at different spatial scales.

    Args:
        num_d: number of discriminators
        num_layers_d: number of Convolution layers (Conv + activation + normalisation + [dropout]) in the first
            discriminator. Each subsequent discriminator has one additional layer, meaning the output size is halved.
        spatial_dims: number of spatial dimensions (1D, 2D etc.)
        channels: number of filters in the first convolutional layer (doubled for each subsequent layer)
        in_channels: number of input channels
        out_channels: number of output channels in each discriminator
        kernel_size: kernel size of the convolution layers
        activation: activation layer type
        norm: normalisation type
        bias: introduction of layer bias
        dropout: probability of dropout applied, defaults to 0.
        minimum_size_im: minimum spatial size of the input image. Introduced to make sure the architecture
            requested isn't going to downsample the input image beyond value of 1.
        last_conv_kernel_size: kernel size of the last convolutional layer.
    """

    def __init__(
        self,
        num_d: int,
        num_layers_d: int,
        spatial_dims: int,
        channels: int,
        in_channels: int,
        out_channels: int = 1,
        kernel_size: int = 4,
        activation: str | tuple = (Act.LEAKYRELU, {"negative_slope": 0.2}),
        norm: str | tuple = "BATCH",
        bias: bool = False,
        dropout: float | tuple = 0.0,
        minimum_size_im: int = 256,
        last_conv_kernel_size: int = 1,
    ) -> None:
        super().__init__()
        self.num_d = num_d
        self.num_layers_d = num_layers_d
        self.num_channels = channels
        self.padding = tuple([int((kernel_size - 1) / 2)] * spatial_dims)
        for i_ in range(self.num_d):
            num_layers_d_i = self.num_layers_d * (i_ + 1)
            output_size = float(minimum_size_im) / (2**num_layers_d_i)
            if output_size < 1:
                raise AssertionError(
                    f"Your image size is too small to take in up to {i_} discriminators with num_layers = {num_layers_d_i}."
                    "Please reduce num_layers, reduce num_D or enter bigger images."
                )
            subnet_d = PatchDiscriminator(
                spatial_dims=spatial_dims,
                channels=self.num_channels,
                in_channels=in_channels,
                out_channels=out_channels,
                num_layers_d=num_layers_d_i,
                kernel_size=kernel_size,
                activation=activation,
                norm=norm,
                bias=bias,
                padding=self.padding,
                dropout=dropout,
                last_conv_kernel_size=last_conv_kernel_size,
            )

            self.add_module(f"discriminator_{i_}", subnet_d)

    def forward(self, i: torch.Tensor) -> tuple[list[torch.Tensor], list[list[torch.Tensor]]]:
        """
        Args:
            i: Input tensor

        Returns:
            list of outputs and another list of lists with the intermediate features
            of each discriminator.
        """

        out: list[torch.Tensor] = []
        intermediate_features: list[list[torch.Tensor]] = []
        for disc in self.children():
            out_d: list[torch.Tensor] = disc(i)
            out.append(out_d[-1])
            intermediate_features.append(out_d[:-1])

        return out, intermediate_features


class PatchDiscriminator(nn.Sequential):
    """
    Patch-GAN discriminator based on Pix2PixHD:
    High-Resolution Image Synthesis and Semantic Manipulation with Conditional GANs (https://arxiv.org/abs/1711.11585)


    Args:
        spatial_dims: number of spatial dimensions (1D, 2D etc.)
        channels: number of filters in the first convolutional layer (doubled for each subsequent layer)
        in_channels: number of input channels
        out_channels: number of output channels
        num_layers_d: number of Convolution layers (Conv + activation + normalisation + [dropout]) in the discriminator.
        kernel_size: kernel size of the convolution layers
        act: activation type and arguments. Defaults to LeakyReLU.
        norm: feature normalization type and arguments. Defaults to batch norm.
        bias: whether to have a bias term in convolution blocks. Defaults to False.
        padding: padding to be applied to the convolutional layers
        dropout: proportion of dropout applied, defaults to 0.
        last_conv_kernel_size: kernel size of the last convolutional layer.
    """

    def __init__(
        self,
        spatial_dims: int,
        channels: int,
        in_channels: int,
        out_channels: int = 1,
        num_layers_d: int = 3,
        kernel_size: int = 4,
        activation: str | tuple = (Act.LEAKYRELU, {"negative_slope": 0.2}),
        norm: str | tuple = "BATCH",
        bias: bool = False,
        padding: int | Sequence[int] = 1,
        dropout: float | tuple = 0.0,
        last_conv_kernel_size: int | None = None,
    ) -> None:
        super().__init__()
        self.num_layers_d = num_layers_d
        self.num_channels = channels
        if last_conv_kernel_size is None:
            last_conv_kernel_size = kernel_size

        self.add_module(
            "initial_conv",
            Convolution(
                spatial_dims=spatial_dims,
                kernel_size=kernel_size,
                in_channels=in_channels,
                out_channels=channels,
                act=activation,
                bias=True,
                norm=None,
                dropout=dropout,
                padding=padding,
                strides=2,
            ),
        )

        input_channels = channels
        output_channels = channels * 2

        # Initial Layer
        for l_ in range(self.num_layers_d):
            if l_ == self.num_layers_d - 1:
                stride = 1
            else:
                stride = 2
            layer = Convolution(
                spatial_dims=spatial_dims,
                kernel_size=kernel_size,
                in_channels=input_channels,
                out_channels=output_channels,
                act=activation,
                bias=bias,
                norm=norm,
                dropout=dropout,
                padding=padding,
                strides=stride,
            )
            self.add_module(f"{l_}", layer)
            input_channels = output_channels
            output_channels = output_channels * 2

        # Final layer
        self.add_module(
            "final_conv",
            Convolution(
                spatial_dims=spatial_dims,
                kernel_size=last_conv_kernel_size,
                in_channels=input_channels,
                out_channels=out_channels,
                bias=True,
                conv_only=True,
                padding=int((last_conv_kernel_size - 1) / 2),
                dropout=0.0,
                strides=1,
            ),
        )

        self.apply(normal_init)

    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
        """
        Args:
            x: input tensor

        Returns:
            list of intermediate features, with the last element being the output.
        """
        out = [x]
        for submodel in self.children():
            intermediate_output = submodel(out[-1])
            out.append(intermediate_output)

        return out[1:]
