from collections import OrderedDict
import numpy as np

from robosuite.utils.transform_utils import convert_quat
import robosuite.utils.transform_utils as T
from robosuite.environments.sawyer import SawyerEnv

from robosuite.models.arenas import TableArena, BinsArena
from robosuite.models.objects import (
    BoxObject, 
    MilkObject,
    BreadObject,
    CerealObject,
    CanObject,
)
from robosuite.models.robots import Sawyer
from robosuite.models.tasks import TableTopTask, UniformRandomSampler
import random


class SawyerLiftBinobs(SawyerEnv):
    """
    This class corresponds to the lifting task for the Sawyer robot arm.
    Objects and observations used are the same to SawyerPickPlace task,
    rewards are the same to SawyerLift.
    """

    def __init__(
        self,
        gripper_type="TwoFingerGripper",
        table_full_size=(0.39, 0.49, 0.82),
        table_friction=(1., 5e-3, 1e-4),
        use_camera_obs=True,
        use_object_obs=True,
        reward_shaping=False,
        placement_initializer=None,
        gripper_visualization=False,
        use_indicator_object=False,
        has_renderer=False,
        has_offscreen_renderer=True,
        render_collision_mesh=False,
        render_visual_mesh=True,
        control_freq=10,
        horizon=1000,
        ignore_done=False,
        camera_name="frontview",
        camera_height=256,
        camera_width=256,
        camera_depth=False,
        object_choice='random',
        arena_type='bin'
    ):
        """
        Args:

            gripper_type (str): type of gripper, used to instantiate
                gripper models from gripper factory.

            table_full_size (3-tuple): x, y, and z dimensions of the table.

            table_friction (3-tuple): the three mujoco friction parameters for
                the table.

            use_camera_obs (bool): if True, every observation includes a
                rendered image.

            use_object_obs (bool): if True, include object (cube) information in
                the observation.

            reward_shaping (bool): if True, use dense rewards.

            placement_initializer (ObjectPositionSampler instance): if provided, will
                be used to place objects on every reset, else a UniformRandomSampler
                is used by default.

            gripper_visualization (bool): True if using gripper visualization.
                Useful for teleoperation.

            use_indicator_object (bool): if True, sets up an indicator object that
                is useful for debugging.

            has_renderer (bool): If true, render the simulation state in
                a viewer instead of headless mode.

            has_offscreen_renderer (bool): True if using off-screen rendering.

            render_collision_mesh (bool): True if rendering collision meshes
                in camera. False otherwise.

            render_visual_mesh (bool): True if rendering visual meshes
                in camera. False otherwise.

            control_freq (float): how many control signals to receive
                in every second. This sets the amount of simulation time
                that passes between every action input.

            horizon (int): Every episode lasts for exactly @horizon timesteps.

            ignore_done (bool): True if never terminating the environment (ignore @horizon).

            camera_name (str): name of camera to be rendered. Must be
                set if @use_camera_obs is True.

            camera_height (int): height of camera frame.

            camera_width (int): width of camera frame.

            camera_depth (bool): True if rendering RGB-D, and RGB otherwise.

            object_choice (str): [ 'milk', 'cereal', 'bread', 'can', 'random']
        """

        # settings for table top
        self.table_full_size = table_full_size
        self.table_friction = table_friction
        self.arena_type = arena_type

        # whether to use ground-truth object states
        self.use_object_obs = use_object_obs

        # reward configuration
        self.reward_shaping = reward_shaping

        # object placement initializer
        if placement_initializer:
            self.placement_initializer = placement_initializer
        elif self.arena_type == 'bin':
            arena = BinsArena(table_full_size, table_friction)
            bin_abs = arena.bin_abs
            self.placement_initializer = UniformRandomSampler(
                x_range=[bin_abs[0]-(table_full_size[0]/2 - 0.05), bin_abs[0]+table_full_size[0]/2 - 0.05],
                y_range=[bin_abs[1]-(table_full_size[1]/2 - 0.05), bin_abs[1]+table_full_size[1]/2 - 0.05],
                #x_range=[bin_abs[0]-0.0, bin_abs[0]+0.0],
                #y_range=[bin_abs[1]-0.0, bin_abs[1]+0.0],
                ensure_object_boundary_in_range=False,
                z_rotation=True,
            )
        elif self.arena_type == 'table':
            self.placement_initializer = UniformRandomSampler(
                x_range=[-0.03, 0.03],
                y_range=[-0.03, 0.03],
                #x_range=[-0.3, 0.3],
                #y_range=[-0.3, 0.3],
                ensure_object_boundary_in_range=False,
                z_rotation=True,
            )
        else:
            raise NotImplementedError
        self.object_choice = object_choice

        # [cls, kwargs]
        self.object_params = {
            #'cube': [BoxObject, {'size_min':[0.020, 0.020, 0.020], 'size_max':[0.022, 0.022, 0.022], 'rgba':[0, 0, 255, 1]}],
            'milk': [MilkObject, {}],
            'bread': [BreadObject, {}],
            'cereal': [CerealObject, {}],
            'can': [CanObject, {}],
        }
        assert self.object_choice in self.object_params or self.object_choice == 'random', \
                'object_choice must in {}'.format(list(self.object_params.keys())+['random'])

        super().__init__(
            gripper_type=gripper_type,
            gripper_visualization=gripper_visualization,
            use_indicator_object=use_indicator_object,
            has_renderer=has_renderer,
            has_offscreen_renderer=has_offscreen_renderer,
            render_collision_mesh=render_collision_mesh,
            render_visual_mesh=render_visual_mesh,
            control_freq=control_freq,
            horizon=horizon,
            ignore_done=ignore_done,
            use_camera_obs=use_camera_obs,
            camera_name=camera_name,
            camera_height=camera_height,
            camera_width=camera_width,
            camera_depth=camera_depth,
        )

    def _load_model(self):
        """
        Loads an xml model, puts it in self.model
        """
        super()._load_model()
        self.mujoco_robot.set_base_xpos([0, 0, 0])

        # load model for table top workspace
        if self.arena_type == 'table':
            self.mujoco_arena = TableArena(
                table_full_size=self.table_full_size, table_friction=self.table_friction
            )
            # The sawyer robot has a pedestal, we want to align it with the table
            self.mujoco_arena.set_origin([0.16 + self.table_full_size[0] / 2, 0, 0])
        elif self.arena_type == 'bin':
            self.mujoco_arena = BinsArena(
                table_full_size=self.table_full_size, table_friction=self.table_friction
            )
            # The sawyer robot has a pedestal, we want to align it with the table
            self.mujoco_arena.set_origin([.5, -0.3, 0])

        if self.use_indicator_object:
            self.mujoco_arena.add_pos_indicator()

        object_cls, object_args = self.object_params[self.object_type]
        object = object_cls(**object_args)
        # initialize objects of interest
        self.mujoco_objects = OrderedDict([("cube", object)])

        # task includes arena, robot, and objects of interest
        self.model = TableTopTask(
            self.mujoco_arena,
            self.mujoco_robot,
            self.mujoco_objects,
            initializer=self.placement_initializer,
        )
        self.model.place_objects()

    def _get_reference(self):
        """
        Sets up references to important components. A reference is typically an
        index or a list of indices that point to the corresponding elements
        in a flatten array, which is how MuJoCo stores physical simulation data.
        """
        super()._get_reference()
        self.cube_body_id = self.sim.model.body_name2id("cube")
        self.l_finger_geom_ids = [
            self.sim.model.geom_name2id(x) for x in self.gripper.left_finger_geoms
        ]
        self.r_finger_geom_ids = [
            self.sim.model.geom_name2id(x) for x in self.gripper.right_finger_geoms
        ]
        self.cube_geom_id = self.sim.model.geom_name2id("cube")

    def _reset_internal(self):
        """
        Resets simulation internal configurations.
        """
        if self.object_choice == 'random':
            self.object_type = random.choice(list(self.object_params.keys()))
        else:
            self.object_type = self.object_choice

        super()._reset_internal()

        # reset positions of objects
        self.model.place_objects()

        # reset joint positions
        init_pos = np.array([-0.5538, -0.8208, 0.4155, 1.8409, -0.4955, 0.6482, 1.9628])
        init_pos += np.random.randn(init_pos.shape[0]) * 0.02
        self.sim.data.qpos[self._ref_joint_pos_indexes] = np.array(init_pos)

    def reward(self, action=None):
        """
        Reward function for the task.

        The dense reward has three components.

            Reaching: in [0, 1], to encourage the arm to reach the cube
            Grasping: in {0, 0.25}, non-zero if arm is grasping the cube
            Lifting: in {0, 1}, non-zero if arm has lifted the cube

        The sparse reward only consists of the lifting component.

        Args:
            action (np array): unused for this task

        Returns:
            reward (float): the reward
        """
        reward = 0.

        # sparse completion reward
        if self._check_success():
            reward = 1.0

        # use a shaping reward
        if self.reward_shaping:

            # reaching reward
            cube_pos = self.sim.data.body_xpos[self.cube_body_id]
            gripper_site_pos = self.sim.data.site_xpos[self.eef_site_id]
            dist = np.linalg.norm(gripper_site_pos - cube_pos)
            reaching_reward = 1 - np.tanh(10.0 * dist)
            reward += reaching_reward

            # grasping reward
            touch_left_finger = False
            touch_right_finger = False
            for i in range(self.sim.data.ncon):
                c = self.sim.data.contact[i]
                if c.geom1 in self.l_finger_geom_ids and c.geom2 == self.cube_geom_id:
                    touch_left_finger = True
                if c.geom1 == self.cube_geom_id and c.geom2 in self.l_finger_geom_ids:
                    touch_left_finger = True
                if c.geom1 in self.r_finger_geom_ids and c.geom2 == self.cube_geom_id:
                    touch_right_finger = True
                if c.geom1 == self.cube_geom_id and c.geom2 in self.r_finger_geom_ids:
                    touch_right_finger = True
            if touch_left_finger and touch_right_finger:
                reward += 0.25

        return reward

    def _get_observation(self):
        """
        Returns an OrderedDict containing observations [(name_string, np.array), ...].

        Important keys:
            robot-state: contains robot-centric information.
            object-state: requires @self.use_object_obs to be True.
                contains object-centric information.
            image: requires @self.use_camera_obs to be True.
                contains a rendered frame from the simulation.
            depth: requires @self.use_camera_obs and @self.camera_depth to be True.
                contains a rendered depth map from the simulation
        """
        di = super()._get_observation()
        # camera observations
        if self.use_camera_obs:
            camera_obs = self.sim.render(
                camera_name=self.camera_name,
                width=self.camera_width,
                height=self.camera_height,
                depth=self.camera_depth,
            )
            if self.camera_depth:
                di["image"], di["depth"] = camera_obs
            else:
                di["image"] = camera_obs

        # low-level object information
        if self.use_object_obs:

            # remember the keys to collect into object info
            object_state_keys = []

            # for conversion to relative gripper frame
            gripper_pose = T.pose2mat((di["eef_pos"], di["eef_quat"]))
            world_pose_in_gripper = T.pose_inv(gripper_pose)

            for obj_str in self.object_params:

                if obj_str == self.object_type:
                    obj_pos = np.array(self.sim.data.body_xpos[self.cube_body_id])
                    obj_quat = convert_quat(
                        np.array(self.sim.data.body_xquat[self.cube_body_id]), to="xyzw"
                    )
                    # get relative pose of object in gripper frame
                    object_pose = T.pose2mat((obj_pos, obj_quat))
                    rel_pose = T.pose_in_A_to_pose_in_B(object_pose, world_pose_in_gripper)
                    rel_pos, rel_quat = T.mat2pose(rel_pose)
                else:
                    obj_pos, obj_quat, rel_pos, rel_quat = \
                            np.zeros(3), np.zeros(4), np.zeros(3), np.zeros(4)

                di["{}_pos".format(obj_str)] = obj_pos
                di["{}_quat".format(obj_str)] = obj_quat

                di["{}_to_eef_pos".format(obj_str)] = rel_pos
                di["{}_to_eef_quat".format(obj_str)] = rel_quat

                object_state_keys.append("{}_pos".format(obj_str))
                object_state_keys.append("{}_quat".format(obj_str))
                object_state_keys.append("{}_to_eef_pos".format(obj_str))
                object_state_keys.append("{}_to_eef_quat".format(obj_str))

            di["object-state"] = np.concatenate([di[k] for k in object_state_keys])

        ## original obs
        ## low-level object information
        #if self.use_object_obs:
        #    # position and rotation of object
        #    cube_pos = np.array(self.sim.data.body_xpos[self.cube_body_id])
        #    cube_quat = convert_quat(
        #        np.array(self.sim.data.body_xquat[self.cube_body_id]), to="xyzw"
        #    )
        #    di["cube_pos"] = cube_pos
        #    di["cube_quat"] = cube_quat

        #    gripper_site_pos = np.array(self.sim.data.site_xpos[self.eef_site_id])
        #    di["gripper_to_cube"] = gripper_site_pos - cube_pos

        #    di["object-state"] = np.concatenate(
        #        [cube_pos, cube_quat, di["gripper_to_cube"]]
        #    )

        return di

    def _check_contact(self):
        """
        Returns True if gripper is in contact with an object.
        """
        collision = False
        for contact in self.sim.data.contact[: self.sim.data.ncon]:
            if (
                self.sim.model.geom_id2name(contact.geom1)
                in self.gripper.contact_geoms()
                or self.sim.model.geom_id2name(contact.geom2)
                in self.gripper.contact_geoms()
            ):
                collision = True
                break
        return collision

    def _check_success(self):
        """
        Returns True if task has been completed.
        """
        cube_height = self.sim.data.body_xpos[self.cube_body_id][2]
        table_height = self.table_full_size[2]

        # cube is higher than the table top above a margin
        return cube_height > table_height + 0.04

    def _gripper_visualization(self):
        """
        Do any needed visualization here. Overrides superclass implementations.
        """

        # color the gripper site appropriately based on distance to cube
        if self.gripper_visualization:
            # get distance to cube
            cube_site_id = self.sim.model.site_name2id("cube")
            dist = np.sum(
                np.square(
                    self.sim.data.site_xpos[cube_site_id]
                    - self.sim.data.get_site_xpos("grip_site")
                )
            )

            # set RGBA for the EEF site here
            max_dist = 0.1
            scaled = (1.0 - min(dist / max_dist, 1.)) ** 15
            rgba = np.zeros(4)
            rgba[0] = 1 - scaled
            rgba[1] = scaled
            rgba[3] = 0.5

            self.sim.model.site_rgba[self.eef_site_id] = rgba

