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):
        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):
        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, low_res=32, high_res=64):
        super(UpscalerBasic, self).__init__()

        self.low_res = low_res
        self.high_res = high_res
        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(low_res * (low_res//2) * 3),
        ])

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

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

        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=(self.high_res, self.high_res//2), method=tf.image.ResizeMethod.LANCZOS5)

        return x




class UpscalerBasicNew(tf.keras.Model):
    
    def __init__(self, low_res=32):
        super(UpscalerBasicNew, self).__init__()

        self.low_res = low_res
        self.flatten = tf.keras.layers.Flatten(name='upscale_enhancer_flattent')

        self.seq = tf.keras.Sequential([
                    tf.keras.layers.Dense(4096, activation='relu', name='upscale_enhancer_dense_1'),
                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
                    tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_2'),
                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
                    tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_3'),
                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
                    tf.keras.layers.Dense(low_res * (low_res//2) * 3, activation='tanh', name='upscale_enhancer_dense_final'),
        ])

        self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3), name='upscale_enhancer_reshape')

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

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

        x = x * (1 - mask[:,:,:,tf.newaxis])
        x = x + img

        return x


class UpscalerBasicConv(tf.keras.Model):
    
    def __init__(self, low_res=32, kernel_size=3):
        super(UpscalerBasicConv, self).__init__()

        self.low_res = low_res

        self.seq = tf.keras.Sequential([
                    tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
                    tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
                    tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
                    tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'),
                    tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'),
        ])


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

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

        x = x * (1 - mask[:,:,:,tf.newaxis])
        x = x + img

        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):
#         x, mask = inputs

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

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

#         return x