import torch
from torch import linalg as LA
import numpy as np
import cv2

from omniisaacgymenvs.tasks.base.rl_task import RLTask
from ur5e_sweeping import UR5eSweeper as Robot

from omni.isaac.core.prims import RigidPrimView
from omni.isaac.core.prims.base_sensor import BaseSensor
from omni.isaac.core.articulations import ArticulationView
from omni.isaac.core.objects import DynamicCylinder, DynamicSphere, DynamicCuboid
from omni.isaac.core.materials import PhysicsMaterial
from omni.isaac.core.utils.string import find_unique_string_name
from omni.isaac.core.utils.prims import get_prim_at_path, is_prim_path_valid
from omni.kit.viewport.utility import get_active_viewport

from skrl.utils import omniverse_isaacgym_utils

from typing import Optional, Tuple


# post_physics_step calls
# - get_observations()
# - get_states()
# - calculate_metrics()
# - is_done()
# - get_extras()


TASK_CFG = {"test": False,
            "device_id": 0,
            "headless": True,
            "sim_device": "gpu",
            "enable_livestream": False,
            "task": {"name": "SweepingUR5e",
                     "physics_engine": "physx",
                     "env": {"numEnvs": 1024,
                             "envSpacing": 1.5,
                             "episodeLength": 200,
                             "enableDebugVis": False,
                             "clipObservations": 1000.0,
                             "clipActions": 1.0,
                             "controlFrequencyInv": 4,
                             "actionScale": 2.5,
                             "dofVelocityScale": 0.1,
                             "controlSpace": "cartesian",
                            #  "goal": [0.5, 0.3, 0.0125]},
                             "goal": [0.5, 0.3, 0.0]},
                     "sim": {"dt": 0.0083,  # 1 / 120
                             "use_gpu_pipeline": True,
                             "gravity": [0.0, 0.0, -9.81],
                             "add_ground_plane": True,
                             "use_flatcache": True,
                             "enable_scene_query_support": False,
                             "enable_cameras": False,
                             "default_physics_material": {"static_friction": 1.0,
                                                         "dynamic_friction": 1.0,
                                                         "restitution": 0.0},
                             "physx": {"worker_thread_count": 4,
                                      "solver_type": 1,
                                      "use_gpu": True,
                                      "solver_position_iteration_count": 4,
                                      "solver_velocity_iteration_count": 1,
                                      "contact_offset": 0.005,
                                      "rest_offset": 0.0,
                                      "bounce_threshold_velocity": 0.2,
                                      "friction_offset_threshold": 0.04,
                                      "friction_correlation_distance": 0.025,
                                      "enable_sleeping": True,
                                      "enable_stabilization": True,
                                      "max_depenetration_velocity": 1000.0,
                                      "gpu_max_rigid_contact_count": 524288,
                                      "gpu_max_rigid_patch_count": 33554432,
                                      "gpu_found_lost_pairs_capacity": 524288,
                                      "gpu_found_lost_aggregate_pairs_capacity": 262144,
                                      "gpu_total_aggregate_pairs_capacity": 1048576,
                                      "gpu_max_soft_body_contacts": 1048576,
                                      "gpu_max_particle_contacts": 1048576,
                                      "gpu_heap_capacity": 33554432,
                                      "gpu_temp_buffer_capacity": 16777216,
                                      "gpu_max_num_partitions": 8},
                             "robot": {"override_usd_defaults": False,
                                       "fixed_base": False,
                                       "enable_self_collisions": True,
                                       "enable_gyroscopic_forces": True,
                                       "solver_position_iteration_count": 4,
                                       "solver_velocity_iteration_count": 1,
                                       "sleep_threshold": 0.005,
                                       "stabilization_threshold": 0.001,
                                       "density": -1,
                                       "max_depenetration_velocity": 1000.0,
                                       "contact_offset": 0.005,
                                       "rest_offset": 0.0},
                             "object": {"override_usd_defaults": False,
                                        "make_kinematic": False,
                                        "fixed_base": False,
                                        # "fixed_base": True,
                                        "enable_self_collisions": False,
                                        "enable_gyroscopic_forces": True,
                                        "solver_position_iteration_count": 4,
                                        "solver_velocity_iteration_count": 1,
                                        "sleep_threshold": 0.005,
                                        "stabilization_threshold": 0.001,
                                        "density": -1,
                                        # "density": 200,
                                        "max_depenetration_velocity": 1000.0,
                                        "contact_offset": 0.005,
                                        "rest_offset": 0.0},
                            #  "goal": {"override_usd_defaults": False,
                            #           "make_kinematic": False,
                            #           "fixed_base": True,
                            #           "enable_self_collisions": False,
                            #           "density": -1
                            "goal": {"override_usd_defaults": False,
                                        "disable_rigid_body_physics": True,
                                        "fixed_base": True,
                                        "enable_self_collisions": False,
                                        "enable_gyroscopic_forces": True,
                                        "solver_position_iteration_count": 4,
                                        "solver_velocity_iteration_count": 1,
                                        "sleep_threshold": 0.005,
                                        "stabilization_threshold": 0.001,
                                        "density": -1,
                                        "max_depenetration_velocity": 1000.0,
                                        "contact_offset": 0.005,
                                        "rest_offset": 0.0
                                      },
                            }
                    }
            }


