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')
        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