# Copyright 2015-2023 - RoboDK Inc. - https://robodk.com/
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
#
#
#
# This is a Python module that allows driving a Kinova robot.
# This Python module can be run directly in console mode to test its functionality.
# This module allows communicating with a robot through the command line.
# The same commands we can input manually are used by RoboDK to drive the robot from the PC.
# RoboDK Drivers are located in /RoboDK/api/Robot/ by default. Drivers can be PY files or EXE files.
#
# Drivers are modular. They are not part of the RoboDK executable but they must be placed in C:/RoboDK/api/robot/, then, linked in the Connection parameters menu:
#   1. right click a robot in RoboDK, then, select "Connect to robot".
#   2. In the "More options" menu it is possible to update the location and name of the driver.
# Driver linking is automatic for currently available drivers.
# More information about robot drivers available here:
#   https://robodk.com/doc/en/Robot-Drivers.html#RobotDrivers
#
# Alternatively to the standard programming methods (where a program is generated, then, transferred to the robot and executed) it is possible to run a program simulation directly on the robot
# The robot movement in the simulator is then synchronized with the real robot.
# Programs generated from RoboDK can be run on the robot by right clicking the program, then selecting "Run on robot".
#   Example:
#   https://www.youtube.com/watch?v=pCD--kokh4s
#
# Example of an online programming project:
#   https://robodk.com/blog/online-programming/
#
# It is possible to control the movement of a robot from the RoboDK API (for example, from a Python or C# program using the RoboDK API).
# The same code is used to simulate and optionally move the real robot.
#   Example:
#   https://robodk.com/offline-programming
#
#   To establish connection from RoboDK API:
#   https://robodk.com/doc/en/PythonAPI/robolink.html#robolink.Item.ConnectSafe
#
# Example of a quick manual test in console mode:
#  User entry: CONNECT 192.168.123.1 7000
#  Response:   SMS:Response from the robot or failure to connect
#  Response:   SMS:Ready 
#  User entry: MOVJ 10 20 30 40 50 60 70
#  Response:   SMS:Working...
#  Response:   SMS:Ready
#  User entry: CJNT
#  Response:   SMS:Working...
#  Response:   JNTS: 10 20 30 40 50 60 70
#
# ---------------------------------------------------------------------------------
import math
import socket
import struct
import sys
import re
from io import BytesIO
import argparse

from robolink import import_install
import_install("mirobot", "mirobot-py")
from mirobot import Mirobot
import time

#import robolink
#robolink.import_install("serial")
#robolink.import_install("kortex_api", os.environ['PYTHONPATH'] + "/kortex_api-2.2.0.post31-py3-none-any.whl")

# ---------------------------------------------------------------------------------
# Set the minimum number of degrees of freedom that are expected
nDOFs_MIN = 6

# Set the driver version
DRIVER_VERSION = "RoboDK Driver for Mirobot v1.0.1"

# ---------------------------------------------------------------------------------


