import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')

import importlib
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())

if args['debug']:
    from configs.iccvw.pyr_64 import config
else:
    config_module = importlib.import_module(args['config'])
    config = config_module.config

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'

import time
import argparse
import wandb
import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image


import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from tqdm import tqdm
from data.dataset import Dataset
from utils.checkpoints import CheckpointSaver, load_model
from torch.utils.data import DataLoader
from utils.utils import text_to_html
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
from pytorch_msssim import MS_SSIM
from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
from smpl.smpl_torch import SMPL
from plotly.subplots import make_subplots
from models.pyramid_model import Pyramid
from models.cvae_model import CVAE


uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])


def step_pyramid(batch):
    model.eval()
    with torch.inference_mode():
        levels_input, conds = group_inputs_for_levels(batch, config)

        out_levels = model([levels_input, conds])
        pred, _, _ = out_levels[-1]

        mesh_pred = uv2mesh(pred)

        if config['use_skinning']:
            shape = batch['pose_shape'][:, 72:]
            gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64)
            shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])

            mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
            mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
        else:
            mesh_pred_posed = mesh_pred

    return mesh_pred_posed


def step(batch):
    return None


def generate():
    for batch in tqdm(test_dataloader):
        if is_pyramid:
            mesh_pred = step_pyramid(batch)
        else: 
            mesh_pred = step(batch)

        ids = batch['id']
        frame_nrs = batch['frame_nr']
        garments = batch['garment']

        for i in range(len(ids)):
            np.savez_compressed(f'{out_folder}/{ids[i]}_{frame_nrs[i]}_{garments[i]}', mesh_pred[i].cpu().numpy())




if __name__ == '__main__':

    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
    skirt_modifier = '_skirt' if config['is_skirt'] else ''
    dataset_resolutions = config['pyramid']['resolutions'].copy()
    if 512 not in dataset_resolutions:
        dataset_resolutions.append(512)

    dataset_test = Dataset(config['test_dir'], config['signature'], device=config['device'], frames_file=config['test_frames'], use_cache=config['dataset_in_memory'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
    test_dataloader = DataLoader(dataset_test, batch_size=4, drop_last=True, shuffle=True, )

    is_pyramid = config['name'].startswith('pyr')

    if is_pyramid:
        model: Pyramid = Pyramid(config)
    else:
        model: CVAE = CVAE(config)

    it = load_model(model, 'checkpoints/iccv/torch/' + config['name'])
    # it = load_model(model, 'checkpoints/iccvw/' + config['name'])

    name = config['name'] + '_' + str(it)

    folder = '/data-local/127_2/results'
    out_folder = f'{folder}/{name}'

    os.makedirs(out_folder, exist_ok=True)

    generate()