class RobotView(ArticulationView):
    def __init__(self, prim_paths_expr: str, name: str = "robot_view") -> None:
        super().__init__(prim_paths_expr=prim_paths_expr, name=name, reset_xform_properties=False)

class SweepingTask(RLTask):
    def __init__(self, name, sim_config, env, offset=None) -> None:
        self._sim_config = sim_config
        self._cfg = sim_config.config
        self._task_cfg = sim_config.task_config
        self.step_num = 0
        self.dt = 1 / 120.0

        self._num_envs = self._task_cfg["env"]["numEnvs"]
        self._env_spacing = self._task_cfg["env"]["envSpacing"]
        self._action_scale = self._task_cfg["env"]["actionScale"]
        self._dof_vel_scale = self._task_cfg["env"]["dofVelocityScale"]
        self._max_episode_length = self._task_cfg["env"]["episodeLength"]
        self._control_space = self._task_cfg["env"]["controlSpace"]
        self._goal_position = self._task_cfg["env"]["goal"]

        self._tip = {0:None, 1:None, 2:None}
        self._num_objects = 4


        self.previous_objs_position = torch.zeros(self._num_objects, self.num_envs, 3)
        self.current_objs_position = torch.zeros(self._num_objects, self.num_envs, 3)

        # each variable has a dimension of N x M
        # N is the number of objects
        # M is the number of environments
        self.previous_sweeper_obj_distance = torch.zeros(self._num_objects, self._num_envs)
        self.current_sweeper_obj_distance = torch.zeros(self._num_objects, self._num_envs)
        self.previous_obj_goal_distance = torch.zeros(self._num_objects, self._num_envs)
        self.current_obj_goal_distance = torch.zeros(self._num_objects, self._num_envs)
        self.obj_moving_distance = torch.zeros(self._num_objects, self._num_envs)
        
        self.alpha = 0.15
        self.beta = 0.5
        self.gamma = 0
        self.zeta = 0.1
        self.eta = 0.25
        self.punishment = -1

        self.completion_reward = torch.zeros(self._num_envs)

        # observation and action space
        self._num_observations = 38
        # self._num_observations = 24
        if self._control_space == "joint":
            self._num_actions = 6
        elif self._control_space == "cartesian":
            self._num_actions = 3
        else:
            raise ValueError("Invalid control space: {}".format(self._control_space))

        RLTask.__init__(self, name, env)

    def set_up_scene(self, scene) -> None:
        super().set_up_scene(scene)
        self.previous_objs_position = self.previous_objs_position.to(device=self._device)
        self.current_objs_position = self.current_objs_position.to(device=self._device)
        self.previous_sweeper_obj_distance = self.previous_sweeper_obj_distance.to(device=self._device)
        self.current_sweeper_obj_distance = self.current_sweeper_obj_distance.to(device=self._device)
        self.previous_obj_goal_distance = self.previous_obj_goal_distance.to(device=self._device)
        self.current_obj_goal_distance = self.current_obj_goal_distance.to(device=self._device)
        self.obj_moving_distance = self.obj_moving_distance.to(device=self._device)
        self.completion_reward = self.completion_reward.to(device=self._device)

        self.get_robot()
        self.get_objects()
        self.get_goal()

        # robot view
        self._robots = RobotView(prim_paths_expr="/World/envs/.*/robot", name="robot_view")
        scene.add(self._robots)

        
        # hands view (cartesian)
        if self._control_space == "cartesian":
            # tip_middle(=end-effectors) view
            self._tip[0] = RigidPrimView(prim_paths_expr=f"/World/envs/.*/robot/tip_middle", name="tip_middle_view")
            scene.add(self._tip[0])
            # self.tip_index = self._robots.body_names.index(self._tip_link)

        self._tip[1] = RigidPrimView(prim_paths_expr=f"/World/envs/.*/robot/tip_1", name="tip_1_view")
        self._tip[2] = RigidPrimView(prim_paths_expr=f"/World/envs/.*/robot/tip_2", name="tip_2_view")
        scene.add(self._tip[1])
        scene.add(self._tip[2])

        # objects view
        self._objects = {}
        for i in range(self._num_objects):
            self._objects[i] = RigidPrimView(prim_paths_expr=f"/World/envs/.*/object_{i}", name=f"object_view_{i}", reset_xform_properties=False)

        # goal view
        self._goals = RigidPrimView(prim_paths_expr="/World/envs/.*/goal", name="goal_view", reset_xform_properties=False)
        scene.add(self._goals)

        self.init_data()

    def get_robot(self):
        robot = Robot(prim_path=self.default_zero_env_path + "/robot",
                      translation=torch.tensor([0.0, 0.0, 0.0]),
                      orientation=torch.tensor([1.0, 0.0, 0.0, 0.0]),
                      name="robot")
        self._sim_config.apply_articulation_settings("robot", get_prim_at_path(robot.prim_path), self._sim_config.parse_actor_config("robot"))

    def get_objects(self):
        physics_material_path = find_unique_string_name(
                    initial_name="/World/Physics_Materials/physics_material",
                    is_unique_fn=lambda x: not is_prim_path_valid(x),
                )
        
        for i in range(self._num_objects):
            obj = DynamicCuboid(
                                prim_path=self.default_zero_env_path + f"/object_{i}",
                                name=f"object_{i}",
                                position=np.array([0.50, -0.2, 0.05]),
                                orientation=np.array([1, 0, 0, 0]),
                                # scale=np.array([0.08, 0.08, 0.08]),
                                scale=np.array([0.04, 0.04, 0.04]),
                                color=np.array([255, 0, 0]),
                                density=1,
                                # physics_material=PhysicsMaterial(prim_path=physics_material_path,
                                #                                  static_friction=0.1, dynamic_friction=0.1),
                                # mass=1.0,
                                )
            # target.enable_rigid_body_physics()
            self._sim_config.apply_articulation_settings(f"object_{i}", get_prim_at_path(obj.prim_path), self._sim_config.parse_actor_config(f"object"))
            # target.set_collision_enabled(True)
            # target.enable_rigid_body_physics()

    def get_goal(self):
        goal = DynamicSphere(prim_path=self.default_zero_env_path + "/goal",
                             name="goal",
                             radius=0.025,
                             color=torch.tensor([0, 255, 0]))
        self._sim_config.apply_articulation_settings("goal", get_prim_at_path(goal.prim_path), self._sim_config.parse_actor_config("goal"))
        goal.set_collision_enabled(False)

    def init_data(self) -> None:
        self.robot_default_dof_pos = torch.tensor(np.radians([-45, -45, 60, -100, -90, -40.0]), device=self._device, dtype=torch.float32)
        # default_dof_pos [0.0, -100.0, 60.0, -90.0, -90.0, 90.0] 
        self.actions = torch.zeros((self._num_envs, self.num_actions), device=self._device)

        if self._control_space == "cartesian":
            # self.jacobians = torch.zeros((self._num_envs, 11, 6, 6), device=self._device)
            if self._robots.body_names == None:
                self.jacobians = torch.zeros((self._num_envs, 11, 6, 6), device=self._device)
            else:
                self.jacobians = torch.zeros((self._num_envs,
                                              self._robots.body_names.index("tip_middle"),
                                              6,
                                              6), device=self._device)
                '''jacobian : (self._num_envs, num_of_bodies-1, wrench, num of joints)
                num_of_bodies - 1 due to start from 0 index'''
            # self.hand_pos, self.hand_rot = torch.zeros((self._num_envs, 3), device=self._device), torch.zeros((self._num_envs, 4), device=self._device)
            self.tip_pos, self.tip_rot = torch.zeros((self._num_envs, 3), device=self._device), torch.zeros((self._num_envs, 4), device=self._device)

    def get_observations(self) -> dict:
        robot_dof_pos = self._robots.get_joint_positions(clone=False)
        robot_dof_vel = self._robots.get_joint_velocities(clone=False)

        tip_pos = {0:None, 1: None, 2: None}    # 0 is a middle part of the sweeper
        tip_rot = {0:None, 1: None, 2: None}
        for i in range(3):
            tip_pos[i], tip_rot[i] = self._tip[i].get_local_poses()

        obj_pos = {}
        obj_rot = {}
        for i in range(self._num_objects):
            obj_pos[i], obj_rot[i] = self._objects[i].get_local_poses()
            print(f"obj_pos[{i}] : \n", self._objects[i].get_world_poses()[0])

        goal_pos, goal_rot = self._goals.get_local_poses()

        # compute objs pose and moving distance
        if torch.mean(self.previous_objs_position) == 0:
            for i in range(self._num_objects):
                self.previous_objs_position[i] = obj_pos[i]
                self.current_objs_position[i] = obj_pos[i]
        else:
            for i in range(self._num_objects):
                self.previous_objs_position[i] = self.current_objs_position[i]
                self.current_objs_position[i] = obj_pos[i]
                self.obj_moving_distance[i] = LA.norm(self.previous_objs_position[i]-self.current_objs_position[i], ord=2, dim=1)

        # compute distance for calculate_metrics() and is_done()
        if torch.mean(self.previous_obj_goal_distance) == 0:
            for i in range(self._num_objects):  # 3 is a number of sweeper parts
                self.current_sweeper_obj_distance[i] = LA.norm(tip_pos[0]-obj_pos[i], ord=2, dim=1)
                self.previous_sweeper_obj_distance[i] = self.current_sweeper_obj_distance[i]
                self.current_obj_goal_distance[i] = LA.norm(obj_pos[i]-goal_pos, ord=2, dim=1)
                self.previous_obj_goal_distance[i] = self.current_obj_goal_distance[i]
        else:                       
            for i in range(self._num_objects):  # 3 is a number of sweeper parts
                self.previous_sweeper_obj_distance[i] = self.current_sweeper_obj_distance[i]
                self.current_sweeper_obj_distance[i] = LA.norm(tip_pos[0]-obj_pos[i], ord=2, dim=1)
                self.previous_obj_goal_distance[i] = self.current_obj_goal_distance[i]
                self.current_obj_goal_distance[i] = LA.norm(obj_pos[i]-goal_pos, ord=2, dim=1)

        dof_pos_scaled = 2.0 * (robot_dof_pos - self.robot_dof_lower_limits) \
            / (self.robot_dof_upper_limits - self.robot_dof_lower_limits) - 1.0
        dof_vel_scaled = robot_dof_vel * self._dof_vel_scale

        generalization_noise = torch.rand((dof_vel_scaled.shape[0], 6), device=self._device) + 0.5

        self.obs_buf[:, 0] = self.progress_buf / self._max_episode_length
        self.obs_buf[:, 1:7] = dof_pos_scaled[:, :6]
        self.obs_buf[:, 7:13] = dof_vel_scaled[:, :6] * generalization_noise
        self.obs_buf[:, 13:16] = tip_pos[0] # 0 is a middle part of the sweeper
        self.obs_buf[:, 16:19] = tip_pos[1]
        self.obs_buf[:, 19:22] = tip_pos[2]

        self.obs_buf[:, 22:25] = obj_pos[0]
        self.obs_buf[:, 25:28] = obj_pos[1]
        self.obs_buf[:, 28:31] = obj_pos[2]
        self.obs_buf[:, 31:34] = obj_pos[2]
        self.obs_buf[:, 34:37] = goal_pos
        
        # for i in range(self._num_objects):
        #     print(f'obj_pos{i}: \n',obj_pos[i])
        
        if self._control_space == "cartesian":
            self.jacobians = self._robots.get_jacobians(clone=False)

        return {self._robots.name: {"obs_buf": self.obs_buf}}

    def pre_physics_step(self, actions) -> None:
        reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
        if len(reset_env_ids) > 0:
            self.reset_idx(reset_env_ids)

        self.actions = actions.clone().to(self._device)
        env_ids_int32 = torch.arange(self._robots.count, dtype=torch.int32, device=self._device)

        if self._control_space == "joint":
            targets = self.robot_dof_targets[:, :6] + self.robot_dof_speed_scales[:6] * self.dt * self.actions * self._action_scale

        elif self._control_space == "cartesian":
            goal_position = self.tip_pos[0] + actions / 100.0
            delta_dof_pos = omniverse_isaacgym_utils.ik(
                                                        jacobian_end_effector=self.jacobians[:, 11-1, :, :],
                                                        current_position=self.tip_pos,
                                                        current_orientation=self.tip_rot,
                                                        goal_position=goal_position,
                                                        goal_orientation=None
                                                        )
            targets = self.robot_dof_targets[:, :6] + delta_dof_pos

        self.robot_dof_targets[:, :6] = torch.clamp(targets, self.robot_dof_lower_limits[:6], self.robot_dof_upper_limits[:6])
        # self.robot_dof_targets[:, 6:] = 0
        self._robots.set_joint_position_targets(self.robot_dof_targets, indices=env_ids_int32)



    def reset_idx(self, env_ids) -> None:
        indices = env_ids.to(dtype=torch.int32)

        # reset robot
        pos = torch.clamp(self.robot_default_dof_pos.unsqueeze(0) + 0.25 * (torch.rand((len(env_ids), self.num_robot_dofs), device=self._device) - 0.5),
                          self.robot_dof_lower_limits, self.robot_dof_upper_limits)
        dof_pos = torch.zeros((len(indices), self._robots.num_dof), device=self._device)
        dof_pos[:, :] = pos
        dof_vel = torch.zeros((len(indices), self._robots.num_dof), device=self._device)
        self.robot_dof_targets[env_ids, :] = pos
        self.robot_dof_pos[env_ids, :] = pos

        self._robots.set_joint_position_targets(self.robot_dof_targets[env_ids], indices=indices)
        self._robots.set_joint_positions(dof_pos, indices=indices)
        self._robots.set_joint_velocities(dof_vel, indices=indices)


        all_env_pos = torch.zeros(self._num_objects, self._num_envs, 3, device=self._device)
        all_env_ori = torch.zeros(self._num_objects, self._num_envs, 4, device=self._device)
        for i in range(self._num_objects):
            single_pos = torch.tensor([0.60 - (0.05*(i+1)), 0.2, 0.08], device=self._device)
            each_env_pos = single_pos.repeat(self._num_envs, 1, 1).transpose(0, 1).squeeze()
            all_env_pos[i] = each_env_pos

            single_ori = torch.tensor([1.0, 0.0, 0.0, 0.0], device=self._device)
            each_env_ori = single_ori.repeat(self._num_envs, 1, 1).transpose(0, 1).squeeze()
            all_env_ori[i] = each_env_ori

            self._objects[i].set_world_poses(each_env_pos + self._env_pos[env_ids],
                                             each_env_ori,
                                             indices=indices)
 

        # reset goal
        # goal_pos_xy_variation = torch.rand((len(env_ids), 2), device=self._device) * 0.1
        # goal_pos_z_variation = torch.mul(torch.ones((len(env_ids), 1), device=self._device), 0.025)
        # goal_pos_variation = torch.cat((goal_pos_xy_variation, goal_pos_z_variation), dim=1)
        # goal_pos = torch.tensor(self._goal_position, device=self._device) + goal_pos_variation
        goal_pos = torch.tensor(self._goal_position, device=self._device)
        self._goals.set_world_poses(goal_pos + self._env_pos[env_ids], indices=indices)

        # reset dummy
        # self._dummy.set_world_poses()

        # bookkeeping
        self.reset_buf[env_ids] = 0
        self.progress_buf[env_ids] = 0

    def post_reset(self):
        self.num_robot_dofs = self._robots.num_dof
        self.robot_dof_pos = torch.zeros((self.num_envs, self.num_robot_dofs), device=self._device)
        dof_limits = self._robots.get_dof_limits()
        self.robot_dof_lower_limits = dof_limits[0, :, 0].to(device=self._device)
        self.robot_dof_upper_limits = dof_limits[0, :, 1].to(device=self._device)
        self.robot_dof_speed_scales = torch.ones_like(self.robot_dof_lower_limits)
        self.robot_dof_targets = torch.zeros((self._num_envs, self.num_robot_dofs), dtype=torch.float, device=self._device)
        # self._targets.enable_rigid_body_physics()
        # self._targets.enable_gravities()

        # randomize all envs
        indices = torch.arange(self._num_envs, dtype=torch.int64, device=self._device)
        self.reset_idx(indices)

    def calculate_metrics(self) -> None:
        # self.rew_buf[:] = -self._computed_tool_target_distance - self._computed_target_goal_distance
        sweeper_obj_distance_reward = 1 / (torch.min(self.current_sweeper_obj_distance, dim=0)[0] + 1e-5)
        obj_goal_distance_reward = 1 / (torch.min(self.current_obj_goal_distance, dim=0)[0] + 1e-5)

        delta_tool_target = torch.mean(self.previous_sweeper_obj_distance, dim=0) - \
                            torch.mean(self.current_sweeper_obj_distance, dim=0)
        delta_target_goal = torch.mean(self.previous_obj_goal_distance, dim=0) - \
                            torch.mean(self.current_obj_goal_distance, dim=0)

        avoiding_non_touch = torch.abs(1 / (torch.mean(self.obj_moving_distance, dim=0) + 0.5))

        self.completion_reward[torch.max(self.current_obj_goal_distance, dim=0)[0]<0.055] = 1
        self.rew_buf[:] = self.alpha * sweeper_obj_distance_reward + \
                          self.beta * obj_goal_distance_reward + \
                          self.gamma * self.completion_reward + \
                          self.zeta * delta_tool_target + \
                          self.eta * delta_target_goal + \
                          self.punishment * avoiding_non_touch        
        

    def is_done(self) -> None:
        self.reset_buf.fill_(0)
        # target reached
        # self.reset_buf = torch.where(self._computed_target_goal_distance <= 0.035, torch.ones_like(self.reset_buf), self.reset_buf)
        self.reset_buf = torch.where(torch.max(self.current_obj_goal_distance, dim=0)[0] <= 0.05,
                                     torch.ones_like(self.reset_buf),
                                     self.reset_buf)
        
        # max episode length
        self.reset_buf = torch.where(self.progress_buf >= self._max_episode_length - 1, torch.ones_like(self.reset_buf), self.reset_buf)
