import os
import time
import random
import torch
from torch import nn
import numpy as np
import scipy as sp
# import plotly.graph_objects as go
# import plotly.express as ex

from utils.checkpoints import load_model
from models.base import AEEncoder, VariationalEncoder, Decoder


class CVAE(nn.Module):
    
    def __init__(self, config):
        super(CVAE, self).__init__()
        self.size = config['resolution']

        self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
        if config['static_model'] != '':
            load_model(self.static_enc, config['static_model'] + '_enc')

        self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
        if config['body_model'] != '':
            load_model(self.body_enc, config['body_model'] + '_enc')

        if config['use_normals']:
            self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device'])
            if config['normals_model'] != '':
                load_model(self.norm_enc, config['normals_model'] + '_enc')

        if config['freeze_encoders']:
            for param in self.static_enc.parameters(): param.requires_grad = False
            for param in self.body_enc.parameters(): param.requires_grad = False
            if config['use_normals']:
                for param in self.norm_enc.parameters(): param.requires_grad = False

        self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
        self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])

        self.latent_shape = (config['cvae']['latent_dim'], 8, 4)

        if config['avg_garment_path'] != '':
            self.avg_garment = torch.permute(torch.tensor(np.load(config['avg_garment_path']), device=config['device']), [2,0,1])
        else:
            self.avg_garment = None

        self.use_normals = config['use_normals']
        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!'

    
    def sample(self, uv_template, uv_body_posed, uv_normals, uv_mask, noise=None):
        if noise is None:
            eps = torch.randn(uv_template.shape[0], self.latent_shape[0], self.latent_shape[1], self.latent_shape[2]).to(uv_template.device)
        else:
            eps = noise

        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
        z = torch.cat([eps, z_cond], dim=1)
        x = self.decoder(z)

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

        x = x * uv_mask[:,None,:,:]

        return x


    def reparameterize(self, mean: torch.Tensor, logvar: torch.Tensor, training: bool = True, fix_noise: bool = False):
        if training == True:
            if not fix_noise:
                eps = torch.randn(size=mean.shape).to(mean.device)
            else:
                eps = torch.randn(size=(1, mean.shape[1]))
                eps = torch.tile(eps, [mean.shape[0],1])
            out = eps * torch.exp(logvar * 0.5) + mean
        else:
            out = torch.exp(logvar * 0.5) + mean
        return out


    def encode_conds(self, uv_template, uv_body_posed, uv_normals):
        z_body = self.body_enc(uv_body_posed)
        z_static = self.static_enc(uv_template)
        if self.use_normals:
            z_norm = self.norm_enc(uv_normals)
            z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
        else:
            z_cond = torch.concat([z_static, z_body], dim=1)

        return z_cond


    def forward(self, inputs):
        uv_unposed, uv_template, uv_body_posed, uv_norm, uv_mask = inputs
        x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_norm], dim=1)
        
        mean, logvar = self.encoder(x)
        z: torch.Tensor = self.reparameterize(mean, logvar)
        z = z.view((mean.shape[0], *self.latent_shape))
        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_norm)

        z = torch.concat([z, z_cond], dim=1)
        x = self.decoder(z)

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

        x = x * uv_mask[:,None,:,:]

        return x, mean, logvar