import copy

import numpy as np


def center_normalize(pcl):
    # Center the point cloud
    centroid = np.mean(pcl, axis=0)
    pcl_centered = pcl - centroid

    # Scale
    max_dist = np.max(pcl_centered)
    pcl_normalized = pcl_centered / (2 * max_dist)

    # Move down till one the lowest point is on z=-0.5
    pcl_normalized[:, 2] += -0.5 - np.min(pcl_normalized[:, 2])

    return pcl_normalized.astype(np.float32)


try:
    import pyrender
    import trimesh
except ImportError:  # Mitsuba-only 脚本只需 center_normalize，可不装 pyrender
    pyrender = None
    trimesh = None


def pretty_color(x, y, z):
    x += 0.5
    y += +0.5
    z = z + 0.5 - 0.0125
    vec = np.array([x, y, z])
    vec = np.clip(vec, 0.001, 1.0)
    norm = np.sqrt(np.sum(vec ** 2))
    vec /= norm
    return [vec[0], vec[1], vec[2]]


# code from https://github.com/nv-nguyen/cnos/blob/main/src/poses/pyrender.py

class Render(object):

    def __init__(self, light_intensity):
        if pyrender is None:
            raise ImportError("Render 类需要安装 pyrender、trimesh：pip install pyrender trimesh")

        # camera pose is fixed as np.eye(4)
        cam_pose = np.eye(4)
        # convert openCV camera
        cam_pose[1, 1] = -1
        cam_pose[2, 2] = -1

        self.cam_pose = cam_pose

        # create scene config
        self.ambient_light = np.array(
            [0.3, 0.3, 0.3, 0.3])  # np.array([0.2, 0.2, 0.2])  # np.array([1.0, 1.0, 1.0, 1.0])

        self.light = pyrender.SpotLight(
            color=np.ones(3),
            intensity=light_intensity,
            innerConeAngle=np.pi / 2.0,
            outerConeAngle=np.pi / 2.0,
        )

        self.render_engine = {}

    def render(
            self,
            mesh,
            obj_poses,
            img_size,
            intrinsic,
            dyn_obj_scale=False,
            factor=1.5,
    ):

        scene = pyrender.Scene(
            bg_color=np.array([0., 0., 0., 1.0]), ambient_light=self.ambient_light
        )

        scene.add(self.light, pose=self.cam_pose)

        # create camera and render engine
        fx, fy, cx, cy = intrinsic[0], intrinsic[4], intrinsic[2], intrinsic[5]

        camera = pyrender.IntrinsicsCamera(
            fx=fx, fy=fy, cx=cx, cy=cy, znear=0.05, zfar=100000
        )
        scene.add(camera, pose=self.cam_pose)

        if not dyn_obj_scale:
            cad_node = scene.add(mesh, pose=np.eye(4), name="cad")

        if (img_size[1], img_size[0]) not in self.render_engine:
            self.render_engine[(img_size[1], img_size[0])] = pyrender.OffscreenRenderer(img_size[1], img_size[0])

        rgbs, depths, masks = [], [], []

        for idx_frame in range(obj_poses.shape[0]):

            if dyn_obj_scale:
                trans = obj_poses[idx_frame][:3, 3]
                dist = ((trans ** 2).sum() ** 0.5) * factor

                obj_poses[idx_frame][:3, 3] /= dist

                scaled_mesh = copy.deepcopy(mesh).apply_scale(1 / dist)
                pymesh = pyrender.Mesh.from_trimesh(as_mesh(scaled_mesh), smooth=False)

                cad_node = scene.add(pymesh, pose=np.eye(4), name="cad")

            scene.set_pose(cad_node, obj_poses[idx_frame])

            rgb, depth = self.render_engine[(img_size[1], img_size[0])].render(scene,
                                                                               pyrender.constants.RenderFlags.RGBA)
            # img = Image.fromarray(np.uint8(rgb))
            # img.save(output_dir + f"_{idx_frame:06d}.png")

            if dyn_obj_scale:
                depth *= dist
                scene.remove_node(cad_node)

            mask = depth > 0
            masks.append(mask)

            rgbs.append(rgb[..., :3])
            depths.append(depth)

        return rgbs, depths, masks


def as_mesh(scene_or_mesh):
    if trimesh is None:
        raise ImportError("as_mesh 需要 trimesh")
    if isinstance(scene_or_mesh, trimesh.Scene):
        result = trimesh.util.concatenate(
            [
                trimesh.Trimesh(vertices=m.vertices, faces=m.faces)
                for m in scene_or_mesh.geometry.values()
            ]
        )
    else:
        result = scene_or_mesh
    return result