import time
import re
import numpy as np
import networkx as nx
# import matplotlib.pyplot as plt

from shapely.geometry import Polygon, Point, LineString
from sklearn.neighbors import KDTree
from planning_utils import GpsLocation, WorldMap, Plot
from scipy.spatial import Voronoi, voronoi_plot_2d
from bresenham import bresenham
from queue import PriorityQueue
from udacidrone.frame_utils import global_to_local
from action_based_plan import ActionPlanner
from PIL import Image, ImageDraw

def distance_heuristic(n1, n2):
    return np.linalg.norm((n1, n2))

class Planner():
    def __init__(self):
        self.init_gps = GpsLocation(0,0,0)
        self.home_gps_pos = GpsLocation(0,0,0)
        self.safety_distance = 3
        self.worldmap = WorldMap('colliders.csv', self.safety_distance)
        self.min_drone_altitude = 5
        self.max_drone_altitude = 20
        self.graph_nodes_kd = None      # KDTree for all graph nodes generated by voronoi. Used for finding closest to start and goal
        self.graph = None               # nx.Graph
        self.start3d = (0, 0, 0)
        self.goal3d = (0, 0, 0)
        self.plot = Plot()
        self.obstacle_tree = None

    # @property
    # def north_min(self):
    #     return self.worldmap.north_min_max[0]

    # @property
    # def north_max(self):
    #     return self.worldmap.north_min_max[1]

    # @property
    # def east_min(self):
    #     return self.worldmap.east_min_max[0]

    # @property
    # def east_max(self):
    #     return self.worldmap.east_min_max[1]
    
    # @property
    # def data(self):
    #     return self.worldmap.data

    @property
    def grid25(self):
        return self.worldmap.grid25

    def set_initial_gps(self, pos):
        self.init_gps = pos

    def load_map(self):
        self.worldmap.load()
        self.home_gps_pos = self.worldmap.home_gps_pos
        self.create_obstacle_tree()

        print("[Planner] Home Gps Location: {}".format(self.home_gps_pos))

    def create_obstacle_tree(self):
        data = self.worldmap.data

        # figure out which possible obstacles might we collide with at the minimum altitude
        height = np.array(np.ceil(data[:, 2] + data[:, 5]))
        collision_index = height > self.min_drone_altitude - self.safety_distance

        # all obstacles that we can collide with
        obstacles = data[collision_index]
        
        centers = np.dstack((obstacles[:, 0], obstacles[:, 1])).squeeze()
        centers = centers.astype(int)

        self.obstacle_tree = KDTree(centers)

    def plan_route(self, start3d, goal3d):
        self.start3d = start3d
        self.goal3d = goal3d

        # Create voronoi graph to create a 2d path
        self.graph, self.graph_nodes_kd = self.create_voronoi_graph()  # graph and its KD

        # find the closest location on the graph that matches the start and the goal states
        indices = self.graph_nodes_kd.query([start3d, goal3d], 
                                k = 1, 
                                return_distance = False)

        # print('closest to goal indices:', indices)
        # print('index 0:', indices[0])
        # print('index 1:', indices[1])

        # find the closest node to the start and goal states

        graph_list = list(self.graph)

        self.c_start = graph_list[indices[0][0]]
        self.c_goal = graph_list[indices[1][0]]

        print('Closest To Start {} is {}'.format(start3d, self.c_start))
        print('Closest To Goal {} is {} '.format(goal3d, self.c_goal))

        print('Running a_star')
        path = self.a_star_graph(self.graph, distance_heuristic, self.c_start, self.c_goal)
        
        # prune path
        print('Pruning path...')
        
        pruned_path = self.prune_path(path)
        landing_starts = len(pruned_path)               # this is returned to the motion_planner to decrease WAYPOINT_HIT_RADIUS

        print('Path pruned. Length: ', len(pruned_path))

        # use action based planner in the end to reach the goal state, specially the
        # height of the goal

        close = np.abs(goal3d - self.c_goal) < 0.1
        if not close.all():
            print("Since closest goal is not on the graph, using action based planner for last node...")
            print("Original goal: {}, closest: {}".format(goal3d, self.c_goal))

            ap = ActionPlanner(self.grid25, self.max_drone_altitude)
            path = ap.find_path(tuple(self.c_goal), tuple(self.goal3d))

            if len(path) > 0:
                # print('adding action based plan at the end of the pruned path')
                pruned_action_path = self.prune_path(path)
                print('Action Path: ', pruned_action_path)

                # print('Before adding pruned path:', pruned_path)

                pruned_path.extend(pruned_action_path)

                # print('-' * 100)
                # print(pruned_path)

        return pruned_path, landing_starts


    def create_grid(self):
        grid25 = self.grid25
        return grid25.create_binary(self.min_drone_altitude)

    def convert_edges_tomap(self, edges):
        n_min = self.worldmap.north_min
        e_min = self.worldmap.east_min

        # Voronoi edges are 0 based as per the grid that was given to it
        # we need to convert it back to world coordinates in the map
        # edges_w = []
        # for p1, p2 in edges:
        #     p1_w = (p1[0] + n_min, p1[1] + e_min)
        #     p2_w = (p2[0] + n_min, p2[1] + e_min)
        #     edges_w.append((p1_w, p2_w))
        
        # make a copy of edges
        edges_w = np.array(edges)

        edges_w[:, 0] += n_min
        edges_w[:, 1] += e_min
        edges_w[:, 2] += n_min
        edges_w[:, 3] += e_min

        return edges_w

    def create_voronoi_graph(self):
        edges = self.get_voronoi_edges()

        # conver the edges back to the map coordinate space (<0 based)
        edges_w = self.convert_edges_tomap(edges)

        graph = nx.Graph()

        # add altitude to each 2d edge returned by voronoi edges
        # for p1, p2 in edges_w:
        #     p1_3d = (p1[0], p1[1], -self.min_drone_altitude)
        #     p2_3d = (p2[0], p2[1], -self.min_drone_altitude)
        #     weight = np.sum((np.array(p2_3d) - np.array(p1_3d)) ** 2) ** 0.5
            
        for p1_north, p1_east, p2_north, p2_east in edges_w:
            p1_3d = (p1_north, p1_east, -self.min_drone_altitude)
            p2_3d = (p2_north, p2_east, -self.min_drone_altitude)
            weight = np.sum((np.array(p2_3d) - np.array(p1_3d)) ** 2) ** 0.5

            graph.add_edge(p1_3d, p2_3d, weight=weight)

            #print('vor: distance between {} and {} = {}'.format(p1_3d, p2_3d, weight))

        # self.plot_vedges(edges)

        graph_nodes_kd = KDTree(graph.nodes)
        return (graph, graph_nodes_kd)

    # def plot_vedges(self, edges):
    #     grid = self.create_grid()
    #     im = Image.fromarray(np.uint8(grid) * 255)
    #     # im = Image.fromarray(np.uint8(np.flip(grid, 0) * 255))
        
    #     draw = ImageDraw.Draw(im)
    #     for y1, x1, y2, x2 in edges:
    #         draw.line((x1, y1, x2, y2), width = 2)

    #     ima = np.array(im)
    #     self.plot.show_image(ima)

    #     plt.rcParams["figure.figsize"] = [12, 12]

    #     fig = plt.figure()
    #     plt.imshow(grid, cmap='Greys', origin='lower')

    #     for y1, x1, y2, x2 in edges:
    #         plt.plot([x1, x2], [y1, y2])
        
    #     plt.show()

    def generate_3dpath(self, local_position, radius = 25):
        # draw random sample of points in 3d
        # find the farthest point in the 2d path that is in radius of given 3d path
        #   and consider that as goal
        # start location is the current location of the drone
        # add all random points in a graph
        # look for a 3d path through the samples using the 2.5d grid as basis for collision
        # in case the cost of following the 2d path is less than the cost for the 3d
        # keep on following 2d

        grid25 = self.grid25
        num_samples = 20

        north_size = int(np.ceil(grid25.north_max - grid25.north_max))
        east_size = int(np.ceil(grid25.east_max - grid25.east_min))

        radius = 25
        n_min, n_max = local_position[0] - radius, local_position[0] + radius
        e_min, e_max = local_position[1] - radius, local_position[1] + radius

        if n_min < 0:
            n_min = 0
        if e_min < 0:
            e_min = 0
        if n_max > north_size:
            n_max = north_size
        if e_max > east_size:
            e_max = east_size

        nvals = np.random.uniform(n_min, n_max, num_samples).astype(int)
        evals = np.random.uniform(e_min, e_max, num_samples).astype(int)
        zvals = np.random.uniform(self.min_drone_altitude, self.max_drone_altitude, num_samples).astype(int)

        samples = list(zip(nvals, evals, zvals))

        t0 = time.time()
        to_keep = []

        grid25 = self.planner.grid25

        for p in samples:
            gp = grid25[p[0], p[1]]
            if gp == 0 and gp < self.max_altitude:
                to_keep.append(p)
                    
        time_taken = time.time() - t0
        print("Time taken for sample point collision: {0} seconds ...".format(time_taken))
        print(len(to_keep))

    # def point(self, p):
    #     return np.array([p[0], p[1], 1.]).reshape(1, -1)

    def collinearity_check(self, p1, p2, p3, epsilon=1e-6):   
        m = np.vstack((p1, p2, p3))
        det = np.linalg.det(m)
        return abs(det) < epsilon

    def a_star_graph(self, graph, heuristic, start, goal):
        """Modified A* to work with NetworkX graphs."""
        path = []
        path_cost = 0

        queue = PriorityQueue()
        queue.put((0, start))
        visited = set(start)

        branch = {}
        found = False
        
        while not queue.empty():
            item = queue.get()
            current_node = item[1]
            
            if current_node == start:
                current_cost = 0.0
            else:              
                current_cost = branch[current_node][0]
                
            if current_node == goal:        
                print('Found a path.')
                found = True
                break
            else:
                #for action in valid_actions(grid, current_node):
                graph_node = graph[current_node]
                
                for next_node in graph_node:
                    if next_node not in visited:                
                        weight = graph_node[next_node]['weight']

                        branch_cost = current_cost + weight
                        queue_cost = branch_cost + heuristic(next_node, goal)
                    
                        visited.add(next_node)               
                        branch[next_node] = (branch_cost, current_node)
                        queue.put((queue_cost, next_node))
                
        if found:
            # retrace steps. Each node will indicate the cost of 
            # reaching the goal from that node
            n = goal
            path_cost = 0
            path.append(goal + (path_cost, ))

            while n != start:
                # Figure out the cost of reaching the Nth node from its parent.
                # Add the cost of reaching goal from the Nth node to the cost of reaching
                #  nth from its parent

                parent_node = branch[n][1]

                # get cost of reaching Nth node from its parent
                graph_node = graph[parent_node]
                weight = graph_node[n]['weight']

                # add the cost of reaching goal from nth to its parent's cost of reaching
                # to the nth
                path_cost += weight

                n = parent_node
                path.append((n[0], n[1], n[2], path_cost))
        else:
            print('**********************')
            print('Failed to find a path!')
            print('**********************') 

        return path[::-1]

    def prune_path(self, pp):
        if pp is None or len(pp) < 4:
            return pp

        pruned_path = []
        
        p1 = pp[0]
        p2 = pp[1]

        pruned_path.append(p1)

        for i in range(2, len(pp)):
            p3 = pp[i]

            p1_ned = (p1[0], p1[1], p1[2])
            p2_ned = (p2[0], p2[1], p2[2])
            p3_ned = (p3[0], p3[1], p3[2])

            if self.collinearity_check(p1_ned, p2_ned, p3_ned):
                #print("Colinear (therefore skipping): ", p1_ned, p2_ned, p3_ned)
                p2 = p3
            else:
                pruned_path.append(p2)

                p1 = p2
                p2 = p3
        
        pruned_path.append(p3)
        # delete the 3rd dimension
        #return np.delete(np.array(pruned_path), 2, axis=1)
        return pruned_path

    def generate_receding_plan(self):
        """This generates a plan using 2d map"""
        pass

    def get_voronoi_edges(self):
        """returns voronoi edges in the image coordinates"""
        north_min = self.worldmap.north_min
        east_min = self.worldmap.east_min

        # Make all data points 0 based instead of from -316  as voronoi can't work on < 0 indices
        north = self.worldmap.data[:, 0] - north_min
        east = self.worldmap.data[:, 1] - east_min
        
        points = np.dstack((north, east)).squeeze()
        
        # create a voronoi graph out of the center points of each obstacle given in the map
        vgraph = Voronoi(points)
        vertices = vgraph.vertices.astype(int)

        # figure out the vertices that are out of bounds (< 0) and remove them from
        # computation by setting them to -1, -1
        nob = vertices[:, 0] < 0
        eob = vertices[:, 1] < 0
        out_of_bounds = np.where(nob | eob)
        vertices[out_of_bounds] = np.array([-1, -1])
        
        grid25 = self.worldmap.grid25

        # set all vertices that are > grid max to -1,-1
        nob = vertices[:, 0] >= grid25.shape[0]
        eob = vertices[:, 1] >= grid25.shape[1]
        out_of_bounds = np.where(nob | eob)
        vertices[out_of_bounds] = np.array([-1, -1])
        
        # create a binary grid given the drone_altitude
        grid = grid25.create_binary(self.min_drone_altitude)
        edges = []
        
        for e in vgraph.ridge_vertices:
            p1 = vertices[e[0]]
            p2 = vertices[e[1]]
            
            collision = False

            # do not consider a vertex that is out of bounds
            if p1[0] == -1 or p2[0] == -1:
                collision = True
            else:
                # check in case there is a collision between the two
                # points given by the vertices in the graph
                cells = bresenham(p1[0], p1[1], p2[0], p2[1])
                for c in cells:
                    if grid[c[0], c[1]] == 1:
                        collision = True
                        break

            if not collision:
                #edges.append((p1, p2))
                edges.append([p1[0], p1[1], p2[0], p2[1]])
        
        return np.array(edges)


