import torch.nn as nn

from ..common import Conv2dReLU


class TransposeX2(nn.Module):

    def __init__(self, in_channels, out_channels, use_batchnorm=True, **batchnorm_params):
        super().__init__()
        layers = []
        layers.append(nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size=4, stride=2, padding=1))
        if use_batchnorm:
            layers.append(nn.BatchNorm2d(out_channels, **batchnorm_params))
        layers.append(nn.ReLU(inplace=True))

        self.block = nn.Sequential(*layers)

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


class DecoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels, use_batchnorm=True):
        super().__init__()

        self.block = nn.Sequential(
            Conv2dReLU(in_channels, in_channels // 4,
                       kernel_size=1, use_batchnorm=use_batchnorm),
            TransposeX2(in_channels // 4, in_channels // 4,
                        use_batchnorm=use_batchnorm),
            Conv2dReLU(in_channels // 4, out_channels,
                       kernel_size=1, use_batchnorm=use_batchnorm),
        )

    def forward(self, x):
        x, skip = x
        x = self.block(x)
        if skip is not None:
            x = x + skip
        return x


class LinkNetHead(nn.Module):
    __name__ = "LinkNetHead"

    def __init__(
            self,
            encoder_channels,
            prefinal_channels=32,
            num_class=1,
            use_batchnorm=True,
    ):
        super().__init__()

        in_channels = encoder_channels

        self.layer1 = DecoderBlock(
            in_channels[-1], in_channels[-2], use_batchnorm=use_batchnorm)
        self.layer2 = DecoderBlock(
            in_channels[-2], in_channels[-3], use_batchnorm=use_batchnorm)
        self.layer3 = DecoderBlock(
            in_channels[-3], in_channels[-4], use_batchnorm=use_batchnorm)
        self.layer4 = DecoderBlock(
            in_channels[-4], in_channels[-5], use_batchnorm=use_batchnorm)
        self.layer5 = DecoderBlock(
            in_channels[-5], prefinal_channels, use_batchnorm=use_batchnorm)
        self.final_conv = nn.Conv2d(
            prefinal_channels, num_class, kernel_size=(1, 1))

    def forward(self, x):
        encoder_head = x[-1]
        skips = x[:-1]

        x = self.layer1([encoder_head, skips[-1]])
        x = self.layer2([x, skips[-2]])
        x = self.layer3([x, skips[-3]])
        x = self.layer4([x, skips[-4]])
        x = self.layer5([x, None])
        x = self.final_conv(x)

        return x
