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




class Bottleneck(tf.keras.Model):
    
    def __init__(self, dim, latent_dim):
        super(Bottleneck, self).__init__()

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

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

        x = self.bottleneck(x)
        x = self.flatten(x)
        x = self.dense(x)
        
        return x


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

        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))

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

        x = self.dense(x)
        x = self.reshape(x)
        x = self.decoder(x)

        return x


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


        bottleneck_size = 512
        latent_dim = 128

        self.encoder = ConvNeXt(depths=[3, 3, 9, 9, 3, 3], dims=[32, 64, 128, 256, 256, bottleneck_size], include_top=False)
        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
        self.decoder = Decoder(bottleneck_size, latent_dim)


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

        x = tf.transpose(x, [0,2,3,1])
        x = self.encoder(x, training=training)
        # x = self.pooling(x)
        # x = self.reshape(x)
        # x = self.flatten(x)

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

        x = self.decoder(x, training=training)
        x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
        x = tf.transpose(x, [0,3,1,2])

        # x = x * uv_mask[:,:, tf.newaxis]
        # x = x * uv_mask_body[tf.newaxis,tf.newaxis,:,:]

        return x



class Bottleneck32(tf.keras.Model):
    
    def __init__(self, dim, latent_dim):
        super(Bottleneck32, self).__init__()

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

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

        x = self.bottleneck(x)
        x = self.flatten(x)
        x = self.dense(x)
        
        return x


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

        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
        self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))

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

        x = self.dense(x)
        x = self.reshape(x)
        x = self.decoder(x)

        return x


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


        bottleneck_size = 512
        latent_dim = 128

        self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
        self.bottleneck = Bottleneck32(bottleneck_size, latent_dim)
        self.decoder = Decoder32(bottleneck_size, latent_dim)


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

        x = tf.transpose(x, [0,2,3,1])
        x = self.encoder(x, training=training)

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

        x = self.decoder(x, training=training)
        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
        x = tf.transpose(x, [0,3,1,2])

        return x




class Bottleneck128(tf.keras.Model):
    
    def __init__(self, dim, latent_dim):
        super(Bottleneck128, self).__init__()

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

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

        x = self.bottleneck(x)
        x = self.flatten(x)
        x = self.dense(x)
        
        return x


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

        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
        self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))

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

        x = self.dense(x)
        x = self.reshape(x)
        x = self.decoder(x)

        return x


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


        bottleneck_size = 512
        latent_dim = 128

        self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
        self.bottleneck = Bottleneck128(bottleneck_size, latent_dim)
        self.decoder = Decoder128(bottleneck_size, latent_dim)


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

        x = tf.transpose(x, [0,2,3,1])
        x = self.encoder(x, training=training)

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

        x = self.decoder(x, training=training)
        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
        x = tf.transpose(x, [0,3,1,2])

        return x