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, CVAE128, CVAE32


class Pyramid(tf.keras.Model):
    
    def __init__(self, low_res=32, latent_dim=128, avg_garment_path=''):
        super(Pyramid, self).__init__()

        if low_res == 32:
            self.model_low = CVAE32(latent_dim=latent_dim)
        if low_res == 128:
            self.model_low = CVAE128(latent_dim=latent_dim)
        self.model_full = CVAE(latent_dim=latent_dim)
        
        if avg_garment_path != '':
            self.avg_garment = np.load(avg_garment_path)
        else:
            self.avg_garment = None

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

        pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training)

        x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)

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

        x = x + pred_full
        
        return x, mean_full, logvar_full, pred_low, mean_low, logvar_low,