# This class handles communication between this driver (PC) and the robot
class ComRobot:
    """Robot class for programming Kinova robots"""    
    LAST_MSG = None  # Keep a copy of the last message received
    CONNECTED = False  # Connection status is known at all times
    MIROBOTAPI = None
    BUFFER_SIZE = None
    #TIMEOUT = None
    #Speed and accelerations
    LINEARSPEED = 100
    LINEARACEL = 30
    JOINTSPEED  = 100
    JOINTACEL  = 80
    LAST_TARGET_JOINTS = []


    # This is executed when the object is created
    def __init__(self):
        self.BUFFER_SIZE = 512  # bytes
        #self.TIMEOUT = 5 * 60  # seconds: it must be enough time for a movement to complete
        # self.TIMEOUT = 10 # seconds
        self.CONNECTED = False

    def __del__(self):
        self.disconnect()

    # Disconnect from robot
    def disconnect(self):
        self.CONNECTED = False
        return True

    # Connect to robot
    def connect(self, ip, port=-1):
        if (self.CONNECTED == True):
            return
        self.disconnect()
        print_message('Connecting to robot %s:%i' % (ip, port))
        
        # Create new socket or serial connection
        if ip == "" or ip == "127.0.0.1":
            self.MIROBOTAPI = Mirobot()
            self.MIROBOTAPI.autoConnect()
        elif ip.startswith("COM"):
            self.MIROBOTAPI = Mirobot(portname=ip)
        else:
            self.MIROBOTAPI = Mirobot(ip)

        self.MIROBOTAPI.home_simultaneous()
        UpdateStatus(ROBOTCOM_WORKING)


        self.CONNECTED = True
        sys.stdout.flush()
        return True

    def MoveJ(self,coordArray):
        # Joint move to the provided angles
        import threading
        def runMove():
            self.MIROBOTAPI.go_to_axis(coordArray[0],coordArray[1],coordArray[2],coordArray[3],coordArray[4],coordArray[5])
        t = threading.Thread(target = runMove)
        t.start()
        while t.is_alive():
            joints = [None] * 7
            joints[0] = self.MIROBOTAPI.status.angle.a
            joints[1] = self.MIROBOTAPI.status.angle.b
            joints[2] = self.MIROBOTAPI.status.angle.c
            joints[3] = self.MIROBOTAPI.status.angle.x
            joints[4] = self.MIROBOTAPI.status.angle.y
            joints[5] = self.MIROBOTAPI.status.angle.z
            joints[6] = self.MIROBOTAPI.status.angle.d
            print_joints(joints,True)
            time.sleep(1.0/1000.0)

        joints = [None] * 7
        joints[0] = self.MIROBOTAPI.status.angle.a
        joints[1] = self.MIROBOTAPI.status.angle.b
        joints[2] = self.MIROBOTAPI.status.angle.c
        joints[3] = self.MIROBOTAPI.status.angle.x
        joints[4] = self.MIROBOTAPI.status.angle.y
        joints[5] = self.MIROBOTAPI.status.angle.z
        joints[6] = self.MIROBOTAPI.status.angle.d
        print_joints(joints)
        return True

    def MoveL(self,xyzabc):
        import threading
        """Linear move (joint values provided in degrees)"""
        # Create new action
        #print(self.MIROBOTAPI.status)
        self.MoveJ(xyzabc)
        return True
    
    def MoveC(self,joints):
        raise Exception("Circular move not implemented")
        #try:
        #    self.UARMAPI.set_mode(mode=1) #Mode 1 corresponds to moving
        #    self.UARMAPI.motion_enable(True)
        #    self.UARMAPI.move_circle(pose1=joints[0:6], pose2=joints[7:], percent=50, speed=self.JOINTSPEED, mvacc=self.JOINTACELL, #wait=True)
            #self.LAST_TARGET_JOINTS = joints
        #except Exception as e:
        #    print_message(str(e))
        #    return False
        #return True

    def GripperPosition(self, gripperPosition):
        # Open the gripper the indicated amount
        # Passed:  Integer position between 0 (open) and 1 (closed)

        # Create the GripperCommand we will send
        gripper_command = Base_pb2.GripperCommand()
        finger = gripper_command.gripper.finger.add()

        # Move the gripper
        gripper_command.mode = Base_pb2.GRIPPER_POSITION
        finger.value = gripperPosition
        finger.finger_identifier = 1

        print("Actuating gripper to {:0.2f}...".format(finger.value))
        self.base.SendGripperCommand(gripper_command)

        # Wait for reported position to be reached or speed reaches zero
        gripper_request1 = Base_pb2.GripperRequest()
        gripper_request1.mode = Base_pb2.GRIPPER_POSITION
        gripper_request2 = Base_pb2.GripperRequest()
        gripper_request2.mode = Base_pb2.GRIPPER_SPEED
        while True:
            gripper_measure = self.base.GetMeasuredGripperMovement(gripper_request1)
            gripper_speed = self.base.GetMeasuredGripperMovement(gripper_request2)
            if len(gripper_measure.finger):
                print("Current position is : {0}".format(gripper_measure.finger[0].value))
                if abs(gripper_measure.finger[0].value - gripperPosition) < 0.01:
                    break
                if gripper_speed == 0.0:
                    break
            else: # Else, no finger present in answer, end loop
                break


    def getJoints(self):
        joints = [None] * 7
        joints[0] = self.MIROBOTAPI.status.angle.a
        joints[1] = self.MIROBOTAPI.status.angle.b
        joints[2] = self.MIROBOTAPI.status.angle.c
        joints[3] = self.MIROBOTAPI.status.angle.x
        joints[4] = self.MIROBOTAPI.status.angle.y
        joints[5] = self.MIROBOTAPI.status.angle.z
        joints[6] = self.MIROBOTAPI.status.angle.d
        return joints

    def setSpeed(self, speed_values):
        # speed_values[0] = speed_values[0] # linear speed in mm/s
        # speed_values[1] = speed_values[1] # joint speed in mm/s
        # speed_values[2] = speed_values[2] # linear acceleration in mm/s2
        # speed_values[3] = speed_values[3] # joint acceleration in deg/s2
        if (speed_values[0] != -1):
            self.LINEARSPEED = speed_values[0]

        if (speed_values[1] != -1):
            self.JOINTSPEED = speed_values[1]

        if (speed_values[2] != -1):
            self.LINEARACEL = speed_values[2]

        if (speed_values[3] != -1):
            self.JOINTACEL = speed_values[3]

        return True

    def setTool(self,tool_pose):
        #self.UARMAPI.set_tcp_offset(tool_pose)
        return True

    def Pause(self,timeMS):
        import time
        time.sleep(timeMS/1000)
        return True

    def setRounding(self,rounding):
        return True

    def setDO(self,digital_IO_State):
        return False

    def WaitDI(self,digital_IO_Num):
        return False
        
        #import time
        #start = time.time()
        #ioNumber = digital_IO_Num[0]
        #ioState = self.UARMAPI.get_tgpio_digital(ioNumber)
        #desiredState = digital_IO_Num[1]
        #try:
        #    timeout = digital_IO_Num[2]
        #except Exception as e:
        #    e = e
        #    timeout = 0