if __name__ == "__main__":
    import time
    from planning_utils import Plot
    from PIL import Image
    import matplotlib.pyplot as plt

    planner = Planner()
    planner.load_map()

    print("map loaded")

    def test_image(grid):
        im = Image.fromarray(np.uint8(np.flip(grid, 0) * 255))
        im = Image.open('test.png')
        print('Grid has been saved to test.png')

    def test_grid():
        p = Planner()
        
        print('Loading map...')
        p.load_map()

        print('Creating grid...')
        grid = p.create_grid()

        print('Plotting grid...')
        fig = plt.figure()
        plt.imshow(grid, origin='lower', cmap='Greys') 
        plt.show()

    def test_voronoi():
        print('test_voronoi')
        
        grid = planner.create_grid()
        edges = planner.get_voronoi_edges()

        print('Voronoi edges computed. Length: ', len(edges))

        n_min = planner.worldmap.north_min
        e_min = planner.worldmap.east_min
        
        fig = plt.figure()
        plt.subplot(121)
        plt.imshow(grid, origin='lower', cmap='Greys') 

        print('Plotting points...')

        for p1, p2 in edges:
            plt.plot([p1[1], p2[1]], [p1[0], p2[0]], 'b-')
    
        print('Converting into map coordinates..')
        edges_w = planner.convert_edges_tomap(edges)

        print('Plotting map points...')

        plt.subplot(122)
        plt.imshow(grid, origin='lower', cmap='Greys') 
        
        for p1, p2 in edges_w:
            p1_i = p1[0] - n_min, p1[1] - e_min
            p2_i = p2[0] - n_min, p2[1] - e_min
            plt.plot([p1_i[1], p2_i[1]], [p1_i[0], p2_i[0]], color='red')
        
        plt.show()

    def test_prune():
        print('test_prune')
        lon0 = -122.397450
        lat0 = 37.792480

        global_home = (lon0, lat0, 0)

        start = (0,0,0)
        goal_gps = GpsLocation(37.794948, -122.396666, 0)
        goal = global_to_local(goal_gps, global_home)

        print('Start: ', start)
        print('Goal: ', goal)

        planner.load_map()
        grid = planner.create_grid()

        s = time.time()
        path2d = planner.plan_route(start, goal)
        e = time.time()
        print("Pruned path. Cost: {}, time taken: {}".format(path2d[0][3], e - s))

    test_prune()

    # --- show grid
    #p = Plot()
    #p.show_grid(grid, planner.min_drone_altitude)


    # -- 3d graph print check
    #grid3d = planner.worldmap.create_grid3d(5)
    #print(grid3d.shape)

    # fig = plt.figure()
    # ax = fig.gca(projection='3d')
    # ax.voxels(grid3d, edgecolor='k')
    # ax.set_xlim(grid3d.shape[0], 0)
    # ax.set_ylim(0, grid3d.shape[1])
    # # add 100 to the height so the buildings aren't so tall
    # ax.set_zlim(0, grid3d.shape[2]+20)

    # plt.xlabel('North')
    # plt.ylabel('East')

    # plt.show()