#!/usr/bin/env python

import rospy
from geometry_msgs.msg import PoseStamped
from styx_msgs.msg import Lane, Waypoint
from std_msgs.msg import Int32
from scipy.spatial import KDTree
import math
import numpy as np
'''
This node will publish waypoints from the car's current position to some `x` distance ahead.

As mentioned in the doc, you should ideally first implement a version which does not care
about traffic lights or obstacles.

Once you have created dbw_node, you will update this node to use the status of traffic lights too.

Please note that our simulator also provides the exact location of traffic lights and their
current status in `/vehicle/traffic_lights` message. You can use this message to build this node
as well as to verify your TL classifier.

TODO (for Yousuf and Aaron): Stopline location for each traffic light.
'''

LOOKAHEAD_WPS = 200 # Number of waypoints we will publish. You can change this number
MAX_DECEL = 1.0

class WaypointUpdater(object):
    def __init__(self):
        rospy.init_node('waypoint_updater')

        rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
        rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)

        # TODO: Add a subscriber for /traffic_waypoint and /obstacle_waypoint below
	rospy.Subscriber('/traffic_waypoint',Int32,self.traffic_cb)

        self.final_waypoints_pub = rospy.Publisher('final_waypoints', Lane, queue_size=1)

        # TODO: Add other member variables you need below
	self.waypoints_2d = None	
	self.base_waypoints = None
	self.waypoints_tree = None
	self.pose = None
	self.stopline_wp_idx = -1

        #rospy.spin()
	self.loop() 

    def loop(self):
	rate = rospy.Rate(10)
	while not rospy.is_shutdown():
		if self.pose and self.base_waypoints:
			closest_idx = self.get_closest_idx()
			self.publish_waypoints(closest_idx)
			rate.sleep()
    def get_closest_idx(self):
	#get current position
	x = self.pose.pose.position.x
	y = self.pose.pose.position.y
	closest_idx = self.waypoints_tree.query([x,y],1)[1]

	#get closest index from tree
	pos_vec = np.array([x,y])
	prev_vec = np.array(self.waypoints_2d[closest_idx-1])
	clos_vec = np.array(self.waypoints_2d[closest_idx])
	ahead = np.dot(clos_vec-prev_vec,pos_vec-clos_vec)
	if ahead > 0:
		closest_idx = (closest_idx + 1) % len(self.waypoints_2d)
	
	#check if closest isahead or behind the car
	# dor product of prev index vs new closest  and closest vs position of car
	#increment if behding the car
	return closest_idx

    def publish_waypoints(self,closest_idx):
	lane = Lane()
	lane.header = self.base_waypoints.header
	farthest_idx = closest_idx + LOOKAHEAD_WPS
	lane.waypoints = self.base_waypoints.waypoints[closest_idx:farthest_idx]
	#rospy.logwarn("stop waypoint index {0}".format(self.stopline_wp_idx))
	if self.stopline_wp_idx != -1 and self.stopline_wp_idx < farthest_idx:
		lane.waypoints = self.decelerate_waypoints(lane.waypoints, closest_idx)
	self.final_waypoints_pub.publish(lane)

    def decelerate_waypoints(self, waypoints, closest_idx):
	
	temp = []
	for i,wp in enumerate(waypoints):
		p = Waypoint()
		p.pose = wp.pose
		stop_idx = max(self.stopline_wp_idx-closest_idx-2,0)
		dist = self.distance(waypoints, i, stop_idx)
		vel = math.sqrt(2 * MAX_DECEL * dist)
		if (vel < 1):
			vel = 0.
		p.twist.twist.linear.x = min(vel, wp.twist.twist.linear.x)
		temp.append(p)
		#if (i%5==0):
		#	rospy.logwarn("decelerating {0} index with velocity {1}".format(i+closest_idx,vel))
	return temp    
    def pose_cb(self, msg):
        # TODO: Implement
	self.pose = msg
        

    def waypoints_cb(self, waypoints):
        # TODO: Implement
	self.base_waypoints = waypoints
	if not self.waypoints_2d:
		self.waypoints_2d = [[waypoint.pose.pose.position.x, waypoint.pose.pose.position.y] for waypoint in waypoints.waypoints]
		self.waypoints_tree = KDTree(self.waypoints_2d)


    def traffic_cb(self, msg):
        # TODO: Callback for /traffic_waypoint message. Implement
        
	self.stopline_wp_idx = msg.data

    def obstacle_cb(self, msg):
        # TODO: Callback for /obstacle_waypoin:t message. We will implement it later
        pass

    def get_waypoint_velocity(self, waypoint):
        return waypoint.twist.twist.linear.x

    def set_waypoint_velocity(self, waypoints, waypoint, velocity):
        waypoints[waypoint].twist.twist.linear.x = velocity

    def distance(self, waypoints, wp1, wp2):
        dist = 0
        dl = lambda a, b: math.sqrt((a.x-b.x)**2 + (a.y-b.y)**2  + (a.z-b.z)**2)
        for i in range(wp1, wp2+1):
            dist += dl(waypoints[wp1].pose.pose.position, waypoints[i].pose.pose.position)
            wp1 = i
        return dist


if __name__ == '__main__':
    try:
        WaypointUpdater()
    except rospy.ROSInterruptException:
        rospy.logerr('Could not start waypoint updater node.')
