import paddle
import paddle.nn as nn
import paddle.nn.functional as F

class KDUNet(nn.Layer):
    def __init__(self, num_classes=50):
        super(KDUNet, self).__init__()
        self.downsample = Downsample()
        self.upsample = Upsample(num_classes)

    def forward(self, inputs):
        x = paddle.to_tensor(inputs[0])
        split_dims_v = inputs[1]
        x, shortcut = self.downsample(x, split_dims_v)
        x = self.upsample(x, shortcut)
        x = x.transpose([0, 2, 1])

        return x

class ConvBNReLU(nn.Layer):
    def __init__(self,
                in_channels,
                out_channels,
                kernel_size,
                stride,
                padding='same',
                **kwargs):
        super().__init__()

        self._conv = nn.Conv1D(
            in_channels, out_channels, kernel_size, stride, padding=padding, **kwargs)

        self._batch_norm = nn.BatchNorm(out_channels)

    def forward(self, x):
        x = self._conv(x)
        x = self._batch_norm(x)
        x = F.relu(x)
        return x

class Downsample(nn.Layer):
    def __init__(self):
        super(Downsample, self).__init__()
        self.convbnrelu1 = ConvBNReLU(3, 32 * 3, 1, 1)
        self.convbnrelu2 = ConvBNReLU(32, 64 * 3, 1, 1)
        self.convbnrelu3 = ConvBNReLU(64, 256 * 3, 1, 1)
        self.convbnrelu4 = ConvBNReLU(256, 512 * 3, 1, 1)
        self.convbnrelu5 = ConvBNReLU(512, 1024 * 3, 1, 1)

    def forward(self, x, split_dims_v):
        def kdconv(x, shortcut, dim, featdim, select, convbnrelu):
            shortcut.append(x)
            x = convbnrelu(x)
            x = paddle.reshape(x, (-1, featdim, 3, dim))
            x = paddle.reshape(x, (-1, featdim, 3 * dim))
            select = paddle.to_tensor(select) + (paddle.arange(0, dim) * 3)
            x = paddle.index_select(x, axis=2, index=select)
            x = paddle.reshape(x, (-1, featdim, int(dim / 2), 2))
            x = paddle.max(x, axis=-1)

            return x, shortcut

        shortcut = []

        x, shortcut = kdconv(x, shortcut, 1024, 32, split_dims_v[0], self.convbnrelu1)
        x, shortcut = kdconv(x, shortcut, 512, 64, split_dims_v[1], self.convbnrelu2)
        x, shortcut = kdconv(x, shortcut, 256, 256, split_dims_v[2], self.convbnrelu3)
        x, shortcut = kdconv(x, shortcut, 128, 512, split_dims_v[3], self.convbnrelu4)
        x, shortcut = kdconv(x, shortcut, 64, 1024, split_dims_v[4], self.convbnrelu5)

        return x, shortcut


class Upsample(nn.Layer):
    def __init__(self, num_classes=50):
        super(Upsample, self).__init__()
        self.deconv1 = nn.Conv1DTranspose(1024, 512, 2, 2)
        self.doubleconv1 = nn.Sequential(
            ConvBNReLU(1024, 512, 1, 1),
            ConvBNReLU(512, 512, 1, 1))
        self.deconv2 = nn.Conv1DTranspose(512, 512, 2, 2)
        self.doubleconv2 = nn.Sequential(
            ConvBNReLU(768, 512, 1, 1),
            ConvBNReLU(512, 512, 1, 1))
        self.deconv3 = nn.Conv1DTranspose(512, 256, 2, 2)
        self.doubleconv3 = nn.Sequential(
            ConvBNReLU(320, 256, 1, 1),
            ConvBNReLU(256, 256, 1, 1))
        self.deconv4 = nn.Conv1DTranspose(256, 256, 2, 2)
        self.doubleconv4 = nn.Sequential(
            ConvBNReLU(288, 128, 1, 1),
            ConvBNReLU(128, 128, 1, 1))
        self.deconv5 = nn.Conv1DTranspose(128, 128, 2, 2)
        self.doubleconv5 = nn.Sequential(
            ConvBNReLU(131, 128, 1, 1),
            nn.Conv1D(128, num_classes, 1, 1))

    def forward(self, x, shortcut):
        x = self.deconv1(x)
        x = paddle.concat([x, shortcut[-1]], axis=1)
        x = self.doubleconv1(x)
        x = self.deconv2(x)
        x = paddle.concat([x, shortcut[-2]], axis=1)
        x = self.doubleconv2(x)
        x = self.deconv3(x)
        x = paddle.concat([x, shortcut[-3]], axis=1)
        x = self.doubleconv3(x)
        x = self.deconv4(x)
        x = paddle.concat([x, shortcut[-4]], axis=1)
        x = self.doubleconv4(x)
        x = self.deconv5(x)
        x = paddle.concat([x, shortcut[-5]], axis=1)
        x = self.doubleconv5(x)

        return x