import torch
import numpy as np
import pyrender
import cv2
from pytorch3d.renderer import look_at_view_transform

def Rotate_y_180(X, pos='right'):
    R = torch.eye(3).to(X.device)
    R[0,0] = -1.0
    R[2,2] = -1.0
    if pos == 'right':
        X = torch.matmul(X, R)
    else:
        X = torch.matmul(R, X)
    return X

def Rotate_z_180(X, pos='right'):
    R = torch.eye(3).to(X.device)
    R[0,0] = -1.0
    R[1,1] = -1.0
    if pos == 'right':
        X = torch.matmul(X, R)
    else:
        X = torch.matmul(R, X)
    return X

class Renderer_pyr():
    def __init__(self, intrinsic, extrinsic, image_size=2048):
        self.name = 'pyrender'
        self.image_size = image_size

        self.R = extrinsic[0:3, 0:3]
        self.T = extrinsic[0:3, 3:4]
        self.K = intrinsic

        R_pyrender = Rotate_y_180(torch.inverse(self.R), 'right').float()
        T_pyrender = -torch.matmul(torch.inverse(self.R), self.T)[:,0].float()

        self.camera = pyrender.IntrinsicsCamera(self.K[0,0], self.K[1,1], self.image_size - self.K[0,2], self.image_size - self.K[1,2])
        self.camera_pose = np.eye(4)
        self.camera_pose[0:3,0:3] = R_pyrender.cpu().numpy()
        self.camera_pose[0:3,3] = T_pyrender.cpu().numpy()

        self.lights = []
        light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=100.0)
        light_pose1 = np.array([[1.0,  0.0,  0.0,   -3.0],
                                [0.0,  1.0,  0.0,   -3.0],
                                [0.0,  0.0,  1.0,   -1.0],
                                [0.0,  0.0,  0.0,   1.0],])
        light_pose2 = np.array([[1.0,  0.0,  0.0,   0.0],
                                [0.0,  1.0,  0.0,   -6.0],
                                [0.0,  0.0,  1.0,   -1.0],
                                [0.0,  0.0,  0.0,   1.0],])
        light_pose3 = np.array([[1.0,  0.0,  0.0,   3.0],
                                [0.0,  1.0,  0.0,   3.0],
                                [0.0,  0.0,  1.0,   -1.0],
                                [0.0,  0.0,  0.0,   1.0],])
        light_pose4 = np.array([[1.0,  0.0,  0.0,   0.0],
                                [0.0,  1.0,  0.0,   1.0],
                                [0.0,  0.0,  1.0,   -5.0],
                                [0.0,  0.0,  0.0,   1.0],])
        self.lights.append((light, light_pose1))
        self.lights.append((light, light_pose2))
        self.lights.append((light, light_pose3))
        self.lights.append((light, light_pose4))

    def render(self, mesh, return_mask=False):
        scene = pyrender.Scene(ambient_light=[0.1, 0.1, 0.1], bg_color=[0.0, 0.0, 0.0])
        scene.add(mesh)
        for light in self.lights:
            scene.add(light[0], pose=light[1])
        scene.add(self.camera, pose=self.camera_pose)
        osr = pyrender.OffscreenRenderer(self.image_size, self.image_size)
        color, depth = osr.render(scene)
        color = cv2.flip(color, -1)
        depth = cv2.flip(depth, -1)
        if return_mask:
            return color, (depth > 0).astype(np.uint8) * 255
        else:
            return color

    def save(self, path):
        intrinsic = self.K.numpy()
        extrinsic = np.hstack((self.R.numpy(), self.T.numpy()))
        np.savez(path, intrinsic=intrinsic, extrinsic=extrinsic)