#!/usr/bin/env python3

from __future__ import print_function

import rospy
from std_msgs.msg import String
from std_msgs.msg import Int32
import odrive
from odrive.enums import *
import time
import math

from nsra_odrive_interface.srv import *

positions = [0, 0, 0, 0, 0, 0]

hz = 20

print("Connecting...")
odrv2 = odrive.find_any(serial_number="20843881304E")
print("Connected to ODrive 2")

rospy.set_param('odrive2', True)

def drive_pub4(data):
    print("Axis 4: " + str(data.data))
    if abs(data.data-positions[3])*hz > 20000:
        odrv2.axis0.controller.config.vel_limit = abs(data.data-positions[3])*hz
    else:
        odrv2.axis0.controller.config.vel_limit = 20000
    odrv2.axis0.controller.input_pos = data.data
    #drive4.publish(int(data.data*(1024000)/math.pi))
    positions[3] = data.data

def drive_pub5(data):
    print("Axis 5: " + str(data.data))
    if abs(data.data-positions[4])*hz > 20000:
        odrv2.axis1.controller.config.vel_limit = abs(data.data-positions[4])*hz
    else:
        odrv2.axis1.controller.config.vel_limit = 20000
    odrv2.axis1.controller.input_pos = data.data
    #drive5.publish(int(data.data*(1024000)/math.pi))
    positions[4] = data.data

def axis_requested_state(req):
    print("Axis requested_state")

    if req.axis == 4:
        if req.requested_state == 0:
            odrv2.axis0.requested_state = AXIS_STATE_FULL_CALIBRATION_SEQUENCE
            while odrv2.axis0.current_state != AXIS_STATE_IDLE:
                time.sleep(0.1)
            odrv2.axis0.requested_state = AXIS_STATE_CLOSED_LOOP_CONTROL
        else:
            odrv2.axis0.requested_state = req.requested_state
    if req.axis == 5:
        if req.requested_state == 0:
            odrv2.axis1.requested_state = AXIS_STATE_FULL_CALIBRATION_SEQUENCE
            while odrv2.axis1.current_state != AXIS_STATE_IDLE:
                time.sleep(0.1)
            odrv2.axis1.requested_state = AXIS_STATE_CLOSED_LOOP_CONTROL
        else:
            odrv2.axis1.requested_state = req.requested_state

    return requested_stateResponse(True)

rospy.init_node('driver')

rospy.Subscriber("/nsra/drive_pub4", Int32, drive_pub4)
rospy.Subscriber("/nsra/drive_pub5", Int32, drive_pub5)

requested_state = rospy.Service('odrv2_requested_state', requested_state, axis_requested_state)

rospy.spin()