#
#        while not (ioState == desiredState) and (time.time() - start) < timeout:
#            ioState = self.UARMAPI.get_tgpio_digital(ioNumber)            
#            time.sleep(0.1)
#        return False




# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------
# Generic RoboDK driver for a specific Robot class
ROBOT = ComRobot()
ROBOT_IP = "172.31.1.147"  # IP of the robot
ROBOT_PORT = 30000  # Communication port of the robot



def Robot_Disconnect():
    global ROBOT
    ROBOT.disconnect()


# ------------ robot connection -----------------
# Establish connection with the robot
def RobotConnect():
    global ROBOT
    global ROBOT_IP
    global ROBOT_PORT
    return ROBOT.connect(ROBOT_IP, ROBOT_PORT)


# Disconnect from the robot
def RobotDisconnect():
    global ROBOT
    ROBOT.disconnect()
    return True


# -----------------------------------------------------------------------------
# Generic RoboDK driver tools

# Note, a simple print() will flush information to the log window of the robot connection in RoboDK
# Sending a print() might not flush the standard output unless the buffer reaches a certain size

def print_message(message):
    """print_message will display a message in the log window (and the connexion status bar)"""
    print("SMS:" + message)
    sys.stdout.flush()  # very useful to update RoboDK as fast as possible


def show_message(message):
    """show_message will display a message in the status bar of the main window"""
    print("SMS2:" + message)
    sys.stdout.flush()  # very useful to update RoboDK as fast as possible


