import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip


class Bottleneck(tf.keras.Model):
    
    def __init__(self, dim, kernel_size=7):
        super(Bottleneck, self).__init__()

        self.bottleneck = tf.keras.Sequential(
                [Block(dim=dim, drop_path=0, kernel_size=kernel_size,
                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
            )

    def call(self, inputs, training=None):
        x = inputs

        x = self.bottleneck(x, training=training)
        
        return x



class Upscaler(tf.keras.Model):
    
    def __init__(self, bottleneck_size):
        super(Upscaler, self).__init__()

        depths = [27, 27]
        dims = [512, 512]
        reverse_dims = [16] + dims[:-1]
        self.encoder = ConvNeXtWithSkip(depths=depths, dims=dims, kernel_size=3)

        self.bottleneck = Bottleneck(dims[-1], kernel_size=3)

        self.decoder = ConvNeXtReverseWithSkip(depths=depths, dims=reverse_dims, kernel_size=3)

    def call(self, inputs, training=None):
        img, mask = inputs

        img = tf.transpose(img, [0,2,3,1])

        x, skip_features = self.encoder(img, training=training)

        x = self.bottleneck(x, training=training)

        x = self.decoder([x, skip_features], training=training)

        x = x * (1 - mask[:,:,:,tf.newaxis])
        x = x + img
        x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)

        return x


class UpscalerBasic(tf.keras.Model):
    
    def __init__(self, bottleneck_size, size=32):
        super(UpscalerBasic, self).__init__()

        self.flatten = tf.keras.layers.Flatten()

        self.seq = tf.keras.Sequential([
                    tf.keras.layers.Dense(1024, activation='relu'),
                    tf.keras.layers.BatchNormalization(),
                    tf.keras.layers.Dense(1024, activation='relu'),
                    tf.keras.layers.BatchNormalization(),
                    tf.keras.layers.Dense(1024, activation='relu'),
                    tf.keras.layers.BatchNormalization(),
                    tf.keras.layers.Dense(size * (size//2) * 3),
        ])

        self.reshape = tf.keras.layers.Reshape((size, (size//2), 3))

    def call(self, inputs, training=None):
        img, mask = inputs

        if img.shape[1] == 3:
            img = tf.transpose(img, [0,2,3,1])

        x = self.flatten(img)
        x = self.seq(x, training=training)
        x = self.reshape(x)

        x = x * (1 - mask[:,:,:,tf.newaxis])
        x = x + img
        x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)

        return x

# class Upscaler(tf.keras.Model):
    
#     def __init__(self, bottleneck_size):
#         super(Upscaler, self).__init__()

#         self.decoder = ConvNeXtReverse(depths=list(reversed([9, 9, 9])), dims=list(reversed([bottleneck_size, 128, 32, 16])))

#     def call(self, inputs, training=None):
#         x, mask = inputs

#         x = tf.transpose(x, [0,2,3,1])
#         x = self.decoder(x)

#         x = x * mask[:,:,:,tf.newaxis]

#         return x