import torch
import torch.nn as nn


class DecoderLayer(nn.Module):
    def __init__(
        self, in_channels, out_channels, kernel_size=2, stride=2, padding=[0, 0]
    ):
        super(DecoderLayer, self).__init__()
        self.up_conv = nn.ConvTranspose2d(
            in_channels=in_channels,
            out_channels=in_channels // 2,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding[0],
        )

        self.bn1 = nn.BatchNorm2d(in_channels)

        self.conv = nn.Sequential(
            *[
                self._conv_relu_layer(
                    in_channels=in_channels if i == 0 else out_channels,
                    out_channels=out_channels,
                    padding=padding[1],
                )
                for i in range(2)
            ]
        )

    def _conv_relu_layer(self, in_channels, out_channels, padding=0):
        return nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=3,
                padding=padding,
            ),
            nn.ReLU(),
            nn.BatchNorm2d(out_channels),
        )

    @staticmethod
    def crop_cat(x, encoder_output):
        delta = (encoder_output.shape[-1] - x.shape[-1]) // 2
        encoder_output = encoder_output[
            :, :, delta : delta + x.shape[-1], delta : delta + x.shape[-1]
        ]
        return torch.cat((encoder_output, x), dim=1)

    def forward(self, x, encoder_output):
        x = self.crop_cat(self.up_conv(x), encoder_output)
        x = self.bn1(x)
        return self.conv(x)


class Decoder(nn.Module):
    def __init__(self, config):
        super(Decoder, self).__init__()
        self.decoder = nn.ModuleDict(
            {
                name: DecoderLayer(
                    in_channels=block["in_channels"],
                    out_channels=block["out_channels"],
                    kernel_size=block["kernel_size"],
                    stride=block["stride"],
                    padding=block["padding"],
                )
                for name, block in config.items()
            }
        )

    def forward(self, x, encoder_output):
        for name, block in self.decoder.items():
            x = block(x, encoder_output[name])
        return x