def print_joints(joints, is_moving=False):
    # if len(joints) > 6:
    #    joints = joints[0:6]
    if is_moving:
        # Display the feedback of the joints when the robot is moving
        print("JNTS_MOVING " + " ".join(format(x, ".5f") for x in joints))  # if joints is a list of float
        # print("JNTS_MOVING " + joints)
    else:
        print("JNTS " + " ".join(format(x, ".5f") for x in joints))  # if joints is a list of float
        # print("JNTS " + joints)
    sys.stdout.flush()  # very useful to update RoboDK as fast as possible


# ---------------------------------------------------------------------------------
# Constant values to display status using UpdateStatus()
ROBOTCOM_UNKNOWN = -1000
ROBOTCOM_CONNECTION_PROBLEMS = -3
ROBOTCOM_DISCONNECTED = -2
ROBOTCOM_NOT_CONNECTED = -1
ROBOTCOM_READY = 0
ROBOTCOM_WORKING = 1
ROBOTCOM_WAITING = 2

# Last robot status is saved
STATUS = ROBOTCOM_DISCONNECTED


# UpdateStatus will send an appropriate message to RoboDK which will result in a specific coloring
# for example, Ready will be displayed in green, Waiting... will be displayed in Yellow and other messages
# will be displayed in red
def UpdateStatus(set_status=None):
    global STATUS
    if set_status is not None:
        STATUS = set_status

    if STATUS == ROBOTCOM_CONNECTION_PROBLEMS:
        print_message("Connection problems")
    elif STATUS == ROBOTCOM_DISCONNECTED:
        print_message("Disconnected")
    elif STATUS == ROBOTCOM_NOT_CONNECTED:
        print_message("Not connected")
    elif STATUS == ROBOTCOM_READY:
        print_message("Ready")
    elif STATUS == ROBOTCOM_WORKING:
        print_message("Working...")
    elif STATUS == ROBOTCOM_WAITING:
        print_message("Waiting...")
    else:
        print_message("Unknown status")


# Sample set of commands that can be provided by RoboDK of through the command line
def TestDriver():
    # try:
    # rob_ip = input("Enter the robot IP: ")
    # rob_port = input("Enter the robot Port (default=1101): ")
    # rob_port = int(rob_port)

    # RunCommand("CONNECT 192.168.0.100 10000")
    RunCommand("CONNECT 192.168.1.10")
    RunCommand("CJNT")

    #RunCommand("RUNPROG -1 SetForceConditionOnce(12)")
    #RunCommand("DISCONNECT")
    # print("Tip: Type 'CJNT' to retrieve")
    # print("Tip: Type 'MOVJ j1 j2 j3 j4 j5 j6 j7' to move the robot (provide joints as angles)")
    # except Exception as e:
    #    print(e)

    # input("Test commands finished. Press enter to continue")

    # RunCommand("SETTOOL -0.025 -41.046 50.920 60.000 -0.000 90.000")
    # RunCommand("MOVJ -5.362010 46.323420 20.746290 74.878840 -50.101680 61.958500")
    # RunCommand("SPEED 250")
    # RunCommand("MOVEL 0 0 0 0 0 0 -5.362010 50.323420 20.746290 74.878840 -50.101680 61.958500")
    # RunCommand("PAUSE 2000") # Pause 2 seconds


# -------------------------- Main driver loop -----------------------------
# Read STDIN and process each command (infinite loop)
# IMPORTANT: This must be run from RoboDK so that RoboDK can properly feed commands through STDIN
# This driver can also be run in console mode providing the commands through the console input
def RunDriver():
    for line in sys.stdin:
        RunCommand(line.strip())


