from robodk.robolink import *   # API to communicate with RoboDK
from robodk.robomath import *   # Basic matrix operations
import os
from time import sleep
import pathlib

class Mitsubishi_Arm:
    def __init__(self, ip_address, port, name, rdk_file_path = pathlib.Path(__file__).parent / "Mitsubishi-RV-2FR.robot"):
        self.ip_address = ip_address
        self.port = port
        self.name = name
        self.rdk_file_path = rdk_file_path
        
    def setArmSpeed(self, speed):
        self.robot.setSpeed(speed_joints=speed, speed_linear=speed, accel_joints=speed, accel_linear=speed)
        self.robot.setSpeedJoints(speed)
        
    def connect(self):
        self.RDK = Robolink(args=["-NOUI", "-NEWINSTANCE"])
        self.RDK.AddFile(str(self.rdk_file_path))
        self.robot = self.RDK.Item(self.name, ITEM_TYPE_ROBOT)
        self.robot.setParam("Driver", "SRVON")
        self.robot.setRounding(1)
        
        if not self.robot.Valid():
            self.RDK.CloseRoboDK()
            raise Exception("No robot selected or available")
        
        self.robot.setConnectionParams(self.ip_address, self.port, "/", "", "")
        self.robot.ConnectSafe()
        sleep(3) 

        status, status_msg = self.robot.ConnectedState()
        if status != ROBOTCOM_READY:
            self.robot.Disconnect()
            self.RDK.CloseRoboDK()
            raise Exception("Failed to connect: " + status_msg)
        
        # Set Run Mode
        self.RDK.setRunMode(RUNMODE_RUN_ROBOT)

        # Get Robot Current Joints
        self.robot.Joints()
        self.robot.ConnectSafe()

        # Set Arm Speed... This only works for moveL NOT moveJ
        self.robot.setSpeed(speed_joints=1000, speed_linear=1000, accel_joints=1000, accel_linear=1000)                         
            
    def disconnect(self):
        sleep(5)
        try:
            self.robot.Disconnect()
            self.RDK.CloseRoboDK()
        except:
            pass
        os.system(f"taskkill /f /im robodk.exe")
        sleep(1)

    def __exit__(self):
        try:
            self.disconnect()
            print("Exited")
        except:
            print("Exit Not Done")

    def __del__(self):
        try:
            self.disconnect()
            print("Deleted")
        except:
            print("Delete Not Done")
        
        
if __name__ == "__main__":
    # mitsubishi Home
    home_pos_vec = [   270.000000,     0.000000,   505.000000,   180.000000,     0.000000,  -90.000000 ]
    home_pose = Fanuc_2_Pose(home_pos_vec)

    # Initialize mitsubishi Arm
    mitsubishi_arm = Mitsubishi_Arm('10.108.40.56', 10001, 'Mitsubishi RV-2FR', pathlib.Path(__file__).parent / "Mitsubishi.rdk")
    
    mitsubishi_arm.connect()

    movements = [(60,0,0),
                 (0,0,60),
                 (-120,0,0),
                 (0,0,-60),
                 (60,0,0)]
    
    def move(pose):
        for _ in range(1):
            print("Need to try another move")
            mitsubishi_arm.robot.MoveL(pose)
            print("That move is done")
    print("-----------Move Started Home------------")
    move(home_pose)
    print("-----------Move Done------------\n\n")
    for _ in range(1):
    
        print("-----------Move Started Left------------")
        move(home_pose*transl(60,0,0))
        print("-----------Move Done------------\n\n")

        print("-----------Move Started Down------------")
        move(home_pose*transl(0,0,60))
        print("-----------Move Done------------\n\n")

        print("-----------Move Started Right------------")
        move(home_pose*transl(-60,0,0))
        print("-----------Move Done------------\n\n")

        print("-----------Move Started Home------------")
        move(home_pose)
        print("-----------Move Done------------\n\n")

    ############################

    # Close RoboDK
    mitsubishi_arm.disconnect()