#!/usr/bin/env python2.7

### Cesar Rodriguez June 2021
### Script to determine payload and drone state from Gazebo

import rospy,  tf
import rosservice
import time
import math
from tf.transformations import *
from offboard_ex.msg import tethered_status
from geometry_msgs.msg import Pose
from gazebo_msgs.srv import GetLinkState
from oscillation_ctrl.srv import WaypointTrack
from std_msgs.msg import Bool

class Main:

	def __init__(self):

		# rate(s)
		rate = 40 # rate for the publisher method, specified in Hz -- 20 Hz

		self.dt		 = 1.0/rate

		# Variables needed for testing start 
		self.tstart		= rospy.get_time() # Keep track of the start time
		while self.tstart == 0.0:				 # Need to make sure get_rostime works
			self.tstart		= rospy.get_time()
		# How long should we wait before before starting test
		self.param_exists = False
		while self.param_exists == False:
			if rospy.has_param('sim/wait'):
				self.wait	= rospy.get_param('sim/wait') # wait time
				self.param_exists = True
			elif rospy.get_time() - self.tstart >= 3.0:
				break
	
		# Will be set to true when test should start
		self.bool = False

		# initialize variables
		self.phi			= 0.0 #	Payload angle of deflection from x-axis
		self.theta		= 0.0 # Payload angle of deflection from y-axis
		self.tetherL	= 0.0 # Tether length
		self.has_run	= 0		# Boolean to keep track of first run instance
		self.phidot		= 0.0 #	
		self.thetadot	= 0.0 #
		self.phibuf	  =	0.0 # Need buffers to determine their rates
		self.thetabuf = 0.0 #
		self.pload		= True # Check if payload exists
		# Max dot values to prevent 'blowup'
		self.phidot_max   = 3.0
		self.thetadot_max = 3.0 

		# variables for message gen
		self.status 					= tethered_status()
		self.status.drone_id	= 'spiri_with_tether::spiri::base'
		self.status.pload_id	= 'spiri_with_tether::mass::payload'

		# service(s)
		self.service1 = '/gazebo/get_link_state'

		# need service list to check if models have spawned
		self.service_list = rosservice.get_service_list()

		# wait for service to exist
		while self.service1 not in self.service_list:
			print "Waiting for models to spawn..."
			self.service_list = rosservice.get_service_list()
			if rospy.get_time() - self.tstart >= 10.0:
				break

		# publisher(s)
		self.publisher = rospy.Publisher('/status/twoBody_status', tethered_status, queue_size=1)
		self.pub_wd	 	 = rospy.Publisher('/status/path_follow', Bool, queue_size=1)

		self.pub_timer  = rospy.Timer(rospy.Duration(1.0/rate), self.link_state)
		self.path_timer = rospy.Timer(rospy.Duration(40.0/rate), self.path_follow)
	
	def cutoff(self,value,ceiling):
		"""
		Takes in value and returns ceiling
		if value > ceiling. Otherwise, it returns
		value back
		"""
	  # initilize sign
		sign   = 1
	
		# check if value is negative
		if value < 0.0:
			sign = -1
		# Cutoff value at ceiling
		if (value > ceiling or value < -ceiling):
			output = sign*ceiling
		else:
			output = value
		return output

	def euler_array(self,pose):
		"""
		Takes in pose msg object and outputs array of euler angs:
			q[0] = Roll
			q[1] = Pitch
			q[2] = Yaw
		"""
		q = euler_from_quaternion([pose.orientation.x,
															pose.orientation.y,
															pose.orientation.z,
															pose.orientation.w])
		return q
		
	# Get link states (drone and pload) and determine angle between them
	def link_state(self,pub_timer):
		
		try:

			# State which service we are querying
			get_P = rospy.ServiceProxy(self.service1,GetLinkState)

			# Set reference frame
			reference	= ''

			# Establish links needed --> Spiri base and payload
			# P = Position vector
			drone_P = get_P(self.status.drone_id,reference)

			# Get orientation of drone in euler angles
			drone_Eul = self.euler_array(drone_P.link_state.pose)

			# Check if payload is part of simulation
			if not drone_P.success:
				self.status.drone_id	= 'spiri::base'
				drone_P = get_P(self.status.drone_id,reference) # i.e. no payload
				self.pload = False

			pload_P = get_P(self.status.pload_id,reference)

			if not self.has_run == 1:
				if self.pload == True:
					# Get tether length based off initial displacement
					self.tetherL = math.sqrt((drone_P.link_state.pose.position.x - 
												pload_P.link_state.pose.position.x)**2 +
											 (drone_P.link_state.pose.position.y - 
												pload_P.link_state.pose.position.y)**2 +
											 (drone_P.link_state.pose.position.z - 
												pload_P.link_state.pose.position.z)**2)
					rospy.set_param('sim/tether_length',self.tetherL)

				else:
					self.tetherL = 0

				self.has_run = 1

			# Need to detemine their location to get angle of deflection
			# Drone
			drone_Px = drone_P.link_state.pose.position.x
			drone_Py = drone_P.link_state.pose.position.y
			drone_Pz = drone_P.link_state.pose.position.z

			if self.pload == True: # If there is payload, determine the variables
				# Pload
				pload_Px = pload_P.link_state.pose.position.x
				pload_Py = pload_P.link_state.pose.position.y	

				# Determine theta (pitch)
				x_sep = pload_Px - drone_Px

				if math.fabs(x_sep) >= self.tetherL or x_sep == 0:
					self.theta = 0
				else:
					self.theta = math.asin(x_sep/self.tetherL)

				# Determine thetadot
				self.thetadot = (self.theta - self.thetabuf)/self.dt
				self.thetadot = self.cutoff(self.thetadot,self.thetadot_max)
				self.thetabuf = self.theta
			
				# Determine phi (roll)
				y_sep = pload_Py - drone_Py

				if math.fabs(y_sep) >= self.tetherL or y_sep == 0:
					self.phi = 0
				else:
					self.phi = math.asin(y_sep/self.tetherL)

				# Determine phidot
				self.phidot = (self.phi - self.phibuf)/self.dt
				self.phidot = self.cutoff(self.phidot,self.phidot_max)
				self.phibuf = self.phi	# Update buffer
				
			else: # Otherwise, vars = 0			
				x_sep = self.phi = self.phidot = self.theta = self.thetadot = 0

			# Print and save results
			print "\n"
			rospy.loginfo("")
			print"Roll:           "+str(round(drone_Eul[0],2)),"\nPitch:         "+str(round(drone_Eul[1],2)),"\nYaw:           "+str(round(drone_Eul[2],2))
			print "drone pos.x:    " + str(round(drone_Px,2))
			print "drone pos.y:    " + str(round(drone_Py,2))
			print "drone pos.z:    " + str(round(drone_Pz,2))
			print "phi:	        " + str(round(self.phi*180/3.14,3))
			print "theta:          " + str(round(self.theta*180/3.14,3))

			# Populate message
			self.status.drone_pos = drone_P.link_state.pose
			self.status.pload_pos = pload_P.link_state.pose
			self.status.length		= self.tetherL
			self.status.phi 			= self.phi
			self.status.phidot		= self.phidot
			self.status.theta			= self.theta
			self.status.thetadot	= self.thetadot

			# Publish message
			self.publisher.publish(self.status)			

		except rospy.ServiceException as e:
			rospy.loginfo("Get Link State call failed: {0}".format(e))

	def path_follow(self,path_timer):
		now = rospy.get_time()
		if now - self.tstart < self.wait:
			self.bool = False
		else:
			self.bool = True
			self.pub_wd.publish(self.bool)
				

if __name__=="__main__":

	# Initiate ROS node
	rospy.init_node('linkStates_node',anonymous=False)
	try:
		Main() # create class object
		rospy.spin() # loop until shutdown signal

	except rospy.ROSInterruptException:
	        pass