import os
import time
import random
import tensorflow as tf
import numpy as np
import scipy as sp
# import plotly.graph_objects as go
# import plotly.express as ex

from custom_utils.custom_utils import load_model
from models.base import AEEncoder, VariationalEncoder, Decoder


class CVAE(tf.keras.Model):
    
    def __init__(self, config):
        super(CVAE, self).__init__()
        self.latent_dim = config['cvae']['latent_dim']
        self.size = config['resolution']

        self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
        if config['static_model'] != '':
            load_model(self.static_enc, config['static_model'] + '_enc')

        self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
        if config['body_model'] != '':
            load_model(self.body_enc, config['body_model'] + '_enc')

        if config['use_curvature']:
            self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
            if config['curvature_model'] != '':
                load_model(self.curv_enc, config['curvature_model'] + '_enc')

        if config['freeze_encoders']:
            self.static_enc.trainable = False
            self.body_enc.trainable = False
            if config['use_curvature']:
                self.curv_enc.trainable = False

        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
                                                 tf.keras.layers.Dense(256, activation='relu'),
                                                 tf.keras.layers.Dense(128, activation='relu'),
                                                 tf.keras.layers.Dense(config['latent_dim'], activation='tanh')])

        self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
        self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])

        if config['avg_garment_path'] != '':
            self.avg_garment = np.load(config['avg_garment_path'])
        else:
            self.avg_garment = None

        self.add_noise = config['cvae']['add_noise']
        self.noise_amount = config['cvae']['noise_amount']
        self.use_curvature = config['use_curvature']
        self.use_offsets_over_body = not config['use_skinning']

        if self.use_offsets_over_body:
            assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
        if self.avg_garment is not None:
            assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'

    @tf.function
    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None):
        if noise is None:
            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
        else:
            eps = noise

        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
        z = tf.concat([eps, z_cond], axis=-1)

        x = self.decoder(z, training=False)
        
        if self.avg_garment is not None:
            x = x + self.avg_garment[tf.newaxis,:,:,:]
        
        x = x * uv_mask[:,:,:,tf.newaxis]

        return x


    def reparameterize(self, mean, logvar, fix_noise=False):
        if not fix_noise:
            eps = tf.random.normal(shape=mean.shape)
        else:
            eps = tf.random.normal(shape=(1, mean.shape[1]))
            eps = tf.tile(eps, [mean.shape[0],1])
        return eps * tf.exp(logvar * .5) + mean


    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training=None):
        z_body = self.body_enc(uv_body_posed, training=training)
        z_static = self.static_enc(uv_template, training=training)
        if self.use_curvature:
            z_curv = self.curv_enc(uv_curv, training=training)
            z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
        else:
            z_cond = tf.concat([z_static, z_body], axis=-1)

        z_cond = self.cond_encoder(z_cond, training=training)

        return z_cond


    def call(self, inputs, training=None):
        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs

        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
        mean, logvar = self.encoder(x, training=training)
        z = self.reparameterize(mean, logvar)
        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)

        if self.add_noise and training == True:
            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)


        z = tf.concat([z, z_cond], axis=-1)
        x = self.decoder(z, training=training)

        if self.avg_garment is not None:
            x = x + self.avg_garment[tf.newaxis,:,:,:]
    
        if self.use_offsets_over_body:
            x = x + uv_body_posed

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

        return x, mean, logvar