#!/usr/bin/env python
import time
from osr_msgs.msg import Joystick, Commands, Encoder, RunStop
from nav_msgs.msg import Odometry
from geometry_msgs.msg import Point, Pose, Quaternion, Twist, Vector3
import rospy
import tf
import math
import numpy

class Odometry2():
    def __init__(self, baseFrame, wheelTrack, mpt, d4, maxTickPerSec, pubTF=False):
        self.encValid = False
        self.priorTime = rospy.Time.now()
        self.priorEncs = [0,0,0,0,0,0]
        self.mpt = mpt
        self.pubTF = pubTF
        
        # distance between wheels
        self.wheelTrack = wheelTrack
        self.d4 = d4
        self.baseFrame = baseFrame
        self.maxTickPerSec = maxTickPerSec

        self.x = 0.0
        self.y = 0.0
        self.th = 0.0

        self.odomPub = rospy.Publisher("/odom", Odometry, queue_size = 1)    
        
        if self.pubTF:
            self.odomBroadcaster = tf.TransformBroadcaster() 

        self.twistCovar = numpy.diag([0.001, 0.001, 0.001, 0.1, 0.1, 0.1]).ravel()
        self.poseCovar = numpy.diag([0.001, 0.001, 0.001, 0.1, 0.1, 0.1]).ravel()


    def onEncoderMessage(self, message):
        self.calculateOdometry(message)

    def isValid(self, message):
        dencLeft = abs(message.rel_enc[1] - self.priorEncs[1])
        dencRight = abs(message.rel_enc[4] - self.priorEncs[4])

        dt = self.getElapsedTime(message.header.stamp)

        if (dencLeft/dt) > self.maxTickPerSec:
            rospy.logwarn("Invalid relative encoder value on left wheel. No odom calculated")
            return False
        
        if (dencRight/dt) > self.maxTickPerSec:
            rospy.logwarn("Invalid relative encoder value on right wheel. No odom calculated")
            return False

        return True

    def publishTransform(self, x, y, quaternion, timestamp):        
        self.odomBroadcaster.sendTransform(
            (x, y, 0),
            (quaternion.x, quaternion.y, quaternion.z, quaternion.w),
            timestamp,
            self.baseFrame,
            "odom")

    def publishOdomMessage(self, x, y, vx, vy, vth, quaternion, timestamp):
        odom = Odometry()
        odom.header.frame_id = "odom"
        odom.child_frame_id = self.baseFrame
        odom.header.stamp = timestamp
        odom.pose.pose.position.x = x
        odom.pose.pose.position.y = y
        odom.pose.pose.position.z = 0
        odom.pose.covariance = self.poseCovar
        odom.pose.pose.orientation = quaternion
        odom.twist.twist.linear.x = vx
        odom.twist.twist.linear.y = vy
        odom.twist.twist.linear.z = 0
        odom.twist.twist.angular.z = vth
        odom.twist.covariance = self.twistCovar        

        self.odomPub.publish(odom)

    def getElapsedTime(self, timestamp, save=False):
        dt = (timestamp - self.priorTime).to_sec()

        if save:
            self.priorTime = timestamp

        return dt

    def calculateTurnRadius(self, dLeft, dRight):
        dlr = dLeft - dRight

         # calculate radius of turn
        if dlr != 0 and dLeft != 0 and dRight != 0:
            lv = self.d4 + dLeft / dRight * self.d4
            # print ("lv: " + str(lv))
            r = lv / (1  - (dLeft / dRight))
        else:
            r = 0

        dist = (dLeft + dRight) / 2

        # calculate angle change    
        if (r != 0):
            dTheta = dist / -r
        else:
            dTheta = 0

        return r, dTheta

    def calculateOdometry(self, message):
        currentTime = message.header.stamp
        encs = message.rel_enc                

        if not self.isValid(message):
            return

        dt = self.getElapsedTime(currentTime, save=True)

        dLeft = self.mpt * (encs[1] - self.priorEncs[1])
        dRight = self.mpt * (encs[4] - self.priorEncs[4])
    
        # dth = (dRight - dLeft) / self.wheelTrack

        radius, dTheta = self.calculateTurnRadius(dLeft, dRight)

        # calculate centre of turn circle
        xOrig = self.x + radius * math.cos(self.th)
        yOrig = self.y + radius * math.sin(self.th)

        # calculate new co-ordinates    
        xNew = xOrig + (self.x - xOrig) * math.cos(dTheta) - (self.y - yOrig) * math.sin(dTheta)
        yNew = yOrig + (self.x - xOrig) * math.sin(dTheta) + (self.y - yOrig) * math.cos(dTheta)

        #calculate change in x,y values
        dx = xNew - self.x
        dy = yNew - self.y     

        self.th += dTheta

        if (self.th > (math.pi * 2)):
            self.th -= (math.pi * 2)
        elif (self.th < (-math.pi * 2)):
            self.th += (math.pi * 2)

        self.x = xNew
        self.y = yNew

        # convert to ros co-ords
        xRos = self.y
        yRos = -self.x

        vxRos = dy / dt
        vyRos = -dx / dt

        vth = dTheta /dt

        quaternion = self.getQuaternion(self.th)
        
        if self.pubTF:
            self.publishTransform(xRos, yRos, quaternion, currentTime)
            
        self.publishOdomMessage(xRos, yRos, vxRos, vyRos, vth, quaternion, currentTime)   

        self.priorEncs = encs    
        

    def getQuaternion(self, th):
        quaternion = Quaternion()
        quaternion.x  = 0.0
        quaternion.y  = 0.0
        quaternion.z = math.sin(th / 2.0)
        quaternion.w = math.cos(th / 2.0)

        return quaternion


