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

from custom_utils.custom_utils import load_model
from models.cvae_model import CVAE
from models.upscaler import UpscalerBasic as Upscaler


class Pyramid(tf.keras.Model):
    
    def __init__(self, low_res=32,
                 latent_dim=128,
                 avg_garment_path='',
                 static_model='',
                 body_model='',
                 curvature_model='',
                 upscaler_model='',
                 freeze_encoders=False,
                 use_full_res_encoders=True):
        super(Pyramid, self).__init__()


        self.model_low = CVAE(latent_dim=latent_dim,
                              size=low_res,
                              static_model=static_model, 
                              body_model=body_model, 
                              curvature_model=curvature_model, 
                              freeze_encoders=freeze_encoders, 
                              encoder_res=512 if use_full_res_encoders else 32)

        self.model_full = CVAE(latent_dim=latent_dim,
                               size=512,
                               static_model=static_model, 
                               body_model=body_model, 
                               curvature_model=curvature_model, 
                               freeze_encoders=freeze_encoders,
                               encoder_res=512)

        self.upscale_enhancer = None
        if upscaler_model != '':
            self.upscale_enhancer = Upscaler(512, low_res)
            load_model(self.upscale_enhancer, upscaler_model)
            self.upscale_enhancer = tf.keras.Sequential([
                self.upscale_enhancer.flatten,
                self.upscale_enhancer.seq,
                self.upscale_enhancer.reshape
            ])

            self.upscale_enhancer.trainable = False
        
        self.use_full_res_encoders = use_full_res_encoders

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

    def enhance(self, low, mask, training=None):
        x = self.upscale_enhancer(low, training=training)
        x = x * (1 - mask[:,:,:,tf.newaxis])
        low = x + low
        return low

    def call(self, inputs, training=None):
        uv_unposed, uv_static, uv_body_posed, uv_curv, uv_mask, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_mask_low = inputs

        if self.use_full_res_encoders:
            pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_static, uv_body_posed, uv_curv], training=training)
        else:
            pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training)

        pred_low = pred_low * uv_mask_low[:,:,:,tf.newaxis]
        
        if self.upscale_enhancer is not None:
            pred_low = self.enhance(pred_low, uv_mask_low, training=training)
        
        x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
        x = x * uv_mask[:,:,:,tf.newaxis]

        pred_full, mean_full, logvar_full, _ = self.model_full([uv_unposed, uv_static, uv_body_posed, uv_curv], training=training)

        x = x + pred_full
        x = x * uv_mask[:,:,:,tf.newaxis]
        
        return x, mean_full, logvar_full, pred_low, mean_low, logvar_low,