import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from models.base import AEEncoder, Bottleneck, Decoder



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

        if size == 512:
            bottleneck_size = 64
            latent_dim = 128
            kernel_size = 7
            encoder_depths = [3, 3, 9, 3, 3]
            encoder_dims = [32, 64, 128, 256, bottleneck_size]
            decoder_input = (8, 4, bottleneck_size)
            decoder_depths = [3, 3, 9, 3, 3]
            decoder_dims = [bottleneck_size, 256, 128, 64, 16]
        elif size == 128:
            bottleneck_size = 512
            latent_dim = 128
            kernel_size = 5
            encoder_depths =[16, 16, 9, 9]
            encoder_dims = [64, 128, 256, bottleneck_size]
            decoder_input = (4, 2, bottleneck_size)
            decoder_depths = [16, 16, 9, 9]
            decoder_dims = [bottleneck_size, 256, 128, 64]
        elif size == 32:
            bottleneck_size = 512
            latent_dim = 128
            kernel_size = 3
            encoder_depths = [27, 27]
            encoder_dims = [256, bottleneck_size]
            decoder_input = (4, 2, bottleneck_size)
            decoder_depths = [27, 27]
            decoder_dims = [bottleneck_size, 256]
        else:
            raise NotImplementedError

        self.encoder : AEEncoder = AEEncoder(latent_dim, encoder_depths, encoder_dims, bottleneck_size, kernel_size)
        self.decoder : Decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size)


    def call(self, inputs, training=None, use_mask=True):
        if use_mask:
            x, mask = inputs
        else:
            x = inputs

        x = self.encoder(x, training=training)
        x = self.decoder(x, training=training)

        if use_mask:
             x = x * mask[:,:,:,tf.newaxis]

        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 = Bottleneck(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.transpose(x, [0,3,1,2])

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

#         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 = Bottleneck(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.transpose(x, [0,3,1,2])

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

#         return x