#!/usr/bin/env python
from collections import defaultdict
import sys
import argparse
import osmnx as ox
import gzip
import subprocess
from time import perf_counter
from io import StringIO
import pandas as pd

epsilon = 0.01

plot_cutoffs = {
    "fast": 80.0,
    "medium": 30.0,
    "slow": 0
}

parser = argparse.ArgumentParser(
                    prog='run',
                    description='Run search from plain text waypoints',
                    epilog='Downloads data from OSM, may take some time!')
parser.add_argument('-s', '--source', help = "Source name, remember to enclose in double quotes!")
parser.add_argument('-g', '--goal', help = "Goal name, remember to enclose in double quotes!")
parser.add_argument('-m', "--map", help = "Map file")
parser.add_argument('-c', "--command", help = "command to run use double quotes if more than signle word")
parser.add_argument('-r', "--reference", help = "Reference command to check optimality", default = None)
parser.add_argument('-p', "--plotter", help="plotting mode fast (Region+) , medium (State), slow (city topological) or pretty (city with geometry)" , default = "medium")
parser.add_argument('-o', "--output", help = "plot output file", default="")

def distance(n1, n2):
    return ((n1.lat - n2.lat)**2 + (n1.lon - n2.lon)**2)**0.5

def check_edge_duration(target, edge_durations):
    mindif = float("inf")
    mindifval = None
    for dur in edge_durations:
        cand = abs(dur - target)
        if mindif > cand:
            mindif = cand
            mindifval = dur
    return mindifval
            

class Node:
    def __init__(self, lat, lon, index, name):
        self.lat = lat
        self.lon = lon
        self.index = index
        self.name = name 
    def __str__(self):
        return self.name
        #return str(self.lat) + " " + str(self.lon)
        

class Graph:
    def __init__(self, filename):
        self.nodes = []
        self.edges = defaultdict(set)
        self.edges_background = defaultdict(lambda: False)
        if(".gz" in filename):
            gfile = gzip.open(filename, "rt")
        else:
            gfile = open(filename)
        num_nodes = int(gfile.readline().split(": ")[1])
        for _ in range(num_nodes):
            line = gfile.readline()
            n = line.split()
            self.nodes.append(Node(float(n[2]), float(n[1]), int(n[0]), line))
        num_edges = int(gfile.readline().split(": ")[1])
        for _ in range(num_edges):
            words = gfile.readline().split()
            index = tuple(map(int,words[0:2]))
            self.edges[index].add(float(words[2]))
            self.edges_background[index] = (words[-1] == "b")
        gfile.close()
    
    def find_nearest(self, node):
        nearest = None
        for i in range(len(self.nodes)):
            n = self.nodes[i]
            if (nearest == None) or (distance(n, node) < distance(self.nodes[nearest], node)):
                nearest = i
        if(distance(node, self.nodes[nearest]) > epsilon):
            print("Warning: no close match found for ", node.lat, node.lon, self.nodes[nearest].lat, self.nodes[nearest].lon)
        return nearest
                

def get_osmids(filename):
    retval = {}
    if(".gz" in filename):
        gfile = gzip.open(filename, "rt")
    else:
        gfile = open(filename)
    num_nodes = int(gfile.readline().split(": ")[1])
    for _ in range(num_nodes):
        line = gfile.readline()
        n = line.split()
        retval[int(n[0])] = int(n[3])
    gfile.close()
    return retval

def check_valid_path(graph, results, start, goal):
    res = results.split("\n")[0:-1]
    if len(res) < 1:
        return True
    words = res[0].split()
    g = float(words[0])
    ind = int(words[1])
    lat = float(words[2])
    lon = float(words[3])
    if g > 0:
        print("Error path starts at g > 0")
        return False
    cur = Node(lat, lon, ind, "\t".join(words[1:])) 
    cur_nearest = cur.index
    if(cur_nearest != start.index):
        print("Error path does not start at start.")
        return False
    for i in range(1, len(res)):
        prev = cur_nearest
        prev_g = g
        words = res[i].split()
        g = float(words[0])
        ind = int(words[1])
        lat = float(words[2])
        lon = float(words[3])
        cur = Node(lat, lon, ind, "\t".join(words[1:])) 
        cur_nearest = cur.index
        index = (prev, cur_nearest)
        if index not in graph.edges:
            print("Edge not found: " + str(graph.nodes[prev]) + "->" + str(cur))
            return False
        elif (g - prev_g - check_edge_duration(g - prev_g, graph.edges[index])) > epsilon:
            print("Incorrect edge duration: " + str(graph.nodes[prev]) + "->" + str(cur))
            print("got: " + str(g - prev_g) + " expected: " + str(check_edge_duration(g - prev_g, graph.edges[index])))
            return False
    if(cur_nearest != goal.index):
        print("Error path does not end at goal.")
        return False
    return True
            

