#!/usr/bin/env python


# MTRX5700 Assignment 3 question 2 transform listener and path publisher
# Author: Tim Martin
# Edited: 07/05/21

import rospy
import math
import numpy as np
import tf2_ros
from tf.transformations import *
import geometry_msgs.msg
from nav_msgs.msg import Odometry, Path

class Turtle_Odometry:

    def __init__(self):
        self.odom_sub = rospy.Subscriber("/odom", Odometry, self.odom_callback)
        self.odom_quat = quaternion_from_euler(0,0,0)
        self.Opos_x = 0
        self.Opos_y = 0
        self.Opos_z = 0

    def odom_callback(self, msg):
        self.Opos_x = msg.pose.pose.position.x
        self.Opos_y = msg.pose.pose.position.y
        self.Opos_z = msg.pose.pose.position.z

        self.odom_quat[0] = msg.pose.pose.orientation.x
        self.odom_quat[1] = msg.pose.pose.orientation.y
        self.odom_quat[2] = msg.pose.pose.orientation.z
        self.odom_quat[3] = msg.pose.pose.orientation.w


if __name__ == '__main__':
    rospy.init_node('tf2_listener')
    Odom_stuff = Turtle_Odometry()
    tfBuffer = tf2_ros.Buffer()
    listener = tf2_ros.TransformListener(tfBuffer)

    map_odom_pub = rospy.Publisher('map_odom', Odometry, queue_size=1)
    map_path_pub = rospy.Publisher('Trajectory_path', Path, queue_size=10)
    map_odompath_pub = rospy.Publisher('Trajectory_odompath', Path, queue_size=10)

    # odom_path_p = np.array([])
    # odom_path_o = np.array([])
    odom_path_p = np.zeros(shape=(200000, 3))
    odom_path_o = np.zeros(shape=(200000, 4))
    map_path = Path()
    map_path.header.frame_id = "odom"
    prev_trans = geometry_msgs.msg.Transform()
    curr_trans = geometry_msgs.msg.Transform()
    first_flag = 0
    i = 0
    pos_num = 0

    rate = rospy.Rate(10)
    while not rospy.is_shutdown():
        try:
            trans = tfBuffer.lookup_transform('map', 'odom', rospy.Time())
	    print('trans is: ', trans)
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException):
            rate.sleep()
            continue

        # retrieve components of map to odom transformation
        #print(trans)
        odom_pos = np.array([Odom_stuff.Opos_x, Odom_stuff.Opos_y, Odom_stuff.Opos_z])
	print('odom_pos', odom_pos)
        odom_quat = np.array([Odom_stuff.odom_quat[0], Odom_stuff.odom_quat[1], Odom_stuff.odom_quat[2], Odom_stuff.odom_quat[3]])
	print('odom_quat', odom_quat)

        tf_x = trans.transform.translation.x
        tf_y = trans.transform.translation.y
        tf_z = trans.transform.translation.z
        tf_trans = np.array([tf_x, tf_y, tf_z])
	print ('tf_trans',tf_trans)
        #print(odom_pos)

        tf_quat = quaternion_from_euler(0,0,0)
        tf_quat[0] = trans.transform.rotation.x
        tf_quat[1] = trans.transform.rotation.y
        tf_quat[2] = trans.transform.rotation.z
        tf_quat[3] = -trans.transform.rotation.w
	print ('tf_quat',tf_quat)
        #print(type(tf_quat))

        # save current odom pose to a path array
        odom_path_p[pos_num] = odom_pos
        odom_path_o[pos_num] = odom_quat
        pos_num += 1
        # if first_flag == 0:
        #     odom_path_o = np.append(odom_path_o, odom_pos, axis=0)
        #     odom_path_p = np.append(odom_path_p, odom_quat, axis=0)
        #     first_flag = 1
        # elif first_flag == 1:
        #     odom_path_o = np.append([odom_path_o], [odom_pos], axis=0)
        #     odom_path_p = np.append([odom_path_p], [odom_quat], axis=0)
        #     first_flag = 2
        # else:
        #     odom_path_o = np.append(odom_path_o, [odom_pos], axis=0)
        #     odom_path_p = np.append(odom_path_p, [odom_quat], axis=0)

        #odom_path.poses.append(Odom_stuff.odom_pose)
        #print(odom_path_o[0])
        #print(odom_path_o[-1])
        #print("\n\n")



        # apply current transform to all path poses
        updated_path = Path()
	nontransformed_path=Path()
        updated_path.header.frame_id = "odom"
	nontransformed_path.header.frame_id = "odom"

        #L = len(odom_path_o)/3

        if prev_trans == trans:
            L = range(int(len(odom_path_o) / 3))
            #print("old")
        else:
            L = [-1]
            #print("new")

        for c in L:
            i += 1
            current_position = odom_path_p[c]
            current_orientation = odom_path_o[c]
            # print("current")
            # print(current_orientation)
            # perform transform
            tmp_transform = tf_quat
            #tmp_transform[3] = tmp_transform[3]
            transformed_orientation = quaternion_multiply(tmp_transform, current_orientation)
            # print("tf_trans")
            # print(tf_trans)
            # print("current_position")
            # print(current_position)
            transformed_position = np.add(tf_trans, current_position)
            # print("Transformed")
            # print(transformed_orientation)
            # print("\n\n")
            # prepare transformed waypoint
            transformed_pose = geometry_msgs.msg.PoseStamped()
            transformed_pose.header.frame_id = "waypoint"
            transformed_pose.pose.orientation.x = transformed_orientation[0]
            transformed_pose.pose.orientation.y = transformed_orientation[1]
            transformed_pose.pose.orientation.z = transformed_orientation[2]
            transformed_pose.pose.orientation.w = transformed_orientation[3]
            transformed_pose.pose.position.x = transformed_position[0]
            transformed_pose.pose.position.y = transformed_position[1]
            transformed_pose.pose.position.z = transformed_position[2]

	# prepare transformed waypoint CV add
            nontransformed_pose = geometry_msgs.msg.PoseStamped()
            nontransformed_pose.header.frame_id = "waypoint"
            nontransformed_pose.pose.orientation.x = current_orientation[0]
            nontransformed_pose.pose.orientation.y = current_orientation[1]
            nontransformed_pose.pose.orientation.z = current_orientation[2]
            nontransformed_pose.pose.orientation.w = current_orientation[3]
            nontransformed_pose.pose.position.x = current_position[0]
            nontransformed_pose.pose.position.y = current_position[1]
            nontransformed_pose.pose.position.z = current_position[2]

            # append waypoint to path
            updated_path.poses.append(transformed_pose)

 	    # append waypoint to path
            nontransformed_path.poses.append(nontransformed_pose)


            #print(updated_path)

        #print(updated_path.poses[0])
        #print(updated_path.poses[1])
        # send path as message

        # apply transformation to current pose
        new_quat = quaternion_multiply(tf_quat,Odom_stuff.odom_quat)
        new_pos = np.add(tf_trans, odom_pos)


        # create the New Map framed odometry message
        map_odom = Odometry()
        map_odom.header.stamp = rospy.Time.now()
        map_odom.header.frame_id = "odom"
        map_odom.pose.pose.position.x = new_pos[0]
        map_odom.pose.pose.position.y = new_pos[1]
        map_odom.pose.pose.position.z = new_pos[2]
        new_quat_msg = geometry_msgs.msg.Quaternion(new_quat[0], new_quat[1], new_quat[2], new_quat[3])
        map_odom.pose.pose.orientation = new_quat_msg
        #print(map_odom.pose.pose)

        # append most recent pose to create the uncorrected path message
        # map_wp = geometry_msgs.msg.PoseStamped()
        # map_wp.header.frame_id = "waypoint"
        # map_wp.pose = map_odom.pose.pose
        #map_wp.pose.position = map_odom.pose.pose.position
        #map_wp.pose.orientation = map_odom.pose.pose.orientation
        #map_path.poses.append(map_wp)
        #print(updated_path.poses[1])
        #print(map_path.poses)
        #print('\n')

        # publish this new odometry and path message
        #map_odom_pub.publish(map_odom)
        #map_path_pub.publish(map_path)
        map_path_pub.publish(updated_path)
	#CV add
	map_odompath_pub.publish(nontransformed_path)
        prev_trans = trans
        rate.sleep()