#!/usr/bin/env python

# Author: Brevin Banks
# Created on: 2023-05-1
#
# Copyright (c) 2015-2021 Johns Hopkins University, University of Washington, Worcester Polytechnic Institute
# Released under MIT License


import crtk
import math
import sys
import rospy
import tf
import numpy as np
import PyKDL


# example of application using device.py
class crtk_servo_cv_admittance:

    # configuration
    def configure(self, device_namespace):
        # ROS initialization
        if not rospy.get_node_uri():
            rospy.init_node('crtk_servo_cv_admittance',
                            anonymous=True, log_level=rospy.WARN)

        print(rospy.get_caller_id() +
              ' -> configuring crtk_device_test for: ' + device_namespace)
        # populate this class with all the ROS topics we need
        self.crtk_utils = crtk.utils(self, device_namespace)
        self.crtk_utils.add_operating_state()
        self.crtk_utils.add_measured_cp()
        self.crtk_utils.add_measured_cf()
        self.crtk_utils.add_servo_cv()
        # for all examples
        self.duration = 10  # 10 seconds
        self.rate = 200    # aiming for 200 Hz
        self.samples = self.duration * self.rate

    def run_servo_cv(self):
        # Not yet implemented for Galen robot
        # if not self.enable(60):
        #     print("Unable to enable the device, make sure it is connected.")
        #     return

        # Read the initial position data
        self.CartesianPose = PyKDL.Frame()
        self.CartesianPose.p = self.measured_cp().p
        self.CartesianPose.M = self.measured_cp().M
        self.Cartesianquaternion = (
            self.CartesianPose.M.x, self.CartesianPose.M.y, self.CartesianPose.M.z, self.CartesianPose.M.w)
        self.CartRotEuler = tf.transformations.euler_from_quaternion(
            self.Cartesianquaternion)

        # Read the initial force data
        self.Wrench = self.measured_cf()
        # self.ReadForce = self.ReadWrench().F
        # self.ReadTorque = self.ReadWrench().T

        # Initialize output measured positions
        # Current time position
        self.Pout = np.array([[0], [0], [0], [0], [0], [0]])
        # 1 previous time step position
        self.Poutp = np.array([[0], [0], [0], [0], [0], [0]])

        # Initialize input reference velocities
        # Current time refernce velocity
        self.Vref = np.array([[0], [0], [0], [0], [0], [0]])
        # 1 previous time step velocity
        self.Vrefp = np.array([[0], [0], [0], [0], [0], [0]])
        # 2 previous time steps velocity
        self.Vrefpp = np.array([[0], [0], [0], [0], [0, [0]]])

        # Initialize output measured velocities
        # Current time output velocity
        self.Vout = np.array([[0], [0], [0], [0], [0], [0]])
        # 1 previous time step velocity
        self.Voutp = np.array([[0], [0], [0], [0], [0], [0]])
        # 2 previous time steps velocity
        self.Voutpp = np.array([[0], [0], [0], [0], [0], [0]])

        # Initialize output measured and filtered velocities
        # Current time output velocity
        self.Voutf = np.array([[0], [0], [0], [0], [0], [0]])
        # 1 previous time step velocity
        self.Voutfp = np.array([[0], [0], [0], [0], [0], [0]])
        # 2 previous time steps velocity
        self.Voutfpp = np.array([[0], [0], [0], [0], [0], [0]])

        # Initialize input reference wrench
        # Current time wrench
        self.Fext = np.array([[0], [0], [0], [0], [0], [0]])
        # 1 previous time step velocity
        self.Fextp = np.array([[0], [0], [0], [0], [0], [0]])
        # 2 previous time steps velocity
        self.Fextpp = np.array([[0], [0], [0], [0], [0], [0]])

        # Velocity Differentiation first order lowpass filter cuttoff frequency
        self.wc = 100  # cutt off frequency

        # Grab the relative simulation starting time
        init_time = rospy.get_time()  # the current time in seconds

        k = 0  # Useful iterator for control history
        # Initialize History Arrays
        rext = []  # History for reference wrench Fext (N)
        rv = []  # History for reference input twist Vref (mm/s)
        cv = []  # History for measured output twist Vout (mm/s)
        cfv = []  # History for filtered measured output twist Vout (mm/s)
        pout = []  # History for measured cartesian positions Pout (mm)
        time_vec = []  # History for the value of each time step (s)

        # Set the desired admittance mass and damping values
        # kg Desired Joint imposed masses
        self.m_ad = np.array([[20], [20], [20], [20], [20], [0]])
        # Ns/m Desired Joint imposed damping
        self.b_ad = np.array([[150], [170], [170], [170], [270], [0]])

        # Create loop rate object. rospy.Rate will hold the rate each loop until the desired rate is met each iteration
        runrate = rospy.Rate(self.rate)  # Ros rate at 200 Hz

        # Generate movement: Run the robot until the desired duration is achieved
        while(self.duration > rospy.get_time()-init_time):

            # TIME
            # Grab relative iteration time
            time_vec.append(rospy.get_time()-init_time)
            if (k != 0):  # After first iteration grab Ts as finite difference between iterations
                self.Ts = time_vec[k]-time_vec[k-1]
            else:  # If the first iteration, assign Ts to be just the first recorded time space
                self.Ts = time_vec[k]

            # FORCE
            # Store prior force history
            self.Fextpp = self.Fextp
            self.Fextp = self.Fext
            # Read the Force Sensor
            self.Wrench = self.measured_cf()
            # Put force into control loop compatible form
            self.Fext = np.array([[self.Wrench.F.x], [self.Wrench.F.y], [self.Wrench.F.z], [
                                 self.Wrench.T.x], [self.Wrench.T.y], [self.Wrench.T.z]])

            # ADMITTANCE
            # Apply Admittance control to force. Fext-->Vref
            self.admittance_control()

            # CONTROL COMMAND
            # assigne reference output of admittance_control to the servo commands
            self.servo_cv(self.Vref)

            # OUTPUT POSITION
            self.grab_position()

            # CONVERT OUTPUT TO VELOCITY
            self.finite_differentiate()

            # STORE ITERATION DATA
            rext.append(self.Fext)
            rv.append(np.reshape(self.Vref, (6, 1)))
            cv.append(self.Vout)
            cfv.append(self.Voutf)
            pout.append(self.Pout)

            # LOOP CONTROL
            runrate.sleep()  # hold the loop until 200Hz is achieved
            k = k+1  # increase iterator

            # END LOOP

        # Send a command to stop any velocity in the robot
        stop_vel = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
        self.servo_cv(stop_vel)

    def admittance_control(self):
        # Store prior velocity history
        self.Vrefpp = self.Vrefp
        self.Vrefp = self.Vref
        # Determine the reference velocity for each joint given the tustin transform of Vref(s)/Fext(s) = 1/(m_ad*s + b_ad)
        for joint in range(5):
            self.Vref[joint] = (-self.Vrefp[joint]*(self.b_ad[joint]-((2/self.Ts)*self.m_ad[joint])) +
                                self.Fextp[joint] + self.Fextp[joint])/(self.b_ad[joint] + (2/self.Ts)*self.m_ad[joint])
        # servo_cv needs a 1x6 vector
        self.Vref = np.reshape(self.Vref, (1, 6))
        # Convert the Velocity from m/s to mm/s
        self.Vref = self.Vref*1000

    def grab_position(self):
        self.Poutp = self.Pout
        self.CartesianPose.p = self.measured_cp().p
        self.CartesianPose.M = self.measured_cp().M
        self.Cartesianquaternion = (
            self.CartesianPose.M.x, self.CartesianPose.M.y, self.CartesianPose.M.z, self.CartesianPose.M.w)
        self.CartRotEuler = tf.transformations.euler_from_quaternion(
            self.Cartesianquaternion)
        self.Pout = np.array([[self.CartesianPose.p.x], [self.CartesianPose.p.y], [self.CartesianPose.p.z], [
                             self.CartRotEuler[0]], [self.CartRotEuler[1]], [self.CartRotEuler[2]]])

    def finite_differentiate(self):
        self.Vout = (self.Pout-self.Poutp)/(self.Ts)
        for joint in range(5):
            self.Voutf[joint] = (self.Ts*self.wc*(self.Vout[joint]+self.Voutp[joint])-self.Voutfp[joint]*(
                self.Ts*self.wc-2))/(2+self.Ts*self.wc)  # first order low pass filter from wc/(s+wc)


# use the class now, i.e. main program
if __name__ == '__main__':
    try:
        if (len(sys.argv) != 2):
            print(
                sys.argv[0], ' requires one argument, i.e. crtk device namespace')
        else:
            example = crtk_servo_cv_admittance()
            example.configure(sys.argv[1])
            example.run_servo_cv()

    except rospy.ROSInterruptException:
        pass