import os
import bpy
import numpy as np
from utils.blender_util import BlenderUtils
from utils.cad_view_sample_util import CADViewSampleUtil
from utils.material_util import MaterialUtil

class DataGenerator:
    def __init__(self, config):
        self.plane_size = config["runner"]["generate"]["plane_size"] 
        self.output_dir = config["runner"]["generate"]["output_dir"]
        self.random_config = config["runner"]["generate"]["random_config"]
        self.light_and_camera_config = config["runner"]["generate"]["light_and_camera_config"]
        self.max_views = config["runner"]["generate"]["max_views"]
        self.min_views = config["runner"]["generate"]["min_views"]
        self.min_diag = config["runner"]["generate"]["min_diag"]
        self.max_diag = config["runner"]["generate"]["max_diag"]
        self.binocular_vision = config["runner"]["generate"]["binocular_vision"]
        self.target_obj = None
        self.stopped = False
        self.random_obj_list = []
        self.display_table_config = {}
        BlenderUtils.setup_scene(self.light_and_camera_config, None, self.binocular_vision)
        self.table = BlenderUtils.get_obj(BlenderUtils.TABLE_NAME)
        
    def put_display_object(self, name):
        obj_mesh_path = BlenderUtils.get_obj_path(self.obj_dir, name)
        obj = BlenderUtils.load_obj(name, obj_mesh_path)
        bpy.ops.rigidbody.object_add()
        bpy.context.object.rigid_body.type = 'ACTIVE'
        MaterialUtil.change_object_material(obj, MaterialUtil.create_mask_material(color=(0, 1.0, 0)))
        self.target_obj = obj

    def reset(self):
        self.target_obj = None
        self.random_obj_list = []
        BlenderUtils.reset_objects_and_platform()
    
    def start_render(self, diag=0):
        object_name = self.target_obj.name
        if "." in object_name:
            object_name = object_name.split(".")[0]
        scene_dir = os.path.join(self.output_dir, object_name)
        if not os.path.exists(scene_dir):
            os.makedirs(scene_dir)
        view_num = int(self.min_views + (diag - self.min_diag)/(self.max_diag - self.min_diag) * (self.max_views - self.min_views))
        view_data = CADViewSampleUtil.sample_view_data_world_space(self.target_obj, distance_range=(0.25,0.5), voxel_size=0.005, max_views=view_num, min_cam_table_included_degree = self.min_cam_table_included_degree, random_view_ratio = self.random_view_ratio )
        object_points = np.array(view_data["voxel_down_sampled_points"])
        normals = np.array(view_data["normals"])
        points_normals = np.concatenate((object_points, normals), axis=1)
        
        np.savetxt(os.path.join(scene_dir, "points_and_normals.txt"), points_normals)
        for i, cam_pose in enumerate(view_data["cam_poses"]):
            BlenderUtils.set_camera_at(cam_pose)
            BlenderUtils.render_mask(scene_dir, f"{i}", binocular_vision=self.binocular_vision, target_object = self.target_obj)
            BlenderUtils.save_cam_params(scene_dir, i, binocular_vision=self.binocular_vision)
        BlenderUtils.save_scene_info(scene_dir, self.display_table_config, object_name)
        
        MaterialUtil.change_object_material(self.target_obj, MaterialUtil.create_normal_material())
        
        for i, cam_pose in enumerate(view_data["cam_poses"]):
            BlenderUtils.set_camera_at(cam_pose)
            BlenderUtils.render_normal_and_depth(scene_dir, f"{i}", binocular_vision=self.binocular_vision, target_object = self.target_obj)
            BlenderUtils.save_cam_params(scene_dir, i, binocular_vision=self.binocular_vision)
        
        depth_dir = os.path.join(scene_dir, "depth")
        for depth_file in os.listdir(depth_dir):
            if not depth_file.endswith(".png"):
                name, _ = os.path.splitext(depth_file)
                file_path = os.path.join(depth_dir, depth_file)
                new_file_path = os.path.join(depth_dir, f"{name}.png")
                os.rename(file_path,new_file_path)
        BlenderUtils.save_blend(scene_dir)
        exit(0)
        
        return True

    def gen_scene_data(self, object_name):
        
        bpy.context.scene.frame_set(0)
        self.put_display_object(object_name)
        diag = BlenderUtils.get_obj_diag(self.target_obj.name)
        self.start_render(diag)