import os
import sys
import numpy as np
import torch
from torch.nn import Module

from utils.data_utils import intrinsic, extrinsic, get_camera
from pytorch3d.renderer import (
    RasterizationSettings, 
    MeshRenderer, 
    MeshRasterizer,  
    SoftPhongShader,
    TexturesVertex,
    HardPhongShader,
)
from pytorch3d.structures import Meshes

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')
        colors = np.load('custom_data/body_colors_interpolated_unweighted.npy')
        
        self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
        self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy()


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

        raster_settings = RasterizationSettings(
                            image_size=(512, 512), 
                            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__':
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    from superSMPL.smpl_torch import SMPL
    import plotly.express as ex
    renderer = PoseRenderer()
    smpl = SMPL(path='superSMPL', device='cuda')
    pose = torch.rand(72, dtype=torch.float32, device='cuda')
    pose[:3] = 0
    shape = torch.rand(11, dtype=torch.float32, device='cuda')
    # pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
    # generate random number between -np.pi, np.pi
    
    # start = time()

    body = smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
    body_to_render = body[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 = renderer(body_to_render, torch.from_numpy(smpl.f_all).to(body_to_render.get_device()))

    ex.imshow(out[0, ..., :3].cpu().numpy()).show()

    print(out)