#class SawyerLiftBinobsRandom(SawyerLiftBinobs):
#    def __init__(self, **kwargs):
#        kwargs['object_choice'] = 'random'
#        super().__init__(**kwargs)
#
#class SawyerLiftBinBinobsRandom(SawyerLiftBinobs):
#    def __init__(self, **kwargs):
#        kwargs['object_choice'] = 'random'
#        kwargs['arena_type'] = 'bin'
#        kwargs['table_full_size'] = (0.39, 0.49, 0.82)
#        super().__init__(**kwargs)
    


if __name__ == '__main__':
    import ipdb
    import cv2
    import os
    import math
    import pcl
    import robosuite.utils.transform_utils as T


    env = SawyerLiftBinobs(has_renderer=True,
                     camera_depth=True,
                     #camera_name='birdview')
                     #camera_name='frontview')
                     camera_name='agentview')

    rgba_color = {
        'blue':  [0, 0, 3, 1],  # 120,255,255
        'green': [0, 3, 0, 1],  # 60,255,255
        'red':   [3, 0, 0, 1], # 0, 255,255
        'yellow':[3, 3, 0, 1],  # 30,255,255
    }
    hsv_range = {
        'blue':   [[115,150,150],[125,255,255]],  # 120,255,255
        'green':  [[55 ,150,150],[65 ,255,255]],  # 60,255,255
        'red':    [[0  ,150,150],[10 ,255,255]],  # 0, 255,255
        'yellow': [[25 ,150,150],[35 ,255,255]],  # 30,255,255

    }

    color_type = 'blue'
    step = 0
    while True:
        lower, upper = np.array(hsv_range[color_type])
        #env.sim.model.geom_rgba[44,:] = rgba_color[color_type]
        env.render()
        """

        obs = env._get_observation()
        color, depth = obs['image'], obs['depth']
        color =cv2.cvtColor(color, cv2.COLOR_RGB2BGR)
        color = cv2.flip(color, 0) # horizontal flip
        depth = cv2.flip(depth, 0) # horizontal flip
        cv2.imshow('color', color)
        cv2.waitKey(100)
        cv2.imshow('depth', depth)
        cv2.waitKey(100)

        hsv = cv2.cvtColor(color, cv2.COLOR_BGR2HSV)
        mask = cv2.inRange(hsv, lower, upper)
        cv2.imshow('mask', mask)
        cv2.waitKey(100)
        
        #x = np.arange(env.camera_width) - (env.camera_width - 1) / 2
        x = np.arange(env.camera_width) - env.camera_width / 2
        x = np.tile(x, (env.camera_height, 1))
        y = np.arange(env.camera_height)[:, np.newaxis] - env.camera_height / 2
        y = np.tile(y, (1, env.camera_width))
        
        cam_id = env.sim.model.camera_name2id(env.camera_name)
        fovy = env.sim.model.cam_fovy[cam_id]
        cam_pos = env.sim.model.cam_pos[cam_id]
        cam_quat = env.sim.model.cam_quat[cam_id]
        cam_mat = T.pose2mat((cam_pos, T.convert_quat(cam_quat, 'xyzw')))
        # can directly get from 
        # cam_mat = env.sim.model.cam_mat0[cam_id]
        
        f = 0.5 * env.camera_height / math.tan(fovy * math.pi / 360)

        #f *= 0.1
        #depth *= 2.7

        x = x * depth / f
        y = y * depth / f

        obj_index = mask > 0
        obj_x = x[obj_index][np.newaxis,:]
        obj_y = -y[obj_index][np.newaxis,:]
        obj_z = -depth[obj_index][np.newaxis,:]
        obj_points = np.concatenate((obj_x, obj_y, obj_z), axis=0)

        # transform
        obj_points = cam_mat[:3, :3].dot(obj_points) + cam_mat[:3, 3:4]
        obj_points = obj_points.T


        #colors = np.ones((obj_points.shape[0], 1)) * 255
        #points = np.concatenate((obj_points, colors), axis=1).astype('float32')
        #cloud = pcl.PointCloud_PointXYZRGB(points)
        #pcl.save(cloud, '../../model/cube.pcd')


        cube_pos = env.sim.data.get_geom_xpos('cube')
        """

        if step % 1000 == 0:
            env.reset()
            action = np.random.rand(8)
            env.step(action)
        step += 1
        

"""
color:
refer to https://www.cnblogs.com/klchang/p/6784856.html
blue: rgba [0, 0, 255, 1], hsv[[115,150,150],[125,255,255]]   # 120,255,255
green:rbga [0, 255, 0, 1], hsv[[55 ,150,150],[65 ,255,255]]   # 60,255,255
red:  rgba [255, 0, 0, 1], hsv[[0  ,150,150],[10 ,255,255]]   # 0, 255,255
yellow:rgba[255,255,0, 1], hsv[[25 ,150,150],[35 ,255,255]]   # 30,255,255
"""
