#!/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
import math
from matplotlib.cbook import flatten
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 tolist(w): return list(flatten(w))
def adjoint(w):
    w=tolist(w)
    return np.array([[0,-w[2],w[1]] , [w[2],0,-w[0]] , [-w[1],w[0],0]])
    
def adjoint_inv(A): 
    return np.array([[A[2,1]],[A[0,2]],[A[1,0]]])

def expw(w): return expm(adjoint(w))
def logw(R): return adjoint_inv(logm(R))

dt = 0.1
p1, p2, p3, p4, p5 = 1, 3, 2, 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

_B = np.array([[p3, 0, 0], [0, p4, 0], [0, 0, p5]]) @ np.array(
    [[-1, -1, -1], [0, math.sin(np.pi / 3), math.sin(-np.pi / 3)], [-1, math.cos(np.pi / 3), math.cos(np.pi / 3)]])

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.R = np.eye(3)
		self.b = 0 # propeler state
		
	def listener(self):
		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/data', 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 control signal
		servo = np.array([[self.s0,self.s1,self.s2]]).T
		u13 = As@(servo-Bs)
		
		# update states 
		self.v = self.v + dt * (p1*((self.esc-1000)/2000)**2- p2*self.v*abs(self.v))
		
		wr = self.v * _B @ u13
		self.R = self.R @ expw(dt * wr)
		self.z = self.z + self.R[2,0]*self.v
		self.x = self.x + self.R[0,0]*self.v
		self.y = self.y + self.R[1,0]*self.v
		self.b = self.b + (self.esc-1000)/100
		
		# update msgs
		
		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
		
		R_hom = np.concatenate((np.concatenate((self.R, np.zeros((3, 1))), 1), np.array([[0., 0., 0., 1.]])))
		quat =  trf.quaternion_from_matrix(R_hom)
		self.t.transform.rotation.x = quat[1]
		self.t.transform.rotation.y = quat[2]
		self.t.transform.rotation.z = quat[3]
		self.t.transform.rotation.w = quat[0]
		self.imu.orientation = self.t.transform.rotation
		
		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
