import numpy as np
import matplotlib.pyplot as plt
from pyquaternion import Quaternion #for specific task



class PolicyBase:
    def __init__(self, inject_noise=False):
        self.inject_noise = inject_noise
        self.step_count = 0
        self.left_trajectory = None
        self.right_trajectory = None
        # print("init base policy")

    def generate_trajectory(self, obs):
        raise NotImplementedError

    @staticmethod
    def interpolate(curr_waypoint, next_waypoint, t):
        t_frac = (t - curr_waypoint["t"]) / (next_waypoint["t"] - curr_waypoint["t"])
        curr_xyz = curr_waypoint['xyz']
        curr_quat = curr_waypoint['quat']
        curr_grip = curr_waypoint['gripper']
        next_xyz = next_waypoint['xyz']
        next_quat = next_waypoint['quat']
        next_grip = next_waypoint['gripper']
        xyz = curr_xyz + (next_xyz - curr_xyz) * t_frac
        quat = curr_quat + (next_quat - curr_quat) * t_frac
        gripper = curr_grip + (next_grip - curr_grip) * t_frac
        return xyz, quat, gripper

    def predict(self, obs, step_count):
        # generate trajectory at first timestep, then open-loop execution
        self.step_count = step_count
        # print("\nstep_counter===",self.step_count,"\n")
        if self.step_count == 0:
            self.generate_trajectory(obs)
            # print("init scripted policy !")
            action_l = np.array([-0.17297014, 1.00485877, 1.32773627,
                         7.06825181e-01,  8.20281078e-06, -7.07388269e-01, -5.20399313e-06])
            action_r = np.array([0.17297014, 0.9951369,  1.32773623,
                         2.59463975e-06, 7.07388269e-01, 5.59551158e-06, 7.06825181e-01])
            action_gripper = np.array([100,100])
            action = np.hstack((action_l,action_r,action_gripper))
            return action
        # obtain left and right waypoints
        # elif self.step_count >= 1:
        #     action_l = np.array([-0.81444094, 1.0, 1.869898,
        #                  1.0,0.0,0.0,0.0])
        #     action_r = np.array([0.81444094, 1.0, 1.869898,
        #                  1.0,0.0,0.0,0.0])
        #     action_gripper = np.array([100,100])
        #     action = np.hstack((action_l,action_r,action_gripper))
        #     return action
        elif self.step_count >= 1: 
            if self.left_trajectory[0]['t'] == self.step_count:
                self.curr_left_waypoint = self.left_trajectory.pop(0)
            next_left_waypoint = self.left_trajectory[0]

            if self.right_trajectory[0]['t'] == self.step_count:
                self.curr_right_waypoint = self.right_trajectory.pop(0)
            next_right_waypoint = self.right_trajectory[0]

            # Inject noise
            if self.inject_noise:
                scale = 0.01
                left_xyz = left_xyz + np.random.uniform(-scale, scale, left_xyz.shape)
                right_xyz = right_xyz + np.random.uniform(-scale, scale, right_xyz.shape)

            # interpolate between waypoints to obtain current pose and gripper command
            left_xyz, left_quat, left_gripper = self.interpolate(self.curr_left_waypoint, next_left_waypoint, self.step_count)
            right_xyz, right_quat, right_gripper = self.interpolate(self.curr_right_waypoint, next_right_waypoint, self.step_count)
            action_gripper = np.array([left_gripper,right_gripper])
            action_left = np.concatenate([left_xyz, left_quat])
            action_right = np.concatenate([right_xyz, right_quat])
            self.step_count += 1
            return np.hstack([action_left, action_right,action_gripper])

        
        else:
            raise AttributeError
        

class TestPickAndTransferPolicy(PolicyBase):
    def generate_trajectory(self, box_pos):
        init_mocap_pose_left  = np.array([-0.17297014, 1.00485877, 1.32773627,
                         7.06825181e-01,  8.20281078e-06, -7.07388269e-01, -5.20399313e-06])  # 双臂直角
        init_mocap_pose_right = np.array([0.17297014, 0.9951369,  1.32773623,
                         2.59463975e-06, 7.07388269e-01, 5.59551158e-06, 7.06825181e-01])

        box_xyz = box_pos
        # box_quat = box_info[3:]
        # print(f"Generate trajectory for {box_xyz=}")

        gripper_quat = Quaternion(init_mocap_pose_right[3:])
        gripper_approach_quat = gripper_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45)
        gripper_pick_quat = gripper_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=90)

        meet_xyz = np.array([0, 1.0, 1.3277])
        meet_right_quat = gripper_quat
        meet_left_quat = Quaternion(init_mocap_pose_left[3:]) * Quaternion(axis=[0.0, 0.0, 1.0], degrees=-90)

        self.right_trajectory = [
            {"t": 1, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": -100}, # sleep
            {"t": 75, "xyz": np.array([(0.8+box_xyz[0])*0.5,(1.0+box_xyz[1])*0.5,init_mocap_pose_right[2]]), "quat": gripper_approach_quat.elements, "gripper": 100},
            {"t": 225, "xyz": box_xyz + np.array([0, 0, 0.3]), "quat": gripper_pick_quat.elements, "gripper": 100}, # approach the cube
            {"t": 275, "xyz": box_xyz + np.array([0, 0, 0.11]), "quat": gripper_pick_quat.elements, "gripper": 100}, # go down
            {"t": 280, "xyz": box_xyz + np.array([0, 0, 0.11]), "quat": gripper_pick_quat.elements, "gripper": -100}, # close gripper
            {"t": 450, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": -100},# approach wait position
            {"t": 500, "xyz": meet_xyz + np.array([0.1, 0, 0.0]), "quat": meet_right_quat.elements, "gripper": -100},# approach meet position
            {"t": 510, "xyz": meet_xyz + np.array([0.1, 0, 0.0]), "quat": meet_right_quat.elements, "gripper": 100}, # open gripper
            {"t": 675, "xyz": meet_xyz + np.array([0.4, 0, 0.0]), "quat": meet_right_quat.elements, "gripper": 100}, # move to right
            {"t": 700, "xyz": meet_xyz + np.array([0.4, 0, 0.0]), "quat": meet_right_quat.elements, "gripper": 100}, # stay
        ]

        self.left_trajectory = [
            {"t": 1, "xyz": init_mocap_pose_left[:3], "quat": init_mocap_pose_left[3:], "gripper": -100},# sleep
            {"t": 250, "xyz": meet_xyz + np.array([-0.5, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": 100}, # approach meet position
            {"t": 500, "xyz": meet_xyz + np.array([-0.14, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": 100}, # move to meet position
            {"t": 505, "xyz": meet_xyz + np.array([-0.14, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": -100}, # close gripper
            {"t": 675, "xyz": meet_xyz + np.array([-0.3, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": -100}, # move left
            {"t": 700, "xyz": meet_xyz + np.array([-0.3, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": -100}, # stay
        ]