if __name__ == '__main__':    

    rospy.init_node('osr_odometry2')
    rospy.loginfo("Starting the osr odometry2 node")	

    baseFrame = rospy.get_param("/odometry/base_frame_id", "base_link")
    # mpt = rospy.get_param("/odometry/mpt", 0.000026322)
    mpt = rospy.get_param("/odometry/mpt",  0.000100708)     
    wheelTrack = rospy.get_param("/odometry/wheel_track", 0.455)
    d4 = rospy.get_param("/odometry/d4", 0.2559)
    maxTickPerSec = rospy.get_param("/odometry/maxTickPerSec", 8000)        
    publishTF = rospy.get_param("~publishTF", False)

    odom = Odometry2(baseFrame, wheelTrack, mpt, d4, maxTickPerSec, pubTF=publishTF)
    encSub  = rospy.Subscriber("/encoder", Encoder, odom.onEncoderMessage)	    

    rate = rospy.Rate(20)

    while not rospy.is_shutdown():        
        rate.sleep()	



[((181, 4, 181, 36), 'rospy.init_node', 'rospy.init_node', ({(181, 20, 181, 35): '"""osr_odometry2"""'}, {}), "('osr_odometry2')", False, 'import rospy\n'), ((182, 4, 182, 52), 'rospy.loginfo', 'rospy.loginfo', ({(182, 18, 182, 51): '"""Starting the osr odometry2 node"""'}, {}), "('Starting the osr odometry2 node')", False, 'import rospy\n'), ((184, 16, 184, 71), 'rospy.get_param', 'rospy.get_param', ({(184, 32, 184, 57): '"""/odometry/base_frame_id"""', (184, 59, 184, 70): '"""base_link"""'}, {}), "('/odometry/base_frame_id', 'base_link')", False, 'import rospy\n'), ((186, 10, 186, 56), 'rospy.get_param', 'rospy.get_param', ({(186, 26, 186, 41): '"""/odometry/mpt"""', (186, 44, 186, 55): '0.000100708'}, {}), "('/odometry/mpt', 0.000100708)", False, 'import rospy\n'), ((187, 17, 187, 64), 'rospy.get_param', 'rospy.get_param', ({(187, 33, 187, 56): '"""/odometry/wheel_track"""', (187, 58, 187, 63): '0.455'}, {}), "('/odometry/wheel_track', 0.455)", False, 'import rospy\n'), ((188, 9, 188, 48), 'rospy.get_param', 'rospy.get_param', ({(188, 25, 188, 39): '"""/odometry/d4"""', (188, 41, 188, 47): '0.2559'}, {}), "('/odometry/d4', 0.2559)", False, 'import rospy\n'), ((189, 20, 189, 68), 'rospy.get_param', 'rospy.get_param', ({(189, 36, 189, 61): '"""/odometry/maxTickPerSec"""', (189, 63, 189, 67): '8000'}, {}), "('/odometry/maxTickPerSec', 8000)", False, 'import rospy\n'), ((190, 16, 190, 52), 'rospy.get_param', 'rospy.get_param', ({(190, 32, 190, 44): '"""~publishTF"""', (190, 46, 190, 51): 'False'}, {}), "('~publishTF', False)", False, 'import rospy\n'), ((193, 14, 193, 74), 'rospy.Subscriber', 'rospy.Subscriber', ({(193, 31, 193, 41): '"""/encoder"""', (193, 43, 193, 50): 'Encoder', (193, 52, 193, 73): 'odom.onEncoderMessage'}, {}), "('/encoder', Encoder, odom.onEncoderMessage)", False, 'import rospy\n'), ((195, 11, 195, 25), 'rospy.Rate', 'rospy.Rate', ({(195, 22, 195, 24): '20'}, {}), '(20)', False, 'import rospy\n'), ((14, 25, 14, 41), 'rospy.Time.now', 'rospy.Time.now', ({}, {}), '()', False, 'import rospy\n'), ((29, 23, 29, 73), 'rospy.Publisher', 'rospy.Publisher', (), '', False, 'import rospy\n'), ((66, 15, 66, 25), 'nav_msgs.msg.Odometry', 'Odometry', ({}, {}), '()', False, 'from nav_msgs.msg import Odometry\n'), ((170, 21, 170, 33), 'geometry_msgs.msg.Quaternion', 'Quaternion', ({}, {}), '()', False, 'from geometry_msgs.msg import Point, Pose, Quaternion, Twist, Vector3\n'), ((173, 23, 173, 41), 'math.sin', 'math.sin', ({(173, 32, 173, 40): 'th / 2.0'}, {}), '(th / 2.0)', False, 'import math\n'), ((174, 23, 174, 41), 'math.cos', 'math.cos', ({(174, 32, 174, 40): 'th / 2.0'}, {}), '(th / 2.0)', False, 'import math\n'), ((197, 14, 197, 33), 'rospy.is_shutdown', 'rospy.is_shutdown', ({}, {}), '()', False, 'import rospy\n'), ((32, 35, 32, 60), 'tf.TransformBroadcaster', 'tf.TransformBroadcaster', ({}, {}), '()', False, 'import tf\n'), ((48, 12, 48, 93), 'rospy.logwarn', 'rospy.logwarn', ({(48, 26, 48, 92): '"""Invalid relative encoder value on left wheel. No odom calculated"""'}, {}), "(\n    'Invalid relative encoder value on left wheel. No odom calculated')", False, 'import rospy\n'), ((52, 12, 52, 94), 'rospy.logwarn', 'rospy.logwarn', ({(52, 26, 52, 93): '"""Invalid relative encoder value on right wheel. No odom calculated"""'}, {}), "(\n    'Invalid relative encoder value on right wheel. No odom calculated')", False, 'import rospy\n'), ((34, 26, 34, 74), 'numpy.diag', 'numpy.diag', ({(34, 37, 34, 73): '[0.001, 0.001, 0.001, 0.1, 0.1, 0.1]'}, {}), '([0.001, 0.001, 0.001, 0.1, 0.1, 0.1])', False, 'import numpy\n'), ((35, 25, 35, 73), 'numpy.diag', 'numpy.diag', ({(35, 36, 35, 72): '[0.001, 0.001, 0.001, 0.1, 0.1, 0.1]'}, {}), '([0.001, 0.001, 0.001, 0.1, 0.1, 0.1])', False, 'import numpy\n'), ((129, 34, 129, 51), 'math.cos', 'math.cos', ({(129, 43, 129, 50): 'self.th'}, {}), '(self.th)', False, 'import math\n'), ((130, 34, 130, 51), 'math.sin', 'math.sin', ({(130, 43, 130, 50): 'self.th'}, {}), '(self.th)', False, 'import math\n'), ((133, 80, 133, 96), 'math.sin', 'math.sin', ({(133, 89, 133, 95): 'dTheta'}, {}), '(dTheta)', False, 'import math\n'), ((134, 80, 134, 96), 'math.cos', 'math.cos', ({(134, 89, 134, 95): 'dTheta'}, {}), '(dTheta)', False, 'import math\n'), ((133, 42, 133, 58), 'math.cos', 'math.cos', ({(133, 51, 133, 57): 'dTheta'}, {}), '(dTheta)', False, 'import math\n'), ((134, 42, 134, 58), 'math.sin', 'math.sin', ({(134, 51, 134, 57): 'dTheta'}, {}), '(dTheta)', False, 'import math\n')]