postscript_header = """%!PS

/m {newpath moveto} bind def
/l {lineto} bind def
/cp {closepath} bind def
/s {stroke} bind def
/sg {setgray} bind def
matrix currentmatrix /originmat exch def
/umatrix {originmat matrix concatmatrix setmatrix} def
[28.3465 0 0 28.3465 10.5 100.0] umatrix
0 sg
0.01 setlinewidth"""


postscript_line_template = """{:.2f} {:.2f} m
{:.2f} {:.2f} l
s"""

postscript_footer = "showpage"

def print_postscript(lats, lons, outfile, header = True, footer = True):
    y_min = lats.min()
    y_span = lats.max() - y_min

    x_min = lons.min()
    x_span = lons.max() - x_min

    scale = 20
    span = max(y_span, x_span)
    t_lats = scale*(lats - y_min)/span
    t_lons = scale*(lons - x_min)/span
   
    if header:
        print(postscript_header, file = outfile)
    for i in range(t_lats.shape[1]):
        print(postscript_line_template.format(t_lons[0, i], t_lats[0,i], t_lons[1, i], t_lats[1,i]), file = outfile)
    if footer:
        print(postscript_footer, file = outfile)
    return x_min, y_min, span

def print_path_postscript(lats, lons, outfile, x_min, y_min, span):
    scale = 20
    t_lats = scale*(lats - y_min)/span
    t_lons = scale*(lons - x_min)/span
    print("1 0 0 setrgbcolor", file = outfile)
    for i in range(1,t_lats.size):
        print(postscript_line_template.format(t_lons[i-1], t_lats[i-1], t_lons[i], t_lats[i]), file = outfile)