# Each line provided through command line or STDIN will be processed by RunCommand
def RunCommand(cmd_line):
    global ROBOT_IP
    global ROBOT_PORT
    global ROBOT

    # strip a line of words into a list of numbers
    def line_2_values(line):
        values = []
        for word in line:
            try:
                number = float(word)
                values.append(number)
            except ValueError:
                pass
        return values

    cmd_words = cmd_line.split(' ')  # [''] if len == 0
    cmd = cmd_words[0]
    cmd_values = line_2_values(cmd_words[1:])  # [] if len <= 1
    n_cmd_values = len(cmd_values)
    n_cmd_words = len(cmd_words)
    success = None

    if cmd_line == "":
        # Skip if no command is provided
        return

    elif cmd_line.startswith("CONNECT"):
        # Connect to robot provided the IP and the port
        global nDOFs_MIN
        if n_cmd_words >= 2:
            ROBOT_IP = cmd_words[1]
        if n_cmd_words >= 3:
            ROBOT_PORT = int(cmd_words[2])
        if n_cmd_words >= 4:
            nDOFs_MIN = int(cmd_words[3])
        success = RobotConnect()

    elif n_cmd_values >= nDOFs_MIN and cmd_line.startswith("MOVJ"):
        UpdateStatus(ROBOTCOM_WORKING)
        # Execute a joint move. RoboDK provides j1,j2,...,j6,j7,x,y,z,w,p,r
        if ROBOT.MoveJ(cmd_values):
            # Notify that we are done with this command
            success = True

    elif n_cmd_values >= nDOFs_MIN and cmd_line.startswith("MOVL"):
        UpdateStatus(ROBOTCOM_WORKING)
        # Execute a linear move. RoboDK provides j1,j2,...,j6,j7,x,y,z,w,p,r
        xyzabc = cmd_values[-6:]
        if ROBOT.MoveL(cmd_values):
            # Notify that we are done with this command
            success = True

    #elif n_cmd_values >= nDOFs_MIN and cmd_line.startswith("MOVLSEARCH"):
    #    UpdateStatus(ROBOTCOM_WORKING)
    #    # Activate the monitor feedback
    #    ROBOT_MOVING = True
    #    # Execute a linear move. RoboDK provides j1,j2,...,j6,x,y,z,w,p,r
    #    if ROBOT.SendCmd(MSG_MOVEL_SEARCH, cmd_values[0:n_cmd_values]):
    #        # Retrieve contact joints
    #        jnts_contact = ROBOT.recv_array()
    #        print_joints(jnts_contact)

    elif n_cmd_values >= 2 * (nDOFs_MIN) and cmd_line.startswith("MOVC"):
        UpdateStatus(ROBOTCOM_WORKING)
        # Activate the monitor feedback
        #ROBOT_MOVING = True
        # Execute a circular move. RoboDK provides j1,j2,...,j6,x,y,z,w,p,r
        xyzwpr12 = cmd_values[-(nDOFs_MIN*2):]
        if ROBOT.MoveC(xyzwpr12):
            # Notify that we are done with this command
            success = True
            UpdateStatus(ROBOTCOM_READY)

    elif cmd_line.startswith("CJNT"):
        UpdateStatus(ROBOTCOM_WORKING)
        # Retrieve the current position of the robot
        curJoints = ROBOT.getJoints()
        if curJoints is not None:
            success = True
            print_joints(curJoints)

    elif n_cmd_values >= 1 and cmd_line.startswith("SPEED"):
        UpdateStatus(ROBOTCOM_WORKING)
        # First value is linear speed in mm/s
        # IMPORTANT! We should only send one "Ready" per instruction
        speed_values = [-1, -1, -1, -1]
        for i in range(max(4, len(cmd_values))):
            speed_values[i] = cmd_values[i]

        # speed_values[0] = speed_values[0] # linear speed in mm/s
        # speed_values[1] = speed_values[1] # joint speed in mm/s
        # speed_values[2] = speed_values[2] # linear acceleration in mm/s2
        # speed_values[3] = speed_values[3] # joint acceleration in deg/s2

        if ROBOT.setSpeed(speed_values):
            # Notify that we are done with this command
            success = True

    elif n_cmd_values >= 6 and cmd_line.startswith("SETTOOL"):
        UpdateStatus(ROBOTCOM_WORKING)
        # Set the Tool reference frame provided the 6 XYZWPR cmd_values by RoboDK
        if ROBOT.setTool(cmd_values):
            # Notify that we are done with this command
            success = True

    elif n_cmd_values >= 1 and cmd_line.startswith("PAUSE"):
        UpdateStatus(ROBOTCOM_WAITING)
        # Run a pause
        if ROBOT.Pause(cmd_values[0]):
            # Notify that we are done with this command
            success = True

    elif n_cmd_values >= 1 and cmd_line.startswith("SETROUNDING"):
        # Set the rounding/smoothing value. Also known as ZoneData in ABB or CNT for Fanuc
        if ROBOT.setRounding(cmd_values[0]):
            # Notify that we are done with this command
            success = True

    elif n_cmd_values >= 2 and cmd_line.startswith("SETDO"):
        UpdateStatus(ROBOTCOM_WORKING)
        # dIO_id = cmd_values[0]
        # dIO_value = cmd_values[1]
        if ROBOT.setDO(cmd_values[0:2]):
            # Notify that we are done with this command
            success = True

    elif n_cmd_values >= 2 and cmd_line.startswith("WAITDI"):
        UpdateStatus(ROBOTCOM_WORKING)
        # dIO_id = cmd_values[0]
        # dIO_value = cmd_values[1]
        if ROBOT.WaitDI(cmd_values[0:3]):
            # Notify that we are done with this command
            success = True

    elif cmd_line.startswith("RUNPROG"): # n_cmd_values >= 1 and n_cmd_words >= 3 and :
        UpdateStatus(ROBOTCOM_WORKING)
        program_id = cmd_values[0]  # Program ID is extracted automatically if the program name is Program ID
        code = cmd_words[2]  # "Program%i" % program_id
        m = re.search(r'^(?P<program_name>.*)\((?P<args>.*)\)', code)
        code_dict = m.groupdict()
        program_name = code_dict['program_name']
        args = code_dict['args'].replace(' ', '').split(',')
        print_message('program_name: ' + program_name)
        print_message('args: ' + str(args))

        print_message(program_name)
        print_message(cmd_line)
        
        try:
            eval(cmd_line[len("RUNPROG"):].strip())

        except Exception as e:
            print_message(str(e))
            UpdateStatus(ROBOTCOM_UNKNOWN)

    elif n_cmd_words >= 2 and cmd_line.startswith("POPUP "):
        UpdateStatus(ROBOTCOM_WORKING)
        message = cmd_line[6:]
        print(message)
        # Wait for command to be executed

        if True:
            # Notify that we are done with this command
            success = True

    elif cmd_line.startswith("DISCONNECT"):
        # Disconnect from robot
        ROBOT.disconnect()
        UpdateStatus(ROBOTCOM_DISCONNECTED)

    elif cmd_line.startswith("STOP") or cmd_line.startswith("QUIT"):
        # Stop the driverç
        ROBOT.disconnect()
        UpdateStatus(ROBOTCOM_DISCONNECTED)
        quit(0)  # Stop the driver

    elif cmd_line.startswith("t"):
        # Call custom procedure for quick testing
        TestDriver()

    else:
        print("Unknown command: " + str(cmd_line))

    if success:
        UpdateStatus(ROBOTCOM_READY)
        
    #else:
    #    UpdateStatus(ROBOTCOM_DISCONNECTED)
    
    # Stop monitoring feedback
    #ROBOT_MOVING = False

    
def RunMain():
    # It is important to disconnect the robot if we force to stop the process
    import atexit

    atexit.register(RobotDisconnect)

    # Flush Disconnected message
    print_message(DRIVER_VERSION)
    UpdateStatus()

    # Run the driver from STDIN
    RunDriver()

    # Test the driver with a sample set of commands
    #TestDriver()

if __name__ == "__main__":
    """Call Main procedure"""
    # Important, leave the Main procedure as RunMain
    RunMain()