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


class Dataset(Dataset):
    def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True):
        self.folder = folder
        self.device = device
        self.frames = read_frames_file(os.path.join(folder, frames_file))
        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.get_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()
        pose_out[1] = np.random.uniform(-np.pi, np.pi)
        # pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
        # generate random number between -np.pi, np.pi
        
        # start = time()

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

        render = self._render_body(body[0])[0,:,:,:3]
        render_out = self._render_body(body_out[0])[0,:,:,:3]

        render = render.permute(2,0,1)
        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_out = torch.min(render_out != 1, dim=0)[0][None,:,:].to(torch.float32)

        render = render * mask
        render_out = render_out * mask_out

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

        data = { 
                 'pose': render,
                 'shape': shape,
                 'pose_out': render_out,
                    # 'body': body,
                    # 'body_out': body_out,
                    # 'render': render,
                    # 'render_out': render_out,
                    'mask': mask,
                    '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 = Dataset('/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