if __name__ == "__main__":
    args = parser.parse_args()
    source = ox.geocoder.geocode(args.source)
    dest = ox.geocoder.geocode(args.goal)
    command = args.command.split() + [
        "-y", str(source[0]),
        "-x", str(source[1]),
        "-Y", str(dest[0]),
        "-X", str(dest[1]),
        "-m", args.map
    ]
    
    if(".gz" in args.map):
        with gzip.open(args.map,"rb") as gzf:
            command[-1] = "-"
            start = perf_counter()
            res = subprocess.run(command, capture_output=True, check=True, input = gzf.read())
            end = perf_counter()
    else:
        start = perf_counter()
        res = subprocess.run(command, capture_output=True, check=True)
        end = perf_counter()
    print("Command: " + " ".join(command), file = sys.stderr)
    results = res.stdout.decode("ascii")
    print(results, end = "")
    print(res.stderr.decode("ascii"), end = "", file = sys.stderr)
    print("Runtime: " + str(end - start) + " seconds")
    
    print("Reading input graph: ", end = "")
    sys.stdout.flush()
    graph = Graph(args.map)
    print("Done")
    print("Finding nearest matching nodes to query: ", end = "")
    sys.stdout.flush()
    source_node = graph.find_nearest(Node(source[0], source[1], None, None))
    dest_node = graph.find_nearest(Node(dest[0], dest[1], None, None))
    print("Done")
    if(args.reference != None):
        print("Running reference solution")
        ref_command = args.reference.split() + [
            "-y", str(source[0]),
            "-x", str(source[1]),
            "-Y", str(dest[0]),
            "-X", str(dest[1]),
            "-m", args.map
        ]
        if(".gz" in args.map):
            with gzip.open(args.map,"rb") as gzf:
                ref_command[-1] = "-"
                ref_res = subprocess.run(ref_command, capture_output=True, check=True, input = gzf.read())
        else:
            ref_res = subprocess.run(ref_command, capture_output=True, check=True)
        ref_results = ref_res.stdout.decode("ascii")
        final_res = results.split("\n")[-2]
        makespan = float(final_res.split()[0])
        final_ref_res = ref_results.split("\n")[-2]
        ref_makespan = float(final_ref_res.split()[0])
        if( makespan > ref_makespan + epsilon):
            print("Optimal: False")
        else:
            print("Optimal: True")
    
    print("Path valid: ", end = "")
    sys.stdout.flush()
    print(str(check_valid_path(graph, results, graph.nodes[source_node], graph.nodes[dest_node])))
    if args.output != "":
        import matplotlib
        import matplotlib.style as mplstyle
        mplstyle.use('fast')
        import numpy as np
        matplotlib.use('PS')
        import matplotlib.pyplot as plt
        if args.plotter == "pretty":
            osmids = get_osmids(args.map)
            osmpath = []
            for res in results.split("\n"):
                if(res != ""):
                    words = res.split()
                    ind = int(words[1])
                    osmpath.append(osmids[ind])
            #g = ox.graph_from_place("Middleton New Hampshire", network_type='drive')
            north = max(source[0], dest[0])
            south = min(source[0], dest[0])
            west = min(source[1], dest[1])
            east = max(source[1], dest[1])
            span = max(abs(north - south), abs(east - west))
            g = ox.graph_from_bbox(north + span, south-span, east+span, west-span,   network_type = 'drive')
            ox.plot_graph_route(g, osmpath,show = False, save = True, filepath = args.output)
        elif "ps" in args.plotter:
            lats = np.empty((2, len(graph.edges)))
            lons = np.empty((2, len(graph.edges)))
            acc = 0
            check = "fast" in args.plotter
            print("Computing Background")
            for edge in graph.edges:
                if not check or graph.edges_background[edge]:
                    source = graph.nodes[edge[0]]
                    dest = graph.nodes[edge[1]]
                    lats[0,acc] = source.lat
                    lats[1,acc] = dest.lat
                    lons[0,acc] = source.lon
                    lons[1,acc] = dest.lon
                    acc += 1
            lats = lats[:,0:acc]
            lons = lons[:,0:acc]
            print("Plotting")
            with open(args.output, "wt") as outfile:
                x_min, y_min, span = print_postscript(lats, lons, outfile, footer = False)
                lats = []
                lons = []
                for res in results.split("\n"):
                    if(res != ""):
                        words = res.split()
                        lons.append(float(words[2]))
                        lats.append(float(words[3]))
                print_path_postscript(lats, lons, outfile, x_min, y_min, span)
                print(postscript_footer, file = outfile)
        else:
            print("Computing Background")
            cutoff = plot_cutoffs[args.plotter]
            durations = []
            lats = np.empty((2, len(graph.edges)))
            lons = np.empty((2, len(graph.edges)))
            acc = 0
            for edge in graph.edges:
                keep = False
                for x in graph.edges[edge]:
                    if x > cutoff and ((args.plotter != "fast") or (graph.edges_background[edge])):
                        keep = True
                if keep:
                    source = graph.nodes[edge[0]]
                    dest = graph.nodes[edge[1]]
                    lats[0,acc] = source.lat
                    lats[1,acc] = dest.lat
                    lons[0,acc] = source.lon
                    lons[1,acc] = dest.lon
                    acc += 1
            lats = lats[:,0:acc]
            lons = lons[:,0:acc]
            print("Plotting Background")
            plt.axis('equal')
            plt.plot(lons, lats, "-", color = "grey")
            print("Plotting Path")
            lats = []
            lons = []
            for res in results.split("\n"):
                if(res != ""):
                    words = res.split()
                    lons.append(float(words[2]))
                    lats.append(float(words[3]))
            plt.plot(lons, lats, marker = '.', color = 'r')
            plt.xlabel("Longitude")
            plt.ylabel("Latitude")
            plt.savefig(args.output)