import os
import json
import bpy
import time
import gc
import numpy as np
import mathutils


class BlenderUtils:

    TABLE_NAME: str = "table"
    CAMERA_NAME: str = "Camera"
    CAMERA_RIGHT_NAME: str = "CameraRight"
    CAMERA_OBJECT_NAME: str = "CameraObject"
    DISPLAY_TABLE_NAME: str = "display_table"
    MESH_FILE_NAME: str = "mesh.obj"

    @staticmethod
    def get_obj_path(obj_dir, name):
        return os.path.join(obj_dir, name, BlenderUtils.MESH_FILE_NAME)

    @staticmethod
    def load_obj(name, mesh_path, scale=1):
        print(mesh_path)
        bpy.ops.wm.obj_import(filepath=mesh_path)
        loaded_object = bpy.context.selected_objects[-1]
        loaded_object.name = name
        loaded_object.data.name = name
        loaded_object.scale = (scale, scale, scale)
        bpy.ops.rigidbody.object_add()
        return loaded_object

    @staticmethod
    def get_obj(name):
        return bpy.data.objects.get(name)

    @staticmethod
    def get_obj_pose(name):
        obj = BlenderUtils.get_obj(name)
        return np.asarray(obj.matrix_world)

    @staticmethod
    def add_plane(name, location, orientation, size=10):
        bpy.ops.mesh.primitive_plane_add(size=size, location=location)
        plane = bpy.context.selected_objects[-1]
        plane.name = name
        plane.rotation_euler = orientation
        bpy.ops.rigidbody.object_add()
        bpy.context.object.rigid_body.type = "PASSIVE"

    @staticmethod
    def add_table(table_model_path):
        table = BlenderUtils.load_obj(
            BlenderUtils.TABLE_NAME, table_model_path, scale=0.01
        )
        bpy.ops.rigidbody.object_add()
        bpy.context.object.rigid_body.type = "PASSIVE"

        mat = bpy.data.materials.new(name="TableYellowMaterial")
        mat.diffuse_color = (1.0, 1.0, 0.0, 1.0)
        if len(table.data.materials) > 0:
            table.data.materials[0] = mat
        else:
            table.data.materials.append(mat)

    @staticmethod
    def setup_scene(init_light_and_camera_config, table_model_path, binocular_vision):
        bpy.context.scene.render.engine = "BLENDER_EEVEE_NEXT"
        bpy.context.scene.display.shading.show_xray = False
        bpy.context.scene.display.shading.use_dof = False
        bpy.context.scene.display.render_aa = "OFF"
        bpy.context.scene.view_settings.view_transform = "Standard"

        bpy.context.scene.eevee.use_ssr = False  # 关闭屏幕空间反射
        bpy.context.scene.eevee.use_bloom = False  # 关闭辉光
        bpy.context.scene.eevee.use_gtao = False  # 关闭环境光遮蔽
        bpy.context.scene.eevee.use_soft_shadows = False  # 关闭软阴影
        bpy.context.scene.eevee.use_shadows = False  # 关闭所有阴影
        bpy.context.scene.world.use_nodes = False  # 如果你不需要环境光，关闭环境节点

        # bpy.context.scene.eevee.use_sss = False            # 关闭次表面散射

        # 2. 设置最低的采样数
        bpy.context.scene.eevee.taa_render_samples = 1
        bpy.context.scene.eevee.taa_samples = 1
        BlenderUtils.init_light_and_camera(
            init_light_and_camera_config, binocular_vision
        )

        BlenderUtils.add_plane("plane_floor", location=(0, 0, 0), orientation=(0, 0, 0))
        BlenderUtils.add_plane("plane_ceil", location=(0, 0, 10), orientation=(0, 0, 0))
        BlenderUtils.add_plane(
            "plane_wall_1", location=(5, 0, 5), orientation=(0, np.pi / 2, 0)
        )
        BlenderUtils.add_plane(
            "plane_wall_2", location=(-5, 0, 5), orientation=(0, np.pi / 2, 0)
        )
        BlenderUtils.add_plane(
            "plane_wall_3", location=(0, 5, 5), orientation=(np.pi / 2, 0, 0)
        )
        BlenderUtils.add_plane(
            "plane_wall_4", location=(0, -5, 5), orientation=(np.pi / 2, 0, 0)
        )

        BlenderUtils.add_table(table_model_path)

    @staticmethod
    def set_light_params(light, config):
        light.location = config["location"]
        light.rotation_euler = config["orientation"]
        if light.type == "SUN":
            light.data.energy = config["power"]
        elif light.type == "POINT":
            light.data.energy = config["power"]

    @staticmethod
    def set_camera_params(camera, config, binocular_vision):

        camera_object = bpy.data.objects.new(BlenderUtils.CAMERA_OBJECT_NAME, None)
        bpy.context.collection.objects.link(camera_object)
        cameras = [bpy.data.objects.get("Camera")]
        camera.location = [0, 0, 0]
        camera.rotation_euler = [0, 0, 0]
        camera.parent = camera_object
        if binocular_vision:
            left_camera = cameras[0]
            right_camera = left_camera.copy()
            right_camera.name = BlenderUtils.CAMERA_RIGHT_NAME
            right_camera.data = left_camera.data.copy()
            right_camera.data.name = BlenderUtils.CAMERA_RIGHT_NAME
            bpy.context.collection.objects.link(right_camera)
            right_camera.parent = camera_object
            right_camera.location = [config["eye_distance"] / 2, 0, 0]
            left_camera.location = [-config["eye_distance"] / 2, 0, 0]
            binocular_angle = config["eye_angle"]
            half_angle = np.radians(binocular_angle / 2)

            left_camera.rotation_euler[1] = -half_angle
            right_camera.rotation_euler[1] = half_angle
            cameras.append(right_camera)

        for camera in cameras:
            camera.data.clip_start = config["near_plane"]
            camera.data.clip_end = config["far_plane"]

            bpy.context.scene.render.resolution_x = config["resolution"][0]
            bpy.context.scene.render.resolution_y = config["resolution"][1]
            sensor_height = 24.0
            focal_length = sensor_height / (
                2 * np.tan(np.radians(config["fov_vertical"]) / 2)
            )
            camera.data.lens = focal_length
            camera.data.sensor_width = (
                sensor_height * config["resolution"][0] / config["resolution"][1]
            )
            camera.data.sensor_height = sensor_height

    @staticmethod
    def init_light_and_camera(init_light_and_camera_config, binocular_vision):

        camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_NAME)
        BlenderUtils.set_camera_params(
            camera,
            init_light_and_camera_config[BlenderUtils.CAMERA_NAME],
            binocular_vision,
        )

    @staticmethod
    def get_obj_diag(name):
        obj = BlenderUtils.get_obj(name)
        return np.linalg.norm(obj.dimensions)

    @staticmethod
    def matrix_to_blender_pose(matrix):
        location = matrix[:3, 3]
        rotation_matrix = matrix[:3, :3]
        rotation_matrix_blender = mathutils.Matrix(rotation_matrix.tolist())
        rotation_euler = rotation_matrix_blender.to_euler()
        return location, rotation_euler

    @staticmethod
    def set_camera_at(pose):
        camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_OBJECT_NAME)
        location, rotation_euler = BlenderUtils.matrix_to_blender_pose(pose)

        camera.location = location
        camera.rotation_euler = rotation_euler

    @staticmethod
    def get_object_bottom_z(obj):
        vertices = [v.co for v in obj.data.vertices]
        vertices_world = [obj.matrix_world @ v for v in vertices]
        min_z = min([v.z for v in vertices_world])
        return min_z

    @staticmethod
    def render_normal(
        output_dir, file_name, binocular_vision=False, target_object=None
    ):
        target_cameras = [BlenderUtils.CAMERA_NAME]
        if binocular_vision:
            target_cameras.append(BlenderUtils.CAMERA_RIGHT_NAME)
            
        for cam_name in target_cameras:
            bpy.context.scene.camera = BlenderUtils.get_obj(cam_name)
            cam_suffix = "L" if cam_name == BlenderUtils.CAMERA_NAME else "R"
            scene = bpy.context.scene
            scene.render.filepath = ""

            mask_dir = os.path.join(output_dir, "normal")
            if not os.path.exists(mask_dir):
                os.makedirs(mask_dir)

            scene.render.filepath = os.path.join(
                output_dir, mask_dir, f"{file_name}_{cam_suffix}.png"
            )
            scene.render.image_settings.color_depth = "8"
            scene.render.resolution_percentage = 100
            scene.render.use_overwrite = False
            scene.render.use_file_extension = False
            scene.render.use_placeholder = False
            bpy.ops.render.render(write_still=True)

        msg = "success"
        return msg
        
    @staticmethod
    def render_mask_and_depth(
        output_dir, file_name, binocular_vision=False, target_object=None
    ):
        target_cameras = [BlenderUtils.CAMERA_NAME]
        if binocular_vision:
            target_cameras.append(BlenderUtils.CAMERA_RIGHT_NAME)
        for cam_name in target_cameras:
            bpy.context.scene.camera = BlenderUtils.get_obj(cam_name)
            cam_suffix = "L" if cam_name == BlenderUtils.CAMERA_NAME else "R"
            scene = bpy.context.scene
            scene.render.filepath = ""

            mask_dir = os.path.join(output_dir, "mask")
            if not os.path.exists(mask_dir):
                os.makedirs(mask_dir)

            scene.render.filepath = os.path.join(
                output_dir, mask_dir, f"{file_name}_{cam_suffix}.png"
            )
            scene.render.image_settings.color_depth = "8"
            scene.render.resolution_percentage = 100
            scene.render.use_overwrite = False
            scene.render.use_file_extension = False
            scene.render.use_placeholder = False

            scene.use_nodes = True
            tree = scene.node_tree

            for node in tree.nodes:
                tree.nodes.remove(node)

            rl = tree.nodes.new("CompositorNodeRLayers")

            map_range = tree.nodes.new("CompositorNodeMapRange")
            map_range.inputs["From Min"].default_value = 0.01
            map_range.inputs["From Max"].default_value = 5
            map_range.inputs["To Min"].default_value = 0
            map_range.inputs["To Max"].default_value = 1
            tree.links.new(rl.outputs["Depth"], map_range.inputs[0])

            output_depth = tree.nodes.new("CompositorNodeOutputFile")

            depth_dir = os.path.join(output_dir, "depth")
            if not os.path.exists(depth_dir):
                os.makedirs(depth_dir)
            output_depth.base_path = depth_dir
            output_depth.file_slots[0].path = f"{file_name}_{cam_suffix}.####"
            output_depth.format.file_format = "PNG"
            output_depth.format.color_mode = "BW"
            output_depth.format.color_depth = "16"
            bpy.ops.render.render(write_still=True)

        msg = "success"
        return msg

    @staticmethod
    def save_cam_params(scene_dir, idx, binocular_vision=False):
        camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_NAME)
        extrinsic = np.array(camera.matrix_world)
        cam_data = camera.data
        focal_length = cam_data.lens
        sensor_width = cam_data.sensor_width
        sensor_height = cam_data.sensor_height
        resolution_x = bpy.context.scene.render.resolution_x
        resolution_y = bpy.context.scene.render.resolution_y
        intrinsic = np.zeros((3, 3))
        intrinsic[0, 0] = focal_length * resolution_x / sensor_width  # fx
        intrinsic[1, 1] = focal_length * resolution_y / sensor_height  # fy
        intrinsic[0, 2] = resolution_x / 2.0  # cx
        intrinsic[1, 2] = resolution_y / 2.0  # cy
        intrinsic[2, 2] = 1.0
        cam_object = BlenderUtils.get_obj(BlenderUtils.CAMERA_OBJECT_NAME)
        extrinsic_cam_object = np.array(cam_object.matrix_world)
        data = {
            "extrinsic": extrinsic.tolist(),
            "extrinsic_cam_object": extrinsic_cam_object.tolist(),
            "intrinsic": intrinsic.tolist(),
            "far_plane": camera.data.clip_end,
            "near_plane": camera.data.clip_start,
        }
        if binocular_vision:
            right_camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_RIGHT_NAME)
            extrinsic_right = np.array(right_camera.matrix_world)
            print("result:", extrinsic_right)

            data["extrinsic_R"] = extrinsic_right.tolist()

        cam_params_dir = os.path.join(scene_dir, "camera_params")
        if not os.path.exists(cam_params_dir):
            os.makedirs(cam_params_dir)
        cam_params_path = os.path.join(cam_params_dir, f"{idx}.json")
        with open(cam_params_path, "w") as f:
            json.dump(data, f, indent=4)

    @staticmethod
    def reset_objects_and_platform():
        all_objects = bpy.data.objects
        keep_objects = {
            "plane_floor",
            "plane_ceil",
            "plane_wall_1",
            "plane_wall_2",
            "plane_wall_3",
            "plane_wall_4",
        }
        keep_objects.add(BlenderUtils.CAMERA_OBJECT_NAME)
        keep_objects.add(BlenderUtils.CAMERA_NAME)
        keep_objects.add(BlenderUtils.CAMERA_RIGHT_NAME)
        keep_objects.add(BlenderUtils.TABLE_NAME)

        for obj in all_objects:
            if obj.name not in keep_objects:
                bpy.data.objects.remove(obj, do_unlink=True)

        for block in bpy.data.meshes:
            if block.users == 0:
                bpy.data.meshes.remove(block)
        for block in bpy.data.materials:
            if block.users == 0:
                bpy.data.materials.remove(block)
        for block in bpy.data.images:
            if block.users == 0:
                bpy.data.images.remove(block)

        gc.collect()
        bpy.context.scene.frame_set(0)

    @staticmethod
    def save_scene_info(scene_root_dir, display_table_config, target_name):
        all_objects = bpy.data.objects
        no_save_objects = {
            "plane_floor",
            "plane_ceil",
            "plane_wall_1",
            "plane_wall_2",
            "plane_wall_3",
            "plane_wall_4",
        }
        no_save_objects.add(BlenderUtils.CAMERA_OBJECT_NAME)
        no_save_objects.add(BlenderUtils.CAMERA_NAME)
        no_save_objects.add(BlenderUtils.CAMERA_RIGHT_NAME)
        no_save_objects.add(BlenderUtils.TABLE_NAME)
        scene_info = {}
        for obj in all_objects:
            if (
                obj.name not in no_save_objects
                and obj.name != BlenderUtils.DISPLAY_TABLE_NAME
            ):
                obj_info = {
                    "location": list(obj.location),
                    "rotation_euler": list(obj.rotation_euler),
                    "scale": list(obj.scale),
                }
                scene_info[obj.name] = obj_info
        scene_info[BlenderUtils.DISPLAY_TABLE_NAME] = display_table_config
        scene_info["target_name"] = target_name
        scene_info_path = os.path.join(scene_root_dir, "scene_info.json")
        with open(scene_info_path, "w") as outfile:
            json.dump(scene_info, outfile)
