import torch
import torch.nn as nn

from .blocks import *


class MTEPainter2(nn.Module):
    def __init__(
        self, in_ch=2, out_ch=1, nf=64, use_spectral_norm=False, norm=nn.BatchNorm2d
    ):
        super(MTEPainter2, self).__init__()

        # Encoder
        self.e_block1 = ResUBlock(
            7, in_ch, nf // 2, nf, use_spectral_norm=use_spectral_norm, norm=norm
        )
        self.e_block2 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlock(
                6, nf, nf // 2, nf * 2, use_spectral_norm=use_spectral_norm, norm=norm
            ),
        )
        self.e_block3 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlock(
                5, nf * 2, nf, nf * 4, use_spectral_norm=use_spectral_norm, norm=norm
            ),
        )

        # Middle
        self.middle = nn.Sequential(
            ResUBlockFlat(
                nf * 4, nf * 2, nf * 4, use_spectral_norm=use_spectral_norm, norm=norm
            ),
            ResUBlockFlat(
                nf * 4, nf * 2, nf * 4, use_spectral_norm=use_spectral_norm, norm=norm
            ),
            ResUBlockFlat(
                nf * 4, nf * 2, nf * 4, use_spectral_norm=use_spectral_norm, norm=norm
            ),
            ResUBlockFlat(
                nf * 4, nf * 2, nf * 4, use_spectral_norm=use_spectral_norm, norm=norm
            ),
        )

        # Decoder
        self.d_block3 = ResUBlock(
            5, nf * 8, nf, nf * 2, use_spectral_norm=use_spectral_norm, norm=norm
        )
        self.d_block2 = ResUBlock(
            6, nf * 4, nf // 2, nf, use_spectral_norm=use_spectral_norm, norm=norm
        )
        self.d_block1 = ResUBlock(
            7, nf * 2, nf // 4, nf, use_spectral_norm=use_spectral_norm, norm=norm
        )

        # Outputs
        self.out = nn.Conv2d(nf + in_ch, out_ch, 1)

    def forward(self, x):
        # Encoder
        e_1 = self.e_block1(x)
        e_2 = self.e_block2(e_1)
        e_3 = self.e_block3(e_2)

        # Middle
        m = self.middle(e_3)

        # Decoder
        d_3 = self.d_block3(torch.cat((m, e_3), 1))
        d_3 = upsample_like(d_3, e_2)
        d_2 = self.d_block2(torch.cat((d_3, e_2), 1))
        d_2 = upsample_like(d_2, e_1)
        d_1 = self.d_block1(torch.cat((d_2, e_1), 1))

        # Combine all scales into final output
        out = self.out(torch.cat((d_1, x), 1))
        out = (torch.tanh(out) + 1) / 2

        return out


