#!/usr/bin/env python
import numpy as np

import rospy
from geometry_msgs.msg import PoseStamped
from styx_msgs.msg import Lane, Waypoint
from std_msgs.msg import Int32


import math
from scipy.spatial import KDTree



'''
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.
'''
MAX_DECEL = 0.5
LOOKAHEAD_WPS = 100 # Number of waypoints we will publish. You can change this number
CONST_DECEL = 1 / LOOKAHEAD_WPS # Deceleration constant for smooth breaking
PUBLISHING_RATE = 20
STOP_LINE_MARGIN = 4 # Padding between stopline and the car's center


class WaypointUpdater(object):
    def __init__(self):
        rospy.init_node('waypoint_updater')
        #rospy.loginfo('created updator node')

        # TODO: Add other member variables you need below
        self.base_Way_Points  = None
        self.waypoints_2d = None
        self.pose = None
        self.waypoint_tree = None
        self.stop_wp_idx = None 

        # TODO: Add a subscriber for /traffic_waypoint and /obstacle_waypoint below
        rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
        rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
        rospy.Subscriber('/traffic_waypoint', Int32, self.traffic_cb)

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

    def final_waypoints_callback(self):
        #rospy.loginfo('inside final waypoint callback')
        rate = rospy.Rate(PUBLISHING_RATE)
        while not rospy.is_shutdown():
            if not None in [self.pose, self.base_Way_Points, self.waypoint_tree, self.waypoints_2d, self.stop_wp_idx]:
                # Get closest waypoint 
                self.publish_waypoints()
            rate.sleep()
    
    def get_closest_waypoint_idx(self):
        x = self.pose.pose.position.x
        y = self.pose.pose.position.y

        closest_idx = self.waypoint_tree.query([x,y], 1)[1]

        # Hyperplane fun to see if point is ahead or not
        closest_coord = self.waypoints_2d[closest_idx]
        prev_coord = self.waypoints_2d[closest_idx - 1]
        cl_vect = np.array(closest_coord)
        prev_vect = np.array(prev_coord)
        pos_vect = np.array([x,y])
        val = np.dot(cl_vect - prev_vect, pos_vect - cl_vect)

        if val > 0:
            closest_idx = (closest_idx + 1) % len(self.waypoints_2d)

        return closest_idx
    
    def publish_waypoints(self) :
        final_lane = self.generate_lane()
        self.final_waypoints_pub.publish(final_lane)

    def generate_lane(self):
        lane = Lane()
        
        closest_idx = self.get_closest_waypoint_idx()
        farthest_idx = closest_idx + LOOKAHEAD_WPS
        base_Way_Points = self.base_Way_Points.waypoints[closest_idx:farthest_idx]
    
         
        rospy.logwarn( "current={}, farthest={}, stop={}".format(closest_idx, farthest_idx, self.stop_wp_idx))
        if (self.stop_wp_idx >= farthest_idx or self.stop_wp_idx < closest_idx):
           rospy.logwarn('traffic light is further than the farthest waypoint')
           lane.waypoints = base_Way_Points
        else:
            rospy.logwarn('red light close, decelerate')
            lane.waypoints = self.decelerate_waypoints(base_Way_Points, closest_idx)
        
        return 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.stop_wp_idx - closest_idx - STOP_LINE_MARGIN, 0) # stopping a couple of waypoints before the end
            dist = self.distance(waypoints, i, stop_idx)
            vel = math.sqrt(2 * MAX_DECEL * dist) + (i * CONST_DECEL)

            if vel < 1:
                vel = 0.
            
            p.twist.twist.linear.x = min(vel, wp.twist.twist.linear.x)
            temp.append(p)
        return temp

    def pose_cb(self, msg):
        #rospy.loginfo('inside pose callback')
        self.pose = msg
        
    def waypoints_cb(self, waypoints):
        #rospy.loginfo('inside waypoint callback')
        self.base_Way_Points = waypoints 
        #rospy.loginfo(self.waypoints_2d)
        if not self.waypoints_2d:
            self.waypoints_2d = [[waypoint.pose.pose.position.x, waypoint.pose.pose.position.y] for waypoint in waypoints.waypoints]
            self.waypoint_tree = KDTree(self.waypoints_2d)
        
        #rospy.loginfo('waypoint tree')
        #rospy.loginfo(self.waypoint_tree)

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

    def obstacle_cb(self, msg):
        # TODO: Callback for /obstacle_waypoint 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.')
