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

class ControlUtil:
    
    __fa = FrankaArm(robot_num=2)
    
    BASE_TO_DISPLAYTBLE:np.ndarray = np.asarray([
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [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]
        ])
    
    @staticmethod
    def get_franka_arm() -> FrankaArm:
        return ControlUtil.__fa
    
    @staticmethod
    def get_pose() -> np.ndarray:
        gripper_to_base = ControlUtil.__fa.get_pose().matrix
        cam_to_world = ControlUtil.BASE_TO_DISPLAYTBLE @ gripper_to_base @ ControlUtil.CAMERA_TO_GRIPPER
        return cam_to_world
    
    @staticmethod
    def set_pose(cam_to_world: np.ndarray) -> None:
        gripper_to_base = np.linalg.inv(ControlUtil.BASE_TO_DISPLAYTBLE) @ cam_to_world @ np.linalg.inv(ControlUtil.CAMERA_TO_GRIPPER)
        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 reset() -> None:
        ControlUtil.__fa.reset_joints()
        
        
# ----------- Debug Test -------------

if __name__ == "__main__":
    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())