import numpy as np
from frankapy import FrankaArm
from autolab_core import RigidTransform
from scipy.spatial.transform import Rotation as R
import serial
import time

class ControlUtil:
    
    __fa:FrankaArm = None
    __ser: serial.Serial = None
    curr_rotation = 0
    
    # BASE_TO_WORLD:np.ndarray = np.asarray([
    #     [1, 0, 0, -0.68834619],
    #     [0, 1, 0, -0.01785319],
    #     [0, 0, 1, -0.11924244],
    #     [0, 0, 0, 1]
    #     ])
    
    # BASELINE_DIS = 0.1
    # LEFT_TO_RIGHT_ANGLE = -25.0
    # CAMERA_TO_LEFT_CAMERA:np.ndarray = np.asarray([
    #     [1, 0, 0, BASELINE_DIS/2 * np.cos(np.radians(-LEFT_TO_RIGHT_ANGLE/2))],
    #     [0, 1, 0, 0],
    #     [0, 0, 1, BASELINE_DIS/2 * np.sin(np.radians(-LEFT_TO_RIGHT_ANGLE/2))],
    #     [0, 0, 0, 1]
    #     ])
    # CAMERA_TO_LEFT_CAMERA[:3, :3] = R.from_euler('y', LEFT_TO_RIGHT_ANGLE/2, degrees=True).as_matrix()
    
    CAMERA_CORRECTION_ANGLE = -13
    CAMERA_CORRECTION:np.ndarray = np.eye(4)
    CAMERA_CORRECTION[:3, :3] = R.from_euler('y', CAMERA_CORRECTION_ANGLE, degrees=True).as_matrix()
    
    BASE_TO_WORLD:np.ndarray = np.asarray([
        [1, 0, 0, -0.49602172],
        [0, 1, 0, 0.06688724],
        [0, 0, 1, -0.2488704 ],
        [0, 0, 0, 1]
        ])
    

    CAMERA_TO_GRIPPER:np.ndarray = np.asarray(
            [[-0.00757086, -0.99984193,  0.01608569,  0.02644546],
            [ 0.99993362, -0.00770935, -0.00856502, -0.04860572],
            [ 0.00868767,  0.01601977,  0.99983391, -0.02169477],
            [ 0.,          0.   ,       0.     ,     1.      ,  ]
        ])

    INIT_GRIPPER_POSE:np.ndarray = np.asarray(
        [[0.41869126  ,0.87596275 , 0.23951774 , 0.36005292],
        [ 0.70787907 ,-0.4800251  , 0.51813998 ,-0.26499909],
        [ 0.56884584, -0.04739109 ,-0.82107382  ,0.66881103],
        [ 0.         , 0.    ,      0.      ,    1.        ]])
    
    # INIT_GRIPPER_POSE:np.ndarray = np.asarray(
    # [[1  ,0 , 0 , 0.5],
    # [ 0 ,-1  , 0 ,0],
    # [ 0, 0 ,-1  ,0.6],
    # [ 0.         , 0.    ,      0.      ,    1.        ]])
    
    INIT_JOINTS:np.ndarray = np.asarray([-1.64987928,-0.73557812,-0.02676473,-2.08365335,0.01794927,1.5182423,0.73276058])
    
    @staticmethod
    def connect_robot():
        if ControlUtil.__fa is None:
            ControlUtil.__fa = FrankaArm(robot_num=1)
        if ControlUtil.__ser is None:
            ControlUtil.__ser = serial.Serial(port="/dev/ttyUSB0", baudrate=115200) 
        
    @staticmethod
    def franka_reset() -> None:
        #ControlUtil.__fa.goto_joints(ControlUtil.INIT_JOINTS, duration=5, use_impedance=False, ignore_errors=False)
        ControlUtil.__fa.reset_joints()
        
    @staticmethod
    def init() -> None:
        ControlUtil.franka_reset()
        ControlUtil.set_gripper_pose(ControlUtil.INIT_GRIPPER_POSE)
        
    @staticmethod
    def get_pose() -> np.ndarray:
        gripper_to_base = ControlUtil.get_curr_gripper_to_base_pose()
        cam_to_world = ControlUtil.BASE_TO_WORLD @ gripper_to_base @ ControlUtil.CAMERA_TO_GRIPPER
        return cam_to_world
    
    @staticmethod
    def set_pose(cam_to_world: np.ndarray) -> None:
        gripper_to_base = ControlUtil.solve_gripper_to_base(cam_to_world)
        ControlUtil.set_gripper_pose(gripper_to_base)
        
    @staticmethod
    def rotate_display_table(degree):
        turn_directions = {
            "left": 1,
            "right": 0
        }
        delta_degree = degree - ControlUtil.curr_rotation
        ControlUtil.curr_rotation = degree
        print(f"Table rotated {ControlUtil.curr_rotation} degree")
        
        if delta_degree >= 0:
            turn_angle = delta_degree
            turn_direction = turn_directions["left"]
        else:
            turn_angle = -delta_degree
            turn_direction = turn_directions["right"]
        write_len = ControlUtil.__ser.write(f"CT+TRUNSINGLE({turn_direction},{turn_angle});".encode('utf-8'))
        print(f"send command: CT+TRUNSINGLE({turn_direction},{turn_angle});")
        
        return delta_degree
        
    
    @staticmethod
    def get_curr_gripper_to_base_pose() -> np.ndarray:
        return ControlUtil.__fa.get_pose().matrix
    
    @staticmethod
    def get_curr_joints() -> np.ndarray:
        return ControlUtil.__fa.get_joints()
    
    @staticmethod
    def set_gripper_pose(gripper_to_base: np.ndarray) -> None:
        gripper_to_base = RigidTransform(rotation=gripper_to_base[:3, :3], translation=gripper_to_base[:3, 3], from_frame="franka_tool", to_frame="world")
        ControlUtil.__fa.goto_pose(gripper_to_base, duration=6, use_impedance=False, ignore_errors=False, ignore_virtual_walls=True)
        
    @staticmethod
    def solve_gripper_to_base(cam_to_world: np.ndarray) -> np.ndarray:
        return np.linalg.inv(ControlUtil.BASE_TO_WORLD) @ cam_to_world @ np.linalg.inv(ControlUtil.CAMERA_TO_GRIPPER)
    
    @staticmethod
    def sovle_cam_to_world(gripper_to_base: np.ndarray) -> np.ndarray:
        return ControlUtil.BASE_TO_WORLD @ gripper_to_base @ ControlUtil.CAMERA_TO_GRIPPER
    
    @staticmethod
    def check_limit(new_cam_to_world):
        if new_cam_to_world[0,3] > 0 or new_cam_to_world[1,3] > 0:
        # if new_cam_to_world[0,3] > 0:
            return False
        x = abs(new_cam_to_world[0,3])
        y = abs(new_cam_to_world[1,3])
        
        tan_y_x = y/x
        min_angle = 20 / 180 * np.pi
        max_angle = 50 / 180 * np.pi
        if tan_y_x < np.tan(min_angle) or tan_y_x > np.tan(max_angle):
            return False
         
        return True
    
    @staticmethod
    def solve_display_table_rot_and_cam_to_world(cam_to_world: np.ndarray) -> tuple:   
        if ControlUtil.check_limit(cam_to_world):
            return 0, cam_to_world
        else:
            min_display_table_rot = 180
            min_new_cam_to_world = None
            for display_table_rot in np.linspace(0.1,360, 1800):
                new_world_to_world = ControlUtil.get_z_axis_rot_mat(display_table_rot)
                new_cam_to_new_world = cam_to_world
                new_cam_to_world = new_world_to_world @ new_cam_to_new_world
                
                if ControlUtil.check_limit(new_cam_to_world):
                    if display_table_rot < min_display_table_rot:
                        min_display_table_rot, min_new_cam_to_world = display_table_rot, new_cam_to_world
                    if abs(display_table_rot - 360) < min_display_table_rot:
                        min_display_table_rot, min_new_cam_to_world = display_table_rot - 360, new_cam_to_world
        
        if min_new_cam_to_world is None:
            raise ValueError("No valid display table rotation found")
        
        return min_display_table_rot, min_new_cam_to_world
                
    @staticmethod
    def get_z_axis_rot_mat(degree):
        radian = np.radians(degree)
        return np.array([
            [np.cos(radian), -np.sin(radian), 0, 0],
            [np.sin(radian), np.cos(radian), 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
            ])
    
    @staticmethod
    def get_gripper_to_base_axis_angle(gripper_to_base: np.ndarray) -> bool:
        rot_mat = gripper_to_base[:3,:3]
        gripper_z_axis = rot_mat[:,2]
        base_x_axis = np.array([1,0,0])
        angle = np.arccos(np.dot(gripper_z_axis, base_x_axis))
        return angle
    
    @staticmethod
    def move_to(pose: np.ndarray):
        rot_degree, cam_to_world = ControlUtil.solve_display_table_rot_and_cam_to_world(pose)
        
        delta_degree = ControlUtil.rotate_display_table(rot_degree)
        exec_time = abs(delta_degree)/9
        start_time = time.time()
        ControlUtil.set_pose(cam_to_world)
        end_time = time.time()
        print(f"Move to pose with rotation {rot_degree} degree, exec time: {end_time - start_time}|exec time: {exec_time}")
        if end_time - start_time < exec_time:
            time.sleep(exec_time - (end_time - start_time))
            
    @staticmethod
    def absolute_rotate_display_table(degree):
        exec_time = abs(degree)/9
        write_len = ControlUtil.__ser.write(f"CT+TRUNSINGLE({0},{degree});".encode('utf-8'))
        print(f"send command: CT+TRUNSINGLE({0},{degree});")
        time.sleep(exec_time)
            
# ----------- Debug Test -------------

if __name__ == "__main__":
    ControlUtil.connect_robot()
    #ControlUtil.franka_reset()
    print(ControlUtil.get_curr_gripper_to_base_pose())
    # ControlUtil.rotate_display_table(180)
    # ControlUtil.init()
    # ControlUtil.franka_reset()
    
    # def main_test():
    #     print(ControlUtil.get_curr_gripper_to_base_pose())
    #     ControlUtil.init()
    
    # def rotate_back(rotation):
    #     ControlUtil.rotate_display_table(-rotation)
    
    # #main_test()
    # import sys; sys.path.append("/home/yan20/nbv_rec/project/franka_control")
    # from utils.communicate_util import CommunicateUtil
    # import ipdb
    # ControlUtil.init()
    # view_data_0 = CommunicateUtil.get_view_data(init=True)
    # depth_extrinsics_0 = view_data_0["depth_extrinsics"]
    # cam_to_world_0 = ControlUtil.get_pose()
    # print("cam_extrinsics_0")
    # print(depth_extrinsics_0)
    # print("cam_to_world_0")
    # print(cam_to_world_0)
    
    # ipdb.set_trace()
    # TEST_POSE:np.ndarray = np.asarray([
    #     [ 0.46532393,  0.62171798,  0.63002284,  0.30230963],
    #     [ 0.43205618, -0.78075723,  0.45136491, -0.29127173],
    #     [ 0.77251656,  0.06217437, -0.63193429,  0.559957  ],
    #     [ 0.        ,  0.        ,  0.        ,  1.        ],
    #     ])
    # TEST_POSE_CAM_TO_WORLD = ControlUtil.BASE_TO_WORLD @ TEST_POSE @ ControlUtil.CAMERA_TO_GRIPPER
    # ControlUtil.move_to(TEST_POSE_CAM_TO_WORLD)
    # view_data_1 = CommunicateUtil.get_view_data()
    # depth_extrinsics_1 = view_data_1["depth_extrinsics"]
    # depth_extrinsics_1[:3,3] = depth_extrinsics_1[:3,3] / 1000
    # cam_to_world_1 = ControlUtil.get_pose()
    # print("cam_extrinsics_1")
    # print(depth_extrinsics_1)
    # print("cam_to_world_1")
    # print(TEST_POSE_CAM_TO_WORLD)
    # actual_cam_to_world_1 = cam_to_world_0 @ depth_extrinsics_1
    # print("actual_cam_to_world_1")
    # print(actual_cam_to_world_1)
    # ipdb.set_trace()
    # TEST_POSE_2:np.ndarray = np.asarray(
    #     [[ 0.74398544, -0.61922696,  0.251049,    0.47000935],
    #     [-0.47287207, -0.75338888, -0.45692666,  0.20843903],
    #     [ 0.47207883 , 0.22123272, -0.85334192,  0.57863381],
    #     [ 0. ,         0.   ,       0.  ,        1.  ,      ]]
    #     )
    # TEST_POSE_CAM_TO_WORLD_2 = ControlUtil.BASE_TO_WORLD @ TEST_POSE_2 @ ControlUtil.CAMERA_TO_GRIPPER
    
    # #ControlUtil.move_to(TEST_POSE_CAM_TO_WORLD_2)
    # ControlUtil.set_pose(TEST_POSE_CAM_TO_WORLD_2)
    # view_data_2 = CommunicateUtil.get_view_data()
    # depth_extrinsics_2 = view_data_2["depth_extrinsics"]
    # depth_extrinsics_2[:3,3] = depth_extrinsics_2[:3,3] / 1000
    # cam_to_world_2 = ControlUtil.get_pose()
    # print("cam_extrinsics_2")
    # print(depth_extrinsics_2)
    # print("cam_to_world_2")
    # print(TEST_POSE_CAM_TO_WORLD_2)
    # actual_cam_to_world_2 = cam_to_world_0 @ depth_extrinsics_2
    # print("actual_cam_to_world_2")
    # print(actual_cam_to_world_2)
    # ipdb.set_trace()