class MTEPainter(nn.Module):
    def __init__(
        self,
        in_ch,
        out_ch=1,
        nf=64,
        n_blocks=8,
        norm=nn.InstanceNorm2d,
        use_spectral_norm=False,
        dilation=2,
    ):
        super(MTEPainter, self).__init__()

        # Encoder
        self.encoder_pad = nn.ReflectionPad2d(3)
        self.e_block1 = ConvBlock(
            in_ch,
            nf,
            kernel_size=7,
            padding=0,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )

        self.e_block2 = ConvBlock(
            nf,
            nf * 2,
            kernel_size=4,
            stride=2,
            padding=1,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )

        self.e_block3 = ConvBlock(
            nf * 2,
            nf * 4,
            kernel_size=4,
            stride=2,
            padding=1,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )

        self.e_block4 = ConvBlock(
            nf * 4,
            nf * 8,
            kernel_size=4,
            stride=2,
            padding=1,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )

        # Resblocks
        resnet_blocks = []
        for i in range(n_blocks):
            resnet_blocks += [ResNetBlock(nf * 8, norm=norm, dilation=dilation)]
        self.resnet_blocks = nn.Sequential(*resnet_blocks)

        # Decoder
        self.d_block3 = ConvBlockUpSample(
            nf * 8,
            nf * 4,
            kernel_size=3,
            stride=1,
            padding=1,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )

        self.d_block2 = ConvBlockUpSample(
            nf * 4,
            nf * 2,
            kernel_size=3,
            stride=1,
            padding=1,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )

        self.d_block1 = ConvBlockUpSample(
            nf * 2,
            nf * 1,
            kernel_size=3,
            stride=1,
            padding=1,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )

        # Skip connection fusion layers
        self.fusion4 = ConvBlock(
            nf * 16,
            nf * 8,
            kernel_size=1,
            stride=1,
            padding=0,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )
        self.fusion3 = ConvBlock(
            nf * 8,
            nf * 4,
            kernel_size=1,
            stride=1,
            padding=0,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )
        self.fusion2 = ConvBlock(
            nf * 4,
            nf * 2,
            kernel_size=1,
            stride=1,
            padding=0,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )
        self.fusion1 = ConvBlock(
            nf * 2,
            nf,
            kernel_size=1,
            stride=1,
            padding=0,
            use_spectral_norm=use_spectral_norm,
            norm=norm,
        )

        self.out = nn.Sequential(
            nn.ReflectionPad2d(3), nn.Conv2d(nf, 1, kernel_size=7, padding=0)
        )

    def forward(self, x):
        # Encoder
        x = self.encoder_pad(x)
        e_1 = self.e_block1(x)  # [512x512]->[512x512]
        e_2 = self.e_block2(e_1)  # [512x512]->[256x256]
        e_3 = self.e_block3(e_2)  # [256x256]->[128x128]
        e_4 = self.e_block4(e_3)  # [128x128]->[64x64]

        # Middle
        m = self.resnet_blocks(e_4)  # [64x64]->[64x64]

        # Decoder
        d_3 = self.d_block3(
            self.fusion4(torch.cat((m, e_4), dim=1))
        )  # [64x64]->[128x128]
        d_2 = self.d_block2(
            self.fusion3(torch.cat((d_3, e_3), dim=1))
        )  # [128x128]->[256x256]
        d_1 = self.d_block1(
            self.fusion2(torch.cat((d_2, e_2), dim=1))
        )  # [256x256]->[512x512]
        d_0 = self.fusion1(torch.cat((d_1, e_1), dim=1))

        out = self.out(d_0)
        out = (torch.tanh(out) + 1) / 2

        return out


class EncoderDecoder(nn.Module):
    def __init__(
        self,
        in_ch=3,
        out_ch=1,
        nf=64,
        residual_blocks=8,
        dilation=2,
        use_spectral_norm=False,
        norm=nn.BatchNorm2d,
    ):
        super(EncoderDecoder, self).__init__()

        self.encoder = nn.Sequential(
            nn.ReflectionPad2d(3),
            spectral_norm(
                nn.Conv2d(in_channels=in_ch, out_channels=nf, kernel_size=7, padding=0),
                use_spectral_norm,
            ),
            norm(nf),
            nn.ReLU(True),
            spectral_norm(
                nn.Conv2d(
                    in_channels=nf,
                    out_channels=nf * 2,
                    kernel_size=4,
                    stride=2,
                    padding=1,
                ),
                use_spectral_norm,
            ),
            norm(nf * 2),
            nn.ReLU(True),
            spectral_norm(
                nn.Conv2d(
                    in_channels=nf * 2,
                    out_channels=nf * 4,
                    kernel_size=4,
                    stride=2,
                    padding=1,
                ),
                use_spectral_norm,
            ),
            norm(nf * 4),
            nn.ReLU(True),
        )

        blocks = []
        for _ in range(residual_blocks):
            block = ResNetBlock(
                nf * 4, dilation=dilation, use_spectral_norm=use_spectral_norm
            )
            blocks.append(block)

        self.middle = nn.Sequential(*blocks)

        self.decoder = nn.Sequential(
            spectral_norm(
                nn.ConvTranspose2d(
                    in_channels=nf * 4,
                    out_channels=nf * 2,
                    kernel_size=4,
                    stride=2,
                    padding=1,
                ),
                use_spectral_norm,
            ),
            norm(nf * 2),
            nn.ReLU(True),
            spectral_norm(
                nn.ConvTranspose2d(
                    in_channels=nf * 2,
                    out_channels=nf,
                    kernel_size=4,
                    stride=2,
                    padding=1,
                ),
                use_spectral_norm,
            ),
            norm(nf),
            nn.ReLU(True),
            nn.ReflectionPad2d(3),
            nn.Conv2d(in_channels=nf, out_channels=out_ch, kernel_size=7, padding=0),
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.middle(x)
        x = self.decoder(x)
        x = torch.sigmoid(x)
        return x
