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

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '1'

import time
import argparse
import wandb
from tqdm 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 torchvision.transforms import Resize
from torchvision.transforms.functional import InterpolationMode
from pathlib import Path

from data.utils import Frame
from data.dataset import read_frames_file
from utils.utils import save_stl, create_temp_frames
from utils.checkpoints import CheckpointSaver, load_model
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
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 data.dataset import Dataset
from utils.checkpoints import CheckpointSaver, load_model
from torch.utils.data import DataLoader
import trimesh as tm
from zipfile import ZipFile

from configs.iccvw.pyr_64_128_256_512_e2e import config as config_1
from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_2

uv2mesh = dict()
uv2mesh['nonskirt'] = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config_1['is_cape'])
uv2mesh['skirt'] = UV2Mesh(res=256, is_skirt=True, is_smpl_orig=config_1['is_cape'])


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

        out = []
        out_levels = model([levels_input, conds])
        for i in range(len(out_levels)):
            pred, _, _ = out_levels[i]

            mesh_pred = uv2mesh[picker](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]
                if not config['is_cape']:
                    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
                else:
                    mesh_pred_posed = mesh_pred_shaped
            else:
                mesh_pred_posed = mesh_pred

            out.append(mesh_pred_posed)

    return out





if __name__ == '__main__':
    results_folder = Path('/data-local/127_2/results')
    out_folder = Path('temp/pyramid_ablation')
    iteration = {'nonskirt': 250000, 'skirt': 200000 }

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

    model = { 'nonskirt': Pyramid(config_1), 'skirt': Pyramid(config_2) }
    load_model(model['nonskirt'], 'checkpoints/iccvw/' + config_1['name'], iteration=iteration['nonskirt'])
    load_model(model['skirt'], 'checkpoints/iccvw/' + config_2['name'], iteration=iteration['skirt'])
    model['nonskirt'].eval()
    model['skirt'].eval()

    frames = ['01070_13_Jumpsuit', '00054_1_Dress', '03849_27_Dress', '06057_152_Dress',
              '03179_274_Dress', '04247_102_Tshirt', '01852_70_Tshirt']
    frames = ['00027_1_Dress']

    frames_corrupt = ['00300_69_Jumpsuit', '07608_77_Jumpsuit', '06955_117_Dress']
    frames = [Frame(frame_str=frame) for frame in frames]
    body_faces = smpl.f

    config = { 'nonskirt': config_1, 'skirt': config_2 }

    for frame in tqdm(frames):
        create_temp_frames(config_1['test_dir'], frame)
        dataset = Dataset(config_1['test_dir'], config_1['signature'], device=config_1['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=dataset_resolutions, is_cape=config_1['is_cape'])
        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )

        picker = 'skirt' if frame.garment in ['Skirt', 'Dress'] else 'nonskirt'

        batch = next(iter(dataloader))
        pred = step_pyramid(model[picker], batch, config[picker])
        
        fix_data = np.load(Path(dataset.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
        mask = fix_data['mask']
        faces = fix_data['faces']

        body = batch['mesh_body_posed'][0].cpu().numpy()
        gt = batch['mesh_posed'][0].cpu().numpy()

        save_stl(out_folder / f'{frame}_gt.stl', gt[mask, :], faces)
        save_stl(out_folder / f'{frame}_gt_body.stl', body, body_faces)

        for i, _pred in enumerate(pred):
            mesh_posed_fixed = fix_collision(_pred, batch['mesh_body_posed'], batch['mesh_mask']).cpu().numpy()

            mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
            m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
            tm.smoothing.filter_humphrey(m, iterations=2)
            verts = m.vertices

            # save_stl(out_folder / f'{frame}_lvl{i}_body.stl', body, body_faces)
            # save_stl(out_folder / f'{frame}_lvl{i}_mesh.stl', verts, faces)

            fig = go.Figure(data=[go.Mesh3d(x=verts[:,0], y=verts[:,2], z=verts[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
                                  go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray', opacity=1)
                                  ])
            fig.update_layout(scene=dict(aspectmode='cube'))
            fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
            fig.show()
            pass

        pass