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"] = '0'
# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'

import torch
from torch.nn import Module
import numpy as np
from superSMPL.smpl_torch import SMPL
import trimesh

import plotly.express as ex
import plotly.graph_objects as go


from data_utils import intrinsic, extrinsic, get_camera

from pytorch3d.renderer import (
    RasterizationSettings, 
    MeshRenderer, 
    MeshRasterizer,  
    SoftPhongShader,
    TexturesVertex,
    HardPhongShader,
    AmbientLights
)
from pytorch3d.structures import Meshes


class PoseRenderer(Module):
    def __init__(self, device='cuda') -> None:
        super().__init__()
        colors = np.load('custom_data/pose_colors.npy')
        self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
        self.cam = torch.tensor([5.0, 0.0, 0.0], device=device).cpu().numpy()


    def forward(self, verts: torch.Tensor, faces: torch.Tensor):
        cameras, K = get_camera(self.cam, [480, 640], verts.get_device())

        raster_settings = RasterizationSettings(
                            image_size=(480, 640), 
                            blur_radius=0.0, 
                            faces_per_pixel=1, 
                            bin_size=0
                        )
        
        renderer = MeshRenderer(
            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
            shader=HardPhongShader(device=verts.device, cameras=cameras, lights=AmbientLights(device=verts.device))
        )

        output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))

        return output
    

if __name__ == '__main__':

    smpl = SMPL(path='superSMPL', device='cuda')
    pose = torch.zeros((1,72), device='cuda')
    # generate random pose
    pose[:,3:] = torch.randn((1,69), device='cuda') * 0.1
    body, _ = smpl(pose, torch.zeros((1,11), device='cuda'), remove=False)
    colors = np.zeros((body.shape[1], 3))


    renderer = PoseRenderer()
    body_to_render = body[0]
    # swap the y,z axes of the body_to_render 
    body_to_render = body_to_render[:,[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='cuda'), body_to_render.T).T

    out = renderer(body_to_render, torch.from_numpy(smpl.f_all).to(body.get_device()))
    ex.imshow(out[0,:,:,:3].cpu().numpy()).show()

    body = body.cpu().numpy()[0]
    faces = smpl.f_all

    # define a gradinet color cube for the body that goes from green to red based on the coordinates of the vertices and not their indices
    # colors = np.zeros((body.shape[0], 3))
    # colors[:,0] = (body[:,0] - np.min(body[:,0])) / (np.max(body[:,0]) - np.min(body[:,0]))
    # colors[:,1] = (body[:,1] - np.min(body[:,1])) / (np.max(body[:,1]) - np.min(body[:,1]))
    # colors[:,2] = (body[:,2] - np.min(body[:,2])) / (np.max(body[:,2]) - np.min(body[:,2]))
    # fig = go.Figure(data=[go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], vertexcolor=colors)])
    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
    # fig.show()  
    # np.save('custom_data/pose_colors.npy', colors.astype(np.float32))

    # same as a scatter plot
    # fig = go.Figure(data=[go.Scatter3d(x=body[:,0], y=body[:,2], z=body[:,1], mode='markers', marker=dict(size=2, color=colors, opacity=0.8))])
    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
    # fig.show()


    
    pass