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


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, features):
        super(VariationalEncoderNew, self).__init__()

        self.encoder = resnet_18(features)

        self.mean_branch = tf.keras.Sequential([
            tf.keras.layers.Conv2D(filters=features,
                                            kernel_size=(3, 3),
                                            strides=1,
                                            padding="same"),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Conv2D(filters=features,
                                            kernel_size=(3, 3),
                                            strides=1,
                                            padding="same")
        ])

        self.var_branch = tf.keras.Sequential([
            tf.keras.layers.Conv2D(filters=features,
                                            kernel_size=(3, 3),
                                            strides=1,
                                            padding="same"),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Conv2D(filters=features,
                                            kernel_size=(3, 3),
                                            strides=1,
                                            padding="same")
        ])

        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, features, activation='tanh'):
        super(AEEncoderNew, self).__init__()

        self.encoder = resnet_18(features)
        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, features, is_label=False, is_cvae=False):
        super(DecoderNew, self).__init__()

        self.is_label = is_label

        if not is_label:
            self.decoder = resnet_18(features, reversed=True, is_cvae=is_cvae)
        else:
            self.decoder = resnet_18(features, reversed=True, labels=6, is_cvae=is_cvae)
            

    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:
            x = x * mask[:,:,:,tf.newaxis]

            background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis])
            x = tf.concat((background, x), axis=-1)

            x = tf.math.softmax(x)

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

        return x