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

class ControlUtil:
    
    #__fa = FrankaArm(robot_num=2)
    
    BASE_TO_WORLD:np.ndarray = np.asarray([
        [1, 0, 0, -0.5],
        [0, 1, 0, 0],
        [0, 0, 1, -0.2],
        [0, 0, 0, 1]
        ])
    
    CAMERA_TO_GRIPPER:np.ndarray = np.asarray([
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]
        ])
    theta = np.radians(25)
    INIT_POSE:np.ndarray = np.asarray([
        [np.cos(theta), 0, -np.sin(theta), 0],
        [0, -1, 0, 0],
        [-np.sin(theta), 0, -np.cos(theta), 0.35],
        [0, 0, 0, 1]
        ])
    
    AXIS_THRESHOLD = (-(np.pi+5e-2)/2, (np.pi+5e-2)/2)
    
    @staticmethod
    def franka_reset() -> None:
        ControlUtil.__fa.reset_joints()
        
    @staticmethod
    def init() -> None:
        ControlUtil.set_pose(ControlUtil.INIT_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)
        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, use_impedance=False, ignore_errors=False)
        
    @staticmethod
    def rotate_display_table(degree):
        pass
    
    @staticmethod
    def get_curr_gripper_to_base_pose() -> np.ndarray:
        return ControlUtil.__fa.get_pose().matrix
    
    @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 solve_display_table_rot_and_cam_to_world(cam_to_world: np.ndarray) -> tuple:
        gripper_to_base = ControlUtil.solve_gripper_to_base(cam_to_world)
        gripper_to_base_axis_angle = ControlUtil.get_gripper_to_base_axis_angle(gripper_to_base)
        
        if ControlUtil.AXIS_THRESHOLD[0] <= gripper_to_base_axis_angle <= ControlUtil.AXIS_THRESHOLD[1]:
            return 0, cam_to_world
        else:
            for display_table_rot in np.linspace(0.1,180, 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.AXIS_THRESHOLD[0] <= ControlUtil.get_gripper_to_base_axis_angle(new_cam_to_world) <= ControlUtil.AXIS_THRESHOLD[1]:
                    return -display_table_rot, new_cam_to_world
                
                display_table_rot = -display_table_rot
                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.AXIS_THRESHOLD[0] <= ControlUtil.get_gripper_to_base_axis_angle(new_cam_to_world) <= ControlUtil.AXIS_THRESHOLD[1]:
                    return -display_table_rot, 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)
        print("table rot degree:", rot_degree)
        ControlUtil.rotate_display_table(rot_degree)
        ControlUtil.set_pose(cam_to_world)
        
        
# ----------- Debug Test -------------

if __name__ == "__main__":
    ControlUtil.init()
    import time
    start = time.time()
    rot_degree, cam_to_world = ControlUtil.solve_display_table_rot_and_cam_to_world(ControlUtil.INIT_POSE)
    end = time.time()
    print(f"Time: {end-start}")
    print(rot_degree, cam_to_world)
    # test_pose = np.asarray([
    #      [1, 0, 0, 0.4],
    #     [0, -1, 0, 0],
    #     [0, 0, -1, 0.6],
    #     [0, 0, 0, 1]
    #      ])
    # ControlUtil.set_pose(test_pose)
    # print(ControlUtil.get_pose())
    # ControlUtil.reset()
    # print(ControlUtil.get_pose())
    
    angle = ControlUtil.get_gripper_to_base_axis_angle(ControlUtil.solve_gripper_to_base(cam_to_world))
    threshold = ControlUtil.AXIS_THRESHOLD
    
    angle_degree = np.degrees(angle)
    threshold_degree = np.degrees(threshold[0]), np.degrees(threshold[1])
    print(f"Angle: {angle_degree}, range: {threshold_degree}")
    ControlUtil.set_pose(cam_to_world)