import sys
import yaml
import omni
from omni.isaac.kit import SimulationApp
import carb
from torch._C import NoneType
import os
import time
import random
import numpy as np
import trimesh
from sklearn import linear_model
import copy
import math
from generate_objects_utils import *

class GENERATE_BOJECTS():
    def __init__(self):
        # Initialize Parameters
        self._read_params()

        # Initialize Isaac Sim Simulator
        self._setup_isaac_sim()

        # Initialize Variables and Parameters Related to Data Generation
        self._setup_data_generatin_params()

        # Start Main Function
        self._main()
    
    def _read_params(self):
        # Load Parameters for Launching Isaac Sim
        _isaac_sim_config_path = sys.path[0]+'/../config/isaac_sim/isaac_sim.yaml'
        with open(_isaac_sim_config_path, 'r') as config:
            self._isaac_sim_config = yaml.safe_load(config)

        # Load Parameters for Data Generation
        _data_generation_config_path = sys.path[0]+'/../config/data_generation/data_generation.yaml'
        with open(_data_generation_config_path, 'r') as config:
            self._data_generation_config = yaml.safe_load(config)
        if(self._data_generation_config["data_save_path"][-1]=="/"):
            self._data_generation_config["data_save_path"][-1] = self._data_generation_config["data_save_path"][:-1]

        # Load Parameters for Replicator
        _replicator_config_path = sys.path[0]+'/../config/replicator/replicator.yaml'
        with open(_replicator_config_path, 'r') as config:
            self._replicator_config = yaml.safe_load(config)
        if(self._replicator_config["writer_config"]["output_dir"][-1]=="/"):
            self._replicator_config["writer_config"]["output_dir"] = self._replicator_config["writer_config"]["output_dir"][:-1]

        # Calculate Parameters
        self._num_r = int((self._data_generation_config["max_radius"]-self._data_generation_config["min_radius"])/self._data_generation_config["dr"])
        # carb.log_error("num_r: "+str(self._data_generation_config["max_radius"]-self._data_generation_config["min_radius"]))
        self._num_h = int(self._data_generation_config["hfov"]/self._data_generation_config["dhfov"])
        self._num_v = int(360/self._data_generation_config["dvfov"])

    def _setup_isaac_sim(self):
        # Start the omniverse application
        self._simulation_app = SimulationApp(launch_config=self._isaac_sim_config)

        # The following items must be imported after simulation is initialized
        from omni.isaac.core import World
        from omni.isaac.dynamic_control import _dynamic_control

        # Acquire dynamic control interface
        self._dc = _dynamic_control.acquire_dynamic_control_interface()

        # Disable joystick to avoid conflict
        carb.settings.get_settings().set("persistent/app/omniverse/gamepadCameraControl", False)

        # Create a new world
        self._world = World(physics_dt=self._isaac_sim_config["physics_dt"],
                            rendering_dt=self._isaac_sim_config["rendering_dt"],
                            stage_units_in_meters=self._isaac_sim_config["meters_unit"])
        if(self._world is NoneType): 
            carb.log_error("Failed to set world.")
        self._stage = omni.usd.get_context().get_stage()
        self._simulation_app.update()

        # Start Simulation
        omni.timeline.get_timeline_interface().play()

    def _setup_replicator(self, scene_id = 0):
        # Initialize Replicator
        import omni.replicator.core as rep
        self._rep = rep

        # Setup Cameras
        self._driver_cam = self._rep.create.camera(name="DriverCam")
        self._driver_rp = self._rep.create.render_product(self._driver_cam, self._replicator_config["resolution"], name="DriverView")

        # Setup Writers
        output_path = self._replicator_config["writer_config"]["output_dir"]
        self._replicator_config["writer_config"]["output_dir"] = output_path + "/scene_" + str(scene_id)
        self._writer = self._rep.WriterRegistry.get(self._replicator_config["writer"])
        self._writer.initialize(**self._replicator_config["writer_config"])
        self._writer.attach(self._driver_rp)

    def _setup_data_generatin_params(self):
        self._trimesh_objects = []
        self._trimesh_occluders = []
        self._occluder_top_z = 0
        self._occluder_spawn_x = 0
        self._occluder_spawn_y = 0
        self._occluder_spawn_z = 0
        self._occluder_x_lb = 0
        self._occluder_x_ub = 0
        self._occluder_y_lb = 0
        self._occluder_y_ub = 0

        self._selected_occluder_index = 0
        self._selected_object_indices = []

        self._replicator_objects = []

    def _extract_model_type(self, path):
        # Specify it according to specific file directory
        NUM_SLASH_BEFORE_TYPE = 2

        num_slash = 0
        object_type_str = []

        for i in range(len(path)):
            index = len(path) -1 - i
            char = path[index]

            if(num_slash == NUM_SLASH_BEFORE_TYPE):
                object_type_str.append(char)

            if(char == "/"):
                num_slash += 1

        object_type_str.reverse()
        object_type_str = ''.join(object_type_str[1:])
        return object_type_str

    def _extract_model_name(self, path):
        # Specify it according to specific file directory
        NUM_SLASH_BEFORE_NAME = 1

        num_slash = 0
        object_name_str = []

        for i in range(len(path)):
            index = len(path) -1 - i
            char = path[index]

            if(num_slash == NUM_SLASH_BEFORE_NAME):
                object_name_str.append(char)

            if(char == "/"):
                num_slash += 1

        object_name_str.reverse()
        object_name_str = ''.join(object_name_str[1:])
        return object_name_str

    def _add_background(self, scene_usd_path):
        from omni.isaac.core.utils.stage import add_reference_to_stage
        from omni.isaac.core.prims import XFormPrim
        add_reference_to_stage(usd_path=scene_usd_path, prim_path="/World/Env")
        self._world.scene.add(XFormPrim(prim_path="/World/Env", name=scene_usd_path, 
                                        position=[0,0,0], orientation=[0,0,0,1], scale=[1,1,1]))
        
    def _set_background_position(self, position):
        _object_prim = self._stage.GetPrimAtPath("/World/Env")
        _object_prim.GetAttribute("xformOp:translate").Set((position[0], position[1], position[2]))

    def _delete_background(self):
        import omni.isaac.core.utils.prims as prims_utils
        prims_utils.delete_prim("/World/Env")

    def _add_object_to_sim(self, obj_index):
        import omni.replicator.core as rep

        # Load Object
        trimesh_obj = self._trimesh_objects[obj_index]

        # Calculate proper spawn pose
        bounding_box = trimesh_obj.bb
        # spawn_z = abs(0-bounding_box[0][1]) + self._occluder_top_z + 0.01
        
        convex_hulls = self._trimesh_occluders[self._selected_occluder_index].convex_hulls
        num_planes = len(convex_hulls)
        if(num_planes >= self._data_generation_config["max_num_planes"]):
            num_planes = self._data_generation_config["max_num_planes"]
        random_plane_index = random.randint(0, num_planes-1)
        convex_hull = convex_hulls[random_plane_index]

        # calculate spawn height
        plane_z = convex_hull.get_average_z()
        spawn_z = self._occluder_spawn_z + plane_z + -bounding_box[0][2] + self._data_generation_config["spawn_z_offset"]
        if(spawn_z > self._data_generation_config["max_spawn_z"]):
            # spawn_z = self._data_generation_config["max_spawn_z"]
            return False
        # carb.log_error("Spawning Object: "+ trimesh_obj.name+", z: "+str(spawn_z))

        # get spawning plane bounding box
        convex_hull_bounds = convex_hull.get_convex_hull()
        convex_hull_x_bounds = np.asarray(convex_hull_bounds).T[0]
        convex_hull_y_bounds = np.asarray(convex_hull_bounds).T[1]
        convex_min_x = min(convex_hull_x_bounds)
        convex_max_x = max(convex_hull_x_bounds)
        convex_min_y = min(convex_hull_y_bounds)
        convex_max_y = max(convex_hull_y_bounds)

        # get object bounding box
        bb_origin = [(bounding_box[0][0]+bounding_box[1][0])/2, (bounding_box[0][0]+bounding_box[1][0])/2]
        bb_width = abs(bounding_box[0][0]-bounding_box[1][0])
        bb_height = abs(bounding_box[0][1]-bounding_box[1][1])
        # carb.log_error("origin: "+str(bb_origin)+", width: "+str(bb_width)+", height: "+str(bb_height))

        for i in range(self._data_generation_config["num_attemp_times"]):
            random_origin_x = random.uniform(convex_min_x, convex_max_x)
            random_origin_y = random.uniform(convex_min_y, convex_max_y)
            bb_origin[0] += random_origin_x
            bb_origin[1] += random_origin_y
            rot_z = random.uniform(-np.pi,np.pi)
            bb = BOUNDING_BOX_2D(bb_origin, bb_width, bb_height, rot_z)
            if(convex_hull.are_points_inside_convex_hull(bb._get_bb_vertices())):
                # carb.log_error("i: "+str(i)+", bb: "+str(bb._get_bb_vertices())+", convex_hull: "+str(convex_hull_bounds[:3]) )
                place_x = bb_origin[0]
                place_y = bb_origin[1]

                # Spawn Object
                object = rep.create.from_usd(trimesh_obj.usd_path, semantics=[("class", trimesh_obj.name)])
                with object:
                    rep.modify.pose(position=(place_x, place_y, spawn_z), rotation_x = -180, rotation_z = -90, scale=self._data_generation_config["model_scaling"])
                self._replicator_objects.append(object)

                return True
        return False

    def _add_occluder_to_sim(self, obj_index):
        from omni.isaac.core.utils.stage import add_reference_to_stage
        from omni.isaac.core.prims import XFormPrim
        from omni.physx.scripts import utils
        from pxr import UsdPhysics

        # Load Object
        trimesh_occluder = self._trimesh_occluders[obj_index]
        add_reference_to_stage(usd_path=trimesh_occluder.usd_path, prim_path="/World/"+trimesh_occluder.name)

        # Make it a rigid body
        occluder_prim = self._stage.GetPrimAtPath("/World/"+trimesh_occluder.name)
        utils.setRigidBody(occluder_prim, "convexDecomposition", False)
        mass_api = UsdPhysics.MassAPI.Apply(occluder_prim)

        # Get Bounding Box
        bounding_box = trimesh_occluder.bb

        self._occluder_top_z =abs(bounding_box[1][1]-bounding_box[0][1])
        self._occluder_x_lb = bounding_box[0][0]
        self._occluder_x_ub = bounding_box[1][0]
        self._occluder_y_lb = bounding_box[0][2]
        self._occluder_y_ub = bounding_box[1][2]

        # Add to Sim
        spawn_z = abs(bounding_box[0][1]) + 0.01
        if(spawn_z > self._data_generation_config["max_spawn_z"]):
            spawn_z = self._data_generation_config["max_spawn_z"]
        self._occluder_spawn_z = spawn_z
        carb.log_warn("Spawning Occluder: "+ trimesh_occluder.name+", z: "+str(spawn_z))
        self._world.scene.add(XFormPrim(prim_path="/World/"+trimesh_occluder.name, 
                                        name=trimesh_occluder.name, 
                                        position=[self._occluder_spawn_x, self._occluder_spawn_y, self._occluder_spawn_z], 
                                        orientation=[0.7071068, 0.7071068, 0, 0], 
                                        scale=self._data_generation_config["model_scaling"]))

    def _read_objects(self, dataset_path):
        for root, dirs, files in os.walk(dataset_path, topdown=False):
            for name in dirs:
                path = os.path.join(root, name)
                obj_label=self._extract_model_type(path)
                obj_name=self._extract_model_name(path)
                if(os.path.join(root, name)[-14:] == "Scan_converted"):
                    # self._objs[path+'/Simp_obj.usd'] = [obj_label, obj_name]
                    self._trimesh_objects.append(Object(obj_path=path[:-10]+'/Simp.obj', 
                                                        usd_path = path+'/Simp_obj.usd',
                                                        scale=self._data_generation_config["model_scaling"],
                                                        eular_angles=[1.57, 0, 0],
                                                        label=obj_label, 
                                                        name=obj_name))
                    
    def _read_occluders(self, dataset_path):
        for root, dirs, files in os.walk(dataset_path, topdown=False):
            for name in dirs:
                path = os.path.join(root, name)
                occluder_label=self._extract_model_type(path)
                occluder_name=self._extract_model_name(path)
                if(os.path.join(root, name)[-14:] == "Scan_converted"):
                    self._trimesh_occluders.append(Object(obj_path=path[:-10]+'/Simp.obj', 
                                                        usd_path = path+'/Simp_obj.usd',
                                                        scale=self._data_generation_config["model_scaling"],
                                                        eular_angles=[1.57, 0, 0],
                                                        label=occluder_label, 
                                                        name=occluder_name))

    def _spawn_objects(self):
        from omni.physx.scripts import utils
        from pxr import UsdPhysics
        import omni.replicator.core as rep

        num_spawned_objects = 0
        usd_spawned = []
        while(num_spawned_objects < self._data_generation_config["object_num"]):
            random_index = random.randint(0, len(self._trimesh_objects)-1)
            if(self._trimesh_objects[random_index].usd_path in usd_spawned):
                pass
            else:
                if(self._simulation_app.is_running()):
                    self._trimesh_objects[random_index]._load_mesh()
                    if(self._add_object_to_sim(random_index)):
                        usd_spawned.append(self._trimesh_objects[random_index].usd_path)
                        num_spawned_objects += 1
                        self._selected_object_indices.append(random_index)
        
        # Make objects all rigid bodies
        for object in self._replicator_objects:
            prim_path = rep.utils.get_node_targets(object.node, "inputs:primsIn")[0]
            object_prim = self._stage.GetPrimAtPath(prim_path)
            utils.setRigidBody(object_prim, "triangleMesh", False)
            mass_api = UsdPhysics.MassAPI.Apply(object_prim)
            self._step(10)
        self._step(100)

    def _spawn_occluder(self):
        random_index = random.randint(0, len(self._trimesh_occluders)-1)
        self._selected_occluder_index = random_index
        if(self._simulation_app.is_running()):
            self._trimesh_occluders[random_index]._load_mesh()
            carb.log_error(self._trimesh_occluders[random_index]._mesh)
            self._trimesh_occluders[random_index]._find_placable_areas()
            self._add_occluder_to_sim(random_index)
            self._step(10)

    def _step(self, step_times, render=True, step_sim=True):
        for i in range(step_times):
            if(self._simulation_app.is_running()):
                self._world.step(render=render, step_sim=step_sim)

    def _export_usd_scene(self, scene_path, scene_id=0):
        self._stage.GetRootLayer().Export(scene_path+"/scene_"+str(scene_id)+"/scene.usd")

    def _generate_dataset(self):
        c = 0
        for r in range(self._num_r):
            for h in range(self._num_h):
                for v in range(self._num_v):
                    sampling_radius = self._data_generation_config["min_radius"] + r*self._data_generation_config["dr"]
                    h_angle_in_rad = (self._data_generation_config["start_h_angle"] + self._data_generation_config["dhfov"] * h) / 180 * math.pi
                    v_angle_in_rad = (0 + self._data_generation_config["dvfov"] * v) / 180 * math.pi
                    look_at_position = (self._data_generation_config["look_at_position"][0]+((random.random()-0.5)*self._data_generation_config["look_at_noise"]),
                                        self._data_generation_config["look_at_position"][1]+((random.random()-0.5)*self._data_generation_config["look_at_noise"]),
                                        self._data_generation_config["look_at_position"][2]+((random.random()-0.5)*self._data_generation_config["look_at_noise"]),)
                    # carb.log_warn("r: "+str(sampling_radius)+", h: "+str(h_angle_in_rad)+", v: "+str(v_angle_in_rad))

                    px = self._data_generation_config["look_at_position"][0] + sampling_radius*math.sin(v_angle_in_rad)*math.cos(h_angle_in_rad)
                    py = self._data_generation_config["look_at_position"][1] + sampling_radius*math.cos(v_angle_in_rad)*math.cos(h_angle_in_rad)
                    pz = self._data_generation_config["look_at_position"][2] + sampling_radius*math.sin(h_angle_in_rad)
                    dx = look_at_position[0] - px
                    dy = look_at_position[1] - py
                    dz = look_at_position[2] - pz
                    roll = (random.random()-0.5)*30
                    pitch = math.atan2(dz, math.sqrt(dx**2+dy**2))/math.pi*180
                    yaw = math.atan2(-dy, -dx)/math.pi*180

                    _camera_prim = self._stage.GetPrimAtPath("/Replicator/DriverCam_Xform")
                    _camera_prim.GetAttribute("xformOp:translate").Set((px, py, pz))
                    _camera_prim.GetAttribute("xformOp:rotateXYZ").Set((roll, pitch, yaw))
                    carb.log_warn("frame: "+str(c)+"x: "+str(px)+", y: "+str(py)+", z: "+str(pz)+", roll: "+str(roll)+", pitch: "+str(pitch)+", yaw: "+str(yaw))
                    self._rep.orchestrator.step()
                    self._step(1)
                    c+=1    

    def _is_scene_valid(self):
        trimesh_occluder = self._trimesh_occluders[self._selected_occluder_index]
        occluder_prim_path = "/World/"+trimesh_occluder.name
        occluder_dc_object = self._dc.get_rigid_body(occluder_prim_path)
        occluder_pose = self._dc.get_rigid_body_pose(occluder_dc_object)

        occluder_actual_position = occluder_pose.p
        occluder_reference_position = [self._occluder_spawn_x, self._occluder_spawn_y, self._occluder_spawn_z]
        # carb.log_error("current occluder pose x: "+str(occluder_pose.p[0])+
        #                ", current occluder pose y: "+str(occluder_pose.p[1])+
        #                ", current occluder pose z: "+str(occluder_pose.p[2]))
        # carb.log_error("reference occluder pose x: "+str(self._occluder_spawn_x)+
        #                ", reference occluder pose y: "+str(self._occluder_spawn_y)+
        #                ", reference occluder pose z: "+str(self._occluder_spawn_z))
        distance = distance_3d(occluder_actual_position, occluder_reference_position)
        carb.log_warn("distance between occluder's reference and actual positions: "+str(distance))
        if(distance >= self._data_generation_config["distance_threshold"]):
            carb.log_error("Occluder is moved too much after spawning objects, scene generation failed!")
            return False
        return True
    
    def _refresh_objects_and_occluders(self):
        self._occluder_top_z = 0
        self._occluder_spawn_x = 0
        self._occluder_spawn_y = 0
        self._occluder_spawn_z = 0
        self._occluder_x_lb = 0
        self._occluder_x_ub = 0
        self._occluder_y_lb = 0
        self._occluder_y_ub = 0

        self._selected_occluder_index = 0
        self._selected_object_indices = []

        self._replicator_objects = []

    def _main(self):
        self._add_background(self._data_generation_config["background_path"])
        self._read_objects(self._data_generation_config["object_dataset_path"])
        self._read_occluders(self._data_generation_config["occluder_dataset_path"])
        # c = 0
        # while(c < self._data_generation_config["num_scenes"]):
        #     self._spawn_occluder()
        #     self._spawn_objects()
        #     if(self._is_scene_valid()):
        #         self._setup_replicator(scene_id=c)
        #         self._generate_dataset()
        #         self._export_usd_scene(self._data_generation_config["data_save_path"], scene_id=c)
        #     self._refresh_objects_and_occluders()

        self._spawn_occluder()
        self._spawn_objects()
        self._setup_replicator(scene_id=0)
        self._generate_dataset()
        self._export_usd_scene(self._data_generation_config["data_save_path"], scene_id=0)
        self._step(10000)
        self._simulation_app.close()


abba = GENERATE_BOJECTS()
