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

class Bottleneck(tf.keras.Model):
    
    def __init__(self, dim, latent_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)]
            )
            
        self.flatten = tf.keras.layers.Flatten()
        self.dense = tf.keras.layers.Dense(latent_dim)

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

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

        return x


class VariationalEncoder(tf.keras.Model):
    
    def __init__(self, latent_dim, depths, dims, kernel_size=7):
        super(VariationalEncoder, self).__init__()

        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)

        self.flatten = tf.keras.layers.Flatten()
        self.mean = tf.keras.layers.Dense(latent_dim)
        self.log_var = tf.keras.layers.Dense(latent_dim)

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

        x = self.encoder(x, training=training)
        x = self.flatten(x)
        mean = self.mean(x)
        log_var = self.log_var(x)

        return mean, log_var


class AEEncoder(tf.keras.Model):
    
    def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7):
        super(AEEncoder, self).__init__()

        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
        self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)


    def call(self, inputs, training):
        if len(inputs) == 2:
            x, _ = inputs
        else:
            x = inputs

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

        if len(inputs) == 2:
            return x, inputs[1]
        else:
            return x


class Decoder(tf.keras.Model):
    
    def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
        super(Decoder, self).__init__()

        self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
        self.reshape = tf.keras.layers.Reshape(input_size)
        self.upscale = upscale
        
        if not is_label:
            self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
        else:
            self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
            

    def call(self, inputs, training):
        if len(inputs) == 2:
            x, mask = inputs
        else:
            x = inputs

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

        if self.upscale:
            x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)


        if len(inputs) == 2:
            if len(mask.shape) == 2:
                x = x * mask[tf.newaxis,:,:,tf.newaxis]
            else:
                x = x * mask[:,:,:,tf.newaxis]

        return x



class VariationalEncoderNew(tf.keras.Model):
    
    def __init__(self, depths, dims, kernel_size=7):
        super(VariationalEncoderNew, self).__init__()

        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)

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

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

        self.flatten_mean = tf.keras.layers.Flatten()
        self.flatten_var = tf.keras.layers.Flatten()



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

        x = self.encoder(x, training=training)
        mean = self.mean_branch(x)
        log_var = self.var_branch(x)

        mean = self.flatten_mean(mean)
        log_var = self.flatten_var(log_var)

        return mean, log_var
        


class AEEncoderNew(tf.keras.Model):
    
    def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
        super(AEEncoderNew, self).__init__()

        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
        self.act = tf.keras.layers.Activation(activation)


    def call(self, inputs, training):
        if len(inputs) == 2:
            x, _ = inputs
        else:
            x = inputs

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

        if len(inputs) == 2:
            return x, inputs[1]
        else:
            return x


class DecoderNew(tf.keras.Model):
    
    def __init__(self, depths, dims, kernel_size=7, is_label=False):
        super(DecoderNew, self).__init__()

        self.is_label = is_label

        if not is_label:
            self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
        else:
            self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
            

    def call(self, inputs, training):
        if len(inputs) == 2:
            x, mask = inputs
        else:
            x = inputs

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

        # if self.is_label:

        if len(inputs) == 2:
            if len(mask.shape) == 2:
                x = x * mask[tf.newaxis,:,:,tf.newaxis]
            else:
                x = x * mask[:,:,:,tf.newaxis]

        return x