#!/usr/bin/env python3

import rospy 
import tf_conversions
import tf2_ros
import transformations as trf
import csv
import geometry_msgs.msg 
import numpy as np
from visualization_msgs.msg import Marker, MarkerArray
from std_msgs.msg import Int32, Float64, Float32
from sensor_msgs.msg import Imu, JointState
from scipy.linalg import expm,logm

def adjoint_inv(A): 
    return np.array([[A[2,1]],[A[0,2]],[A[1,0]]])
def logw(R): return adjoint_inv(logm(R))

dt = 0.1
p1, p2, p3, p4, p5 = 1, 3, 0.5, 0.5, 0.5
As = np.array([[0, -np.pi/2000,0],[-np.pi/2000,0,0],[0,0,-np.pi/2000]])
Bs = np.array([[1500, 1500, 1500]]).T

class Riptide_states:
	def __init__(self):
		self.imu   = Imu()
		self.mbar  = float()
		self.mbar0 = float(1000)
		
		self.s0 = int(1500)
		self.s1 = int(1500)
		self.s2 = int(1500)
		self.esc = int(1000)
		
		# initial state
		self.v = 0
		self.x = 0
		self.y = 0 
		self.z = 0
		self.z_bis = 0
		self.b = 0 # propeler state
		
	def listener(self):
		rospy.Subscriber("/imu/data", Imu, self.imu_callback)
		rospy.Subscriber("/pressure_mbar", Float32, self.mbar_callback)
		rospy.Subscriber("/pololu/esc", Int32, self.esc_callback)
		rospy.Subscriber("/pololu/servo0", Int32, self.s0_callback)
		rospy.Subscriber("/pololu/servo1", Int32, self.s1_callback)
		rospy.Subscriber("/pololu/servo2", Int32, self.s2_callback)
		
	def publisher(self):
		self.br = tf2_ros.TransformBroadcaster()
		self.path = rospy.Publisher('riptide_path', Marker, queue_size=10)
		self.path_bis = rospy.Publisher('riptide_path_bis', Marker, queue_size=10)
		self.joint = rospy.Publisher('joint_states', JointState, queue_size=10)
		self.present_imu_2_rviz = rospy.Publisher('imu_2_rviz', Imu, queue_size=10)
		
		# msgs
		self.t = geometry_msgs.msg.TransformStamped()
		
		self.P = Marker()
		self.P.type = 4 # LINE_STRIP
		self.P.action = 0 # ADD
		self.P.scale.x = 0.05 # line width
		self.P.color.r = 1.0
		self.P.color.g = 1.0 
		self.P.color.a = 1.0 
		self.P.header.frame_id = "fixed_frame"
		self.P.id = 0
		self.P.pose.orientation.w = 1.0
		
		self.P_bis = Marker()
		self.P_bis.type = 4 # LINE_STRIP
		self.P_bis.action = 0 # ADD
		self.P_bis.scale.x = 0.05 # line width
		self.P_bis.color.g = 1.0 
		self.P_bis.color.a = 1.0 
		self.P_bis.header.frame_id = "fixed_frame"
		self.P_bis.id = 10
		self.P_bis.pose.orientation.w = 1.0
		
		self.J = JointState()
		self.J.header.frame_id = "base_link"
		self.J.name = ["base_to_propel","base_to_top_fin","base_to_right_fin","base_to_left_fin"]
	
	def publish(self):
	
		# update states 
		self.v = self.v + dt * (p1*((self.esc-1000)/2000)**2- p2*self.v*abs(self.v))
		
		quat = self.imu.orientation
		R = trf.quaternion_matrix([quat.w,quat.x,quat.y,quat.z])[0:3,0:3]
		Rd = np.eye(3)
		ew = R.T @ logw(Rd @ R.T)
		self.z = -float((self.mbar-self.mbar0)/100)
		self.z_bis = self.z_bis + R[2,0]*self.v
		self.x = self.x + R[0,0]*self.v
		self.y = self.y + R[1,0]*self.v
		self.b = self.b + (self.esc-1000)/100
		
		# update msgs
		servo = np.array([[self.s0,self.s1,self.s2]]).T
		u13 = As@(servo-Bs)
		
		self.t.header.stamp = rospy.Time.now()
		self.t.header.frame_id = "fixed_frame"
		self.t.child_frame_id = "base_link"

		self.t.transform.translation.x = self.x
		self.t.transform.translation.y = self.y
		self.t.transform.translation.z = self.z
		
		
		if quat.x+quat.y+quat.z+quat.w == 0: # if no imu singal
			quat.w = 1.0
		self.t.transform.rotation = quat
		
		p = geometry_msgs.msg.Point()
		p.x = self.x
		p.y = self.y
		p.z = self.z
		
		self.P.header.stamp = rospy.Time.now()
		self.P.points.append(p)
		
		p_bis = geometry_msgs.msg.Point()
		p_bis.x = self.x
		p_bis.y = self.y
		p_bis.z = self.z_bis
		
		self.P_bis.header.stamp = rospy.Time.now()
		self.P_bis.points.append(p_bis)
		
		self.J.position = [self.b,u13[0,0],u13[1,0],u13[2,0]] 
		self.J.header.stamp = rospy.Time.now()
		self.imu.header.stamp = rospy.Time.now()
		self.imu.header.frame_id = "imu"
		
		# publish
		self.br.sendTransform(self.t)
		self.path.publish(self.P)
		self.path_bis.publish(self.P_bis)
		self.joint.publish(self.J)
		self.present_imu_2_rviz.publish(self.imu)
	
	def imu_callback(self,data):
		self.imu = data
		
	def mbar_callback(self,data):
		self.mbar = data.data
		
	def esc_callback(self,data):
		self.esc = data.data
		
	def s0_callback(self,data):
		self.s0 = data.data
		
	def s1_callback(self,data):
		self.s1 = data.data
		
	def s2_callback(self,data):
		self.s2 = data.data
	
			
def talker():
	
	#start node and publishers 
	rospy.init_node('csv_2_rviz')
	
	rip = Riptide_states()
	rip 
	rip.listener()  # subscribers
	rip.publisher() # publishers
	rospy.sleep(2)	
	rip.mbar0 = rip.mbar
	# loop
	while not rospy.is_shutdown():
		rip.publish()
		rospy.sleep(dt)
	

if __name__ == '__main__':
	
	try:
		talker()
	except rospy.ROSInterruptException:
		pass
