import torch
import torch.nn as nn

from .blocks import *


class U2NetS(nn.Module):
    """
    U^2-Net Small varient
    https://arxiv.org/pdf/2005.09007.pdf
    """

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

        # Encoder
        self.e_block1 = ResUBlock(
            7, in_ch, 16, 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, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm),
        )
        self.e_block3 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlock(5, nf, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm),
        )
        self.e_block4 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlock(4, nf, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm),
        )
        self.e_block5 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlockFlat(nf, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm),
        )
        self.e_block6 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlockFlat(nf, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm),
        )

        # Decoder
        self.d_block5 = ResUBlockFlat(
            nf * 2, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm
        )
        self.d_block4 = ResUBlock(
            4, nf * 2, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm
        )
        self.d_block3 = ResUBlock(
            5, nf * 2, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm
        )
        self.d_block2 = ResUBlock(
            6, nf * 2, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm
        )
        self.d_block1 = ResUBlock(
            7, nf * 2, 16, nf, use_spectral_norm=use_spectral_norm, norm=norm
        )

        # Outputs
        self.side1 = nn.Conv2d(nf, out_ch, 3, padding=1)
        self.side2 = nn.Conv2d(nf, out_ch, 3, padding=1)
        self.side3 = nn.Conv2d(nf, out_ch, 3, padding=1)
        self.side4 = nn.Conv2d(nf, out_ch, 3, padding=1)
        self.side5 = nn.Conv2d(nf, out_ch, 3, padding=1)
        self.side6 = nn.Conv2d(nf, out_ch, 3, padding=1)
        self.out = nn.Conv2d(6 * out_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)
        e_4 = self.e_block4(e_3)
        e_5 = self.e_block5(e_4)
        e_6 = self.e_block6(e_5)
        e_6 = upsample_like(e_6, e_5)

        # Decoder
        d_5 = self.d_block5(torch.cat((e_6, e_5), 1))
        d_5 = upsample_like(d_5, e_4)
        d_4 = self.d_block4(torch.cat((d_5, e_4), 1))
        d_4 = upsample_like(d_4, e_3)
        d_3 = self.d_block3(torch.cat((d_4, 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))

        # Outputs at each scale
        out1 = self.side1(d_1)
        out2 = upsample_like(self.side2(d_2), d_1)
        out3 = upsample_like(self.side3(d_3), d_1)
        out4 = upsample_like(self.side4(d_4), d_1)
        out5 = upsample_like(self.side5(d_5), d_1)
        out6 = upsample_like(self.side6(e_6), d_1)

        # Combine all scales into final output
        out = self.out(torch.cat((out1, out2, out3, out4, out5, out6), 1))
        out = torch.sigmoid(out)

        return out


class U2NetM(nn.Module):
    """
    U^2-Net medium varient
    """

    def __init__(
        self, in_ch=2, out_ch=1, nf=64, use_spectral_norm=False, norm=nn.BatchNorm2d
    ):
        super(U2NetM, 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 // 2,
                nf * 2,
                use_spectral_norm=use_spectral_norm,
                norm=norm,
            ),
        )
        self.e_block4 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlock(
                4,
                nf * 2,
                nf // 2,
                nf * 2,
                use_spectral_norm=use_spectral_norm,
                norm=norm,
            ),
        )
        self.e_block5 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlockFlat(
                nf * 2, nf // 2, nf * 2, use_spectral_norm=use_spectral_norm, norm=norm
            ),
        )
        self.e_block6 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlockFlat(
                nf * 2, nf // 2, nf * 2, use_spectral_norm=use_spectral_norm, norm=norm
            ),
        )

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

        # Outputs
        self.side1 = nn.Conv2d(nf, out_ch, 3, padding=1)
        self.side2 = nn.Conv2d(nf * 2, out_ch, 3, padding=1)
        self.side3 = nn.Conv2d(nf * 2, out_ch, 3, padding=1)
        self.side4 = nn.Conv2d(nf * 2, out_ch, 3, padding=1)
        self.side5 = nn.Conv2d(nf * 2, out_ch, 3, padding=1)
        self.side6 = nn.Conv2d(nf * 2, out_ch, 3, padding=1)
        self.out = nn.Conv2d(6 * out_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)
        e_4 = self.e_block4(e_3)
        e_5 = self.e_block5(e_4)
        e_6 = self.e_block6(e_5)
        e_6 = upsample_like(e_6, e_5)

        # Decoder
        d_5 = self.d_block5(torch.cat((e_6, e_5), 1))
        d_5 = upsample_like(d_5, e_4)
        d_4 = self.d_block4(torch.cat((d_5, e_4), 1))
        d_4 = upsample_like(d_4, e_3)
        d_3 = self.d_block3(torch.cat((d_4, 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))

        # Outputs at each scale
        out1 = self.side1(d_1)
        out2 = upsample_like(self.side2(d_2), d_1)
        out3 = upsample_like(self.side3(d_3), d_1)
        out4 = upsample_like(self.side4(d_4), d_1)
        out5 = upsample_like(self.side5(d_5), d_1)
        out6 = upsample_like(self.side6(e_6), d_1)

        # Combine all scales into final output
        out = self.out(torch.cat((out1, out2, out3, out4, out5, out6), 1))
        out = torch.sigmoid(out)

        return out


class U2NetL(nn.Module):
    """
    U^2-Net large varient
    https://arxiv.org/pdf/2005.09007.pdf
    """

    def __init__(
        self, in_ch=2, out_ch=1, nf=64, use_spectral_norm=False, norm=nn.BatchNorm2d
    ):
        super(U2NetL, 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
            ),
        )
        self.e_block4 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlock(
                4,
                nf * 4,
                nf * 2,
                nf * 8,
                use_spectral_norm=use_spectral_norm,
                norm=norm,
            ),
        )
        self.e_block5 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlockFlat(
                nf * 8, nf * 4, nf * 8, use_spectral_norm=use_spectral_norm, norm=norm
            ),
        )
        self.e_block6 = nn.Sequential(
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
            ResUBlockFlat(
                nf * 8, nf * 4, nf * 8, use_spectral_norm=use_spectral_norm, norm=norm
            ),
        )

        # Decoder
        self.d_block5 = ResUBlockFlat(
            nf * 16, nf * 4, nf * 8, use_spectral_norm=use_spectral_norm, norm=norm
        )
        self.d_block4 = ResUBlock(
            4, nf * 16, nf * 2, nf * 4, use_spectral_norm=use_spectral_norm, norm=norm
        )
        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.side1 = nn.Conv2d(nf, out_ch, 3, padding=1)
        self.side2 = nn.Conv2d(nf, out_ch, 3, padding=1)
        self.side3 = nn.Conv2d(nf * 2, out_ch, 3, padding=1)
        self.side4 = nn.Conv2d(nf * 4, out_ch, 3, padding=1)
        self.side5 = nn.Conv2d(nf * 8, out_ch, 3, padding=1)
        self.side6 = nn.Conv2d(nf * 8, out_ch, 3, padding=1)
        self.out = nn.Conv2d(6 * out_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)
        e_4 = self.e_block4(e_3)
        e_5 = self.e_block5(e_4)
        e_6 = self.e_block6(e_5)
        e_6 = upsample_like(e_6, e_5)

        # Decoder
        d_5 = self.d_block5(torch.cat((e_6, e_5), 1))
        d_5 = upsample_like(d_5, e_4)
        d_4 = self.d_block4(torch.cat((d_5, e_4), 1))
        d_4 = upsample_like(d_4, e_3)
        d_3 = self.d_block3(torch.cat((d_4, 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))

        # Outputs at each scale
        out1 = self.side1(d_1)
        out2 = upsample_like(self.side2(d_2), d_1)
        out3 = upsample_like(self.side3(d_3), d_1)
        out4 = upsample_like(self.side4(d_4), d_1)
        out5 = upsample_like(self.side5(d_5), d_1)
        out6 = upsample_like(self.side6(e_6), d_1)

        # Combine all scales into final output
        out = self.out(torch.cat((out1, out2, out3, out4, out5, out6), 1))
        out = torch.sigmoid(out)

        return out
