import numpy as np
from frankapy import FrankaArm
from autolab_core import RigidTransform
import serial
import time

class ControlUtil:
    
    __fa:FrankaArm = None
    __ser: serial.Serial = None
    cnt_rotation = 0
    
    BASE_TO_WORLD:np.ndarray = np.asarray([
        [1, 0, 0, -0.7323],
        [0, 1, 0, 0.05926],
        [0, 0, 1, -0.21],
        [0, 0, 0, 1]
        ])
    
    CAMERA_TO_GRIPPER:np.ndarray = np.asarray([
        [0, -1, 0, 0.01],
        [1, 0, 0, 0],
        [0, 0, 1, 0.075],
        [0, 0, 0, 1]
        ])
    INIT_GRIPPER_POSE:np.ndarray = np.asarray([
        [ 0.44808722 , 0.61103352 , 0.65256787 , 0.36428118],
        [ 0.51676868 , -0.77267257 , 0.36866054,  -0.26519364],
        [ 0.72948524 , 0.17203456 ,-0.66200043 , 0.60938969],
        [ 0.    ,      0.    ,      0.       ,   1.       ]
        ])

    
    @staticmethod
    def connect_robot():
        if ControlUtil.__fa is None:
            ControlUtil.__fa = FrankaArm(robot_num=2)
        if ControlUtil.__ser is None:
            ControlUtil.__ser = serial.Serial(port="/dev/ttyUSB0", baudrate=115200) 
        
    @staticmethod
    def franka_reset() -> None:
        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": 0,
            "right": 1
        }
        ControlUtil.cnt_rotation += degree
        print(f"Table rotated {ControlUtil.cnt_rotation} degree")
        if degree >= 0:
            turn_angle = degree
            turn_direction = turn_directions["right"]
        else:
            turn_angle = -degree
            turn_direction = turn_directions["left"]
        write_len = ControlUtil.__ser.write(f"CT+TRUNSINGLE({turn_direction},{turn_angle});".encode('utf-8'))
        
    
    @staticmethod
    def get_curr_gripper_to_base_pose() -> np.ndarray:
        return ControlUtil.__fa.get_pose().matrix
    
    @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=5, use_impedance=False, ignore_errors=False)
        
    @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:
            return False
        x = abs(new_cam_to_world[0,3])
        y = abs(new_cam_to_world[1,3])
        
        tan_y_x = y/x
        if tan_y_x < np.sqrt(3)/3 or tan_y_x > np.sqrt(3):
            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):
                display_table_rot_z_pose = ControlUtil.get_z_axis_rot_mat(display_table_rot)
                new_cam_to_world = display_table_rot_z_pose @ cam_to_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)
        exec_time = rot_degree/9
        start_time = time.time()
        ControlUtil.rotate_display_table(rot_degree)
        ControlUtil.set_pose(cam_to_world)
        end_time = time.time()
        if end_time - start_time < exec_time:
            time.sleep(exec_time - (end_time - start_time))
            
# ----------- Debug Test -------------

if __name__ == "__main__":
    ControlUtil.connect_robot()
    
    def main_test():
        print(ControlUtil.get_curr_gripper_to_base_pose())
        ControlUtil.init()
    
    def rotate_back(rotation):
        ControlUtil.rotate_display_table(-rotation)
    
    rotate_back(122.0661478599)