import os

if "PYOPENGL_PLATFORM" not in os.environ:
    os.environ["PYOPENGL_PLATFORM"] = "egl"
import cv2
import numpy as np
import pyrender
import torch
import trimesh
from torchvision.utils import make_grid

from .render_openpose import render_openpose


def create_raymond_lights():
    import pyrender

    thetas = np.pi * np.array([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0])
    phis = np.pi * np.array([0.0, 2.0 / 3.0, 4.0 / 3.0])

    nodes = []

    for phi, theta in zip(phis, thetas):
        xp = np.sin(theta) * np.cos(phi)
        yp = np.sin(theta) * np.sin(phi)
        zp = np.cos(theta)

        z = np.array([xp, yp, zp])
        z = z / np.linalg.norm(z)
        x = np.array([-z[1], z[0], 0.0])
        if np.linalg.norm(x) == 0:
            x = np.array([1.0, 0.0, 0.0])
        x = x / np.linalg.norm(x)
        y = np.cross(z, x)

        matrix = np.eye(4)
        matrix[:3, :3] = np.c_[x, y, z]
        nodes.append(
            pyrender.Node(
                light=pyrender.DirectionalLight(color=np.ones(3), intensity=1.0), matrix=matrix
            )
        )

    return nodes


class MeshRenderer:

    def __init__(self, cfg, faces=None):
        self.cfg = cfg
        self.focal_length = cfg.EXTRA.FOCAL_LENGTH
        self.img_res = cfg.MODEL.IMAGE_SIZE
        self.renderer = pyrender.OffscreenRenderer(
            viewport_width=self.img_res, viewport_height=self.img_res, point_size=1.0
        )

        self.camera_center = [self.img_res // 2, self.img_res // 2]
        self.faces = faces

    def visualize(self, vertices, camera_translation, images, focal_length=None, nrow=3, padding=2):
        images_np = np.transpose(images, (0, 2, 3, 1))
        rend_imgs = []
        for i in range(vertices.shape[0]):
            fl = self.focal_length
            rend_img = torch.from_numpy(
                np.transpose(
                    self.__call__(
                        vertices[i],
                        camera_translation[i],
                        images_np[i],
                        focal_length=fl,
                        side_view=False,
                    ),
                    (2, 0, 1),
                )
            ).float()
            rend_img_side = torch.from_numpy(
                np.transpose(
                    self.__call__(
                        vertices[i],
                        camera_translation[i],
                        images_np[i],
                        focal_length=fl,
                        side_view=True,
                    ),
                    (2, 0, 1),
                )
            ).float()
            rend_imgs.append(torch.from_numpy(images[i]))
            rend_imgs.append(rend_img)
            rend_imgs.append(rend_img_side)
        rend_imgs = make_grid(rend_imgs, nrow=nrow, padding=padding)
        return rend_imgs

    def visualize_tensorboard(
        self,
        vertices,
        camera_translation,
        images,
        pred_keypoints,
        gt_keypoints,
        focal_length=None,
        nrow=5,
        padding=2,
    ):
        images_np = np.transpose(images, (0, 2, 3, 1))
        rend_imgs = []
        pred_keypoints = np.concatenate(
            (pred_keypoints, np.ones_like(pred_keypoints)[:, :, [0]]), axis=-1
        )
        pred_keypoints = self.img_res * (pred_keypoints + 0.5)
        gt_keypoints[:, :, :-1] = self.img_res * (gt_keypoints[:, :, :-1] + 0.5)
        # keypoint_matches = [(1, 12), (2, 8), (3, 7), (4, 6), (5, 9), (6, 10), (7, 11), (8, 14), (9, 2), (10, 1), (11, 0), (12, 3), (13, 4), (14, 5)]
        for i in range(vertices.shape[0]):
            fl = self.focal_length
            rend_img = torch.from_numpy(
                np.transpose(
                    self.__call__(
                        vertices[i],
                        camera_translation[i],
                        images_np[i],
                        focal_length=fl,
                        side_view=False,
                    ),
                    (2, 0, 1),
                )
            ).float()
            rend_img_side = torch.from_numpy(
                np.transpose(
                    self.__call__(
                        vertices[i],
                        camera_translation[i],
                        images_np[i],
                        focal_length=fl,
                        side_view=True,
                    ),
                    (2, 0, 1),
                )
            ).float()
            hand_keypoints = pred_keypoints[i, :21]
            # extra_keypoints = pred_keypoints[i, -19:]
            # for pair in keypoint_matches:
            #    hand_keypoints[pair[0], :] = extra_keypoints[pair[1], :]
            pred_keypoints_img = render_openpose(255 * images_np[i].copy(), hand_keypoints) / 255
            hand_keypoints = gt_keypoints[i, :21]
            # extra_keypoints = gt_keypoints[i, -19:]
            # for pair in keypoint_matches:
            #    if extra_keypoints[pair[1], -1] > 0 and hand_keypoints[pair[0], -1] == 0:
            #        hand_keypoints[pair[0], :] = extra_keypoints[pair[1], :]
            gt_keypoints_img = render_openpose(255 * images_np[i].copy(), hand_keypoints) / 255
            rend_imgs.append(torch.from_numpy(images[i]))
            rend_imgs.append(rend_img)
            rend_imgs.append(rend_img_side)
            rend_imgs.append(torch.from_numpy(pred_keypoints_img).permute(2, 0, 1))
            rend_imgs.append(torch.from_numpy(gt_keypoints_img).permute(2, 0, 1))
        rend_imgs = make_grid(rend_imgs, nrow=nrow, padding=padding)
        return rend_imgs

    def __call__(
        self,
        vertices,
        camera_translation,
        image,
        focal_length=5000,
        text=None,
        resize=None,
        side_view=False,
        baseColorFactor=(1.0, 1.0, 0.9, 1.0),
        rot_angle=90,
    ):
        renderer = pyrender.OffscreenRenderer(
            viewport_width=image.shape[1], viewport_height=image.shape[0], point_size=1.0
        )
        material = pyrender.MetallicRoughnessMaterial(
            metallicFactor=0.0, alphaMode="OPAQUE", baseColorFactor=baseColorFactor
        )

        camera_translation[0] *= -1.0

        mesh = trimesh.Trimesh(vertices.copy(), self.faces.copy())
        if side_view:
            rot = trimesh.transformations.rotation_matrix(np.radians(rot_angle), [0, 1, 0])
            mesh.apply_transform(rot)
        rot = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0])
        mesh.apply_transform(rot)
        mesh = pyrender.Mesh.from_trimesh(mesh, material=material)

        scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0], ambient_light=(0.3, 0.3, 0.3))
        scene.add(mesh, "mesh")

        camera_pose = np.eye(4)
        camera_pose[:3, 3] = camera_translation
        camera_center = [image.shape[1] / 2.0, image.shape[0] / 2.0]
        camera = pyrender.IntrinsicsCamera(
            fx=focal_length, fy=focal_length, cx=camera_center[0], cy=camera_center[1]
        )
        scene.add(camera, pose=camera_pose)

        light_nodes = create_raymond_lights()
        for node in light_nodes:
            scene.add_node(node)

        color, rend_depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
        color = color.astype(np.float32) / 255.0
        valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
        if not side_view:
            output_img = color[:, :, :3] * valid_mask + (1 - valid_mask) * image
        else:
            output_img = color[:, :, :3]
        if resize is not None:
            output_img = cv2.resize(output_img, resize)

        output_img = output_img.astype(np.float32)
        renderer.delete()
        return output_img
