from typing import Tuple, List
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms


class ConvolutionBlock(nn.Module):
    """
    A block with two convolutional layers followed by batch norm
    layers and LeakyReLU as activation function.
    """

    def __init__(self, input_channel: int, output_channel: int):
        super(ConvolutionBlock, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(input_channel, output_channel, kernel_size=3, padding=1),
            nn.BatchNorm2d(output_channel),
            nn.LeakyReLU(),
            nn.Conv2d(output_channel, output_channel, kernel_size=3, padding=1),
            nn.BatchNorm2d(output_channel),
            nn.LeakyReLU()
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.conv(x)


class DownSamplingBlock(nn.Module):
    """
    A down-sampling block with a max pooling layer followed by
    a convolutional block.
    """

    def __init__(self, input_channel: int, output_channel: int):
        super(DownSamplingBlock, self).__init__()
        self.down_sampler = nn.Sequential(
            nn.MaxPool2d(2),
            ConvolutionBlock(input_channel, output_channel)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_sampler(x)


class UpSamplingBlock(nn.Module):
    """
    An up-sampling block with a 2d transpose convolution layer
    for up-sampling followed by a ConvolutionBlock.
    """

    def __init__(self, input_channel1, input_channel2, output_channel):
        super(UpSamplingBlock, self).__init__()
        self.up = nn.ConvTranspose2d(input_channel1, input_channel2, kernel_size=2, stride=2)
        self.conv = ConvolutionBlock(input_channel2 * 2, output_channel)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        x1 = self.up(x1)
        height = x1.size()[2]  # Finding dimensions of the up-sampled image for cropping
        width = x1.size()[3]
        x2 = transforms.CenterCrop((height, width))(x2)  # Cropping the image from the skip connection from the center
        x = torch.cat([x2, x1], dim=1)  # Concatenating both images together
        return self.conv(x)


class DownSampler(nn.Module):
    """
    The module used for down-sampling process consisted of
    multiple down-sampling blocks and one convolutional block
    at the first part.
    """

    def __init__(self, input_channel: int, scale_channels: list):
        super(DownSampler, self).__init__()
        self.input_channel = input_channel
        self.scale_channels = scale_channels

        self.conv = ConvolutionBlock(self.input_channel, self.scale_channels[0])
        self.down_sampler1 = DownSamplingBlock(self.scale_channels[0], self.scale_channels[1])
        self.down_sampler2 = DownSamplingBlock(self.scale_channels[1], self.scale_channels[2])
        self.down_sampler3 = DownSamplingBlock(self.scale_channels[2], self.scale_channels[3])
        self.down_sampler4 = DownSamplingBlock(self.scale_channels[3], self.scale_channels[4])

    def forward(self, x: torch.Tensor) -> List:
        features = self.conv(x)
        down_sampled1 = self.down_sampler1(features)
        down_sampled2 = self.down_sampler2(down_sampled1)
        down_sampled3 = self.down_sampler3(down_sampled2)
        down_sampled4 = self.down_sampler4(down_sampled3)

        return [features, down_sampled1, down_sampled2, down_sampled3, down_sampled4]


class UpSampler(nn.Module):
    """
    The module used for up-sampling process consisted of
    multiple up-sampling blocks and one convolutional layer
    at the last part for the probability map.
    """

    def __init__(self, scale_channels: list, class_num: int, shape: tuple):
        super(UpSampler, self).__init__()
        self.scale_channels = scale_channels
        self.class_num = class_num
        self.shape = shape

        self.up_sampler4 = UpSamplingBlock(scale_channels[4], scale_channels[3], scale_channels[3])
        self.up_sampler3 = UpSamplingBlock(scale_channels[3], scale_channels[2], scale_channels[2])
        self.up_sampler2 = UpSamplingBlock(scale_channels[2], scale_channels[1], scale_channels[1])
        self.up_sampler1 = UpSamplingBlock(scale_channels[1], scale_channels[0], scale_channels[0])

        self.probability_map_conv0 = nn.Sequential(
            nn.Conv2d(self.scale_channels[0], self.class_num, kernel_size=1)
        )
        self.probability_map_conv1 = nn.Sequential(
            nn.Conv2d(self.scale_channels[1], self.class_num, kernel_size=1)
        )
        self.probability_map_conv2 = nn.Sequential(
            nn.Conv2d(self.scale_channels[2], self.class_num, kernel_size=1)
        )
        self.probability_map_conv3 = nn.Sequential(
            nn.Conv2d(self.scale_channels[3], self.class_num, kernel_size=1)
        )

    def forward(self, down_outputs: list) -> Tuple:
        """
        NOTE: Numbers in front of the var names indicate the scale number.
        """
        
        down_outputs0 = down_outputs.pop(0)
        down_outputs1 = down_outputs.pop(0)
        down_outputs2 = down_outputs.pop(0)
        down_outputs3 = down_outputs.pop(0)
        down_outputs4 = down_outputs.pop(0)

        up_sampled4 = self.up_sampler4(down_outputs4, down_outputs3)
        up_sampled4 = F.interpolate(up_sampled4, (64, 64))
        output3 = self.probability_map_conv3(up_sampled4)
        interpolated_output3 = F.interpolate(output3, self.shape)

        up_sampled3 = self.up_sampler3(up_sampled4, down_outputs2)
        output2 = self.probability_map_conv2(up_sampled3)
        interpolated_output2 = F.interpolate(output2, self.shape)

        up_sampled2 = self.up_sampler2(up_sampled3, down_outputs1)
        output1 = self.probability_map_conv1(up_sampled2)
        interpolated_output1 = F.interpolate(output1, self.shape)

        up_sampled1 = self.up_sampler1(up_sampled2, down_outputs0)
        output0 = self.probability_map_conv0(up_sampled1)

        return torch.softmax(output0, dim=1), torch.softmax(interpolated_output1, dim=1), \
               torch.softmax(interpolated_output2, dim=1), torch.softmax(interpolated_output3, dim=1)
