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"] = '3'

import torch
from torch.utils.data import Dataset
import numpy as np
from tqdm import tqdm
from time import time
from cloth3d.read import DataReader
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
from superSMPL.smpl_torch import SMPL
from data.pose_renderer import PoseRenderer
import trimesh
import random


class PoseDataset(Dataset):
    def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True, ae_mode=False):
        self.folder = folder
        self.device = device
        self.ae_mode = ae_mode
        self.frames = read_frames_file(os.path.join(folder, frames_file))
        random.shuffle(self.frames)
        self.reader = DataReader(folder, 'superSMPL')
        
        self.cache = dict() if use_cache else None

        self.poses = dict()
        self.shapes = dict()

        ids = set()
        for frame in self.frames:
            ids.add(frame.id)

        for id in tqdm(ids):
            info = self.reader.read_info(id)
            self.poses[id] = info['poses'].astype(np.float32)
            self.shapes[id] = np.concatenate((info['shape'], [info['gender']])).astype(np.float32)

        self.smpl = SMPL(path='superSMPL', device=device)
        self.renderer = PoseRenderer()
        # self.colors = np.load('custom_data/pose_colors.npy')


    def _render_body(self, body_to_render: torch.Tensor):
        # swap the y,z axes of the body_to_render 
        body_to_render = body_to_render[0,:,[0,2,1]]
        # and rotate it 90 degress along the z axis
        body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device=body_to_render.device), body_to_render.T).T

        out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.device))

        return out


    def __len__(self):
        return len(self.frames)
    

    def __getitem__(self, idx):
        frame : Frame = self.frames[idx]
        if self.cache is not None:
            if frame in self.cache.keys():
                return self.cache[frame]

        pose = torch.tensor(self.poses[frame.id][:,frame.frame_nr], device=self.device)
        shape = torch.tensor(self.shapes[frame.id], device=self.device)
        pose[:3] = 0
        pose_out = pose.clone()

        angle = torch.rand(1) * 2 * np.pi - np.pi
        matrix = self.smpl.rodrigues(torch.tensor([0.0, angle.numpy()[0], 0.0], device=self.device)[None, None, :]).to(self.device)[0]
        angle = angle.to(self.device)

        pose_out[1] = angle
        # pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
        # generate random number between -np.pi, np.pi
        
        # start = time()

        if self.ae_mode:
            body = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
        else:
            body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
            body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)

        # mesh = trimesh.Trimesh(vertices=body[0][0,:,:].cpu().numpy(), faces=self.smpl.f_all, process=False, validate=False)
        # mesh.apply_transform(trimesh.transformations.rotation_matrix(angle=3.14159, direction=[0, 1, 0], point=[0, 0, 0]))
        # body_back = torch.tensor(mesh.vertices, device=body[0].device).unsqueeze(0).to(torch.float32)
        # body_back[0,:,0] = -body_back[0,:,0] 

        # print(f'smpl took {time() - start} seconds')
        # start = time()

        render = self._render_body(body[0])[0,:,:,:3]
        # render_back = self._render_body(body_back)[0,:,:,:3]
        if not self.ae_mode:
            render_out = self._render_body(body_out[0])[0,:,:,:3]

        render = render.permute(2,0,1)
        # render_back = render_back.permute(2,0,1)
        if not self.ae_mode:
            render_out = render_out.permute(2,0,1)

        # import plotly.express as ex
        # ex.imshow(render[0,:,:,:3].cpu().numpy()).show()
        # ex.imshow(render_out[0,:,:,:3].cpu().numpy()).show()

        mask = torch.min(render != 1, dim=0)[0][None,:,:].to(torch.float32)
        # mask_back = torch.min(render_back != 1, dim=0)[0][None,:,:].to(torch.float32)
        if not self.ae_mode:
            mask_out = torch.min(render_out != 1, dim=0)[0][None,:,:].to(torch.float32)

        render = render * mask
        # render_back = render_back * mask_back
        if not self.ae_mode:
            render_out = render_out * mask_out

        # print(f'render took {time() - start} seconds')

        data = { 
                 'pose': render,
                 'shape': shape,
                #  'pose_back': render_back,
                    # 'body': body,
                    # 'body_out': body_out,
                    # 'render': render,
                    # 'render_out': render_out,
                 'mask': mask,
                 'angle': angle,
                 'matrix': matrix,
               }
        
        if not self.ae_mode:
            data['pose_out'] = render_out
            data['mask_out'] = mask_out

        if self.cache is not None:
            self.cache[frame] = data
        return data
    

if __name__ == '__main__':
    from torch.utils.data import DataLoader
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device)
    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)

    for batch in train_dataloader:
        batch
        pass

    pass