#	Modifications and additions to code written by brooksandrew

import osmnx as ox
import networkx as nx
import pandas as pd
import matplotlib.pyplot as plt
import gpxpy
from collections import Counter

def circuit_path_string_to_int(circuit_rpp):
    """
    Converts nodes in path lists from strings to integers
    Args:
        circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp
    Returns:
        circuit_rpp (list): modified circuit
    """
    for e in circuit_rpp:
        if type(e[3]['path']) == str:
            exec('e[3]["path"]=' + e[3]["path"])
    return(circuit_rpp)

def create_req_and_opt_graph(req_comp_g_contracted, complete_g, circuit_rpp, GranularConnector_EdgeList):
    """
    Creates a graph with required and optional edges delineated for visualization
    Args:
        req_comp_g_contracted (NetworkX MultiDiGraph): required component graph containing contracted edges
            generated by initialize_rpp.InnerAndOuterToEdgeListFile
        complete_g (NetworkX MultiDiGraph): complete street network graph
            generated by initialize_rpp.InnerAndOuterToEdgeListFile
        circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp and
            edited by circuit_path_string_to_int
    Returns:
        final_graph (NetworkX DiGraph): graph of route with optional and required edges delineated
        print statements with required and optional edge breakdown
    """
    final_graph = req_comp_g_contracted.copy()
    unexpanded_edges = 0
    unexpanded_edges_list = []
    granular_connector_edges = 0
    granular_connector_edges_list = []
    granular_req_edges = 0
    granular_req_edges_list = []
    optional_edges = 0
    optional_edges_list = []
    for e in circuit_rpp:
        if [e[0], e[1]] not in unexpanded_edges_list:
            unexpanded_edges+=1
            unexpanded_edges_list+=[[e[0], e[1]]]
        # add granular optional edges to final_graph
        path = e[3]['path']
        for pair in list(zip(path[:-1], path[1:])):
            if (req_comp_g_contracted.has_edge(pair[0], pair[1])):
                edge = req_comp_g_contracted[pair[0]][pair[1]][0]
                if [pair[0], pair[1]] in GranularConnector_EdgeList:
                    final_graph[pair[0]][pair[1]][0]['granular_type'] = 'req street and connector'
                    if [pair[0], pair[1]] not in granular_connector_edges_list:
                        granular_connector_edges+=1
                        granular_connector_edges_list+=[[pair[0], pair[1]]]
                    elif 1 in req_comp_g_contracted[pair[0]][pair[1]]:
                            granular_connector_edges+=1
                else:
                    if [pair[0], pair[1]] not in granular_req_edges_list:
                        final_graph[pair[0]][pair[1]][0]['granular_type'] = 'req street'
                        granular_req_edges+=1
                        granular_req_edges_list+=[[pair[0], pair[1]]]
            else:
                if [pair[0], pair[1]] in GranularConnector_EdgeList:
                    if [pair[0], pair[1]] not in granular_connector_edges_list:
                        final_graph.add_edge(pair[0], pair[1], granular=True, granular_type='connector')
                        granular_connector_edges+=1
                        granular_connector_edges_list+=[[pair[0], pair[1]]]
                elif [pair[0], pair[1]] not in optional_edges_list:
                    final_graph.add_edge(pair[0], pair[1], granular=True, granular_type='optional')
                    optional_edges+=1
                    optional_edges_list+=[[pair[0], pair[1]]]
                else:
                    print(pair)
        for n in path:
            final_graph.add_node(n, y=complete_g.nodes[n]['y'], x=complete_g.nodes[n]['x'])
    print('Edges in Circuit')
    print('\tTotal Unexpanded Edges: {}'.format(unexpanded_edges))
    print('\tTotal Edges (All Contracted Edges Granularized): {}'.format(granular_connector_edges+granular_req_edges+optional_edges))
    print('\t\tGranular Connector Edges: {}'.format(granular_connector_edges))
    print('\t\tGranular Required Edges: {}'.format(granular_req_edges))
    print('\t\tGranular Optional Edges: {}'.format(optional_edges))
    return final_graph

def create_number_of_passes_graph(circuit_rpp, complete_g):
    """
    labels each edge with number of passes for visualization
    Args:
        circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp and
            edited by circuit_path_string_to_int
         complete_g (NetworkX MultiDiGraph): complete street network graph
            generated by initialize_rpp.InnerAndOuterToEdgeListFile
    Returns:
        grppviz(NetworkX Graph): route graph with number of pass attribute added to each edge
    """
    color_seq = [None, 'black', 'magenta', 'orange', 'yellow']
    grppviz = nx.Graph()
    for e in circuit_rpp:
        for n1, n2 in zip(e[3]['path'][:-1], e[3]['path'][1:]):
            if grppviz.has_edge(n1, n2):
                grppviz[n1][n2]['cnt'] += 1
                if grppviz[n1][n2]['cnt'] < 5:
                    grppviz[n1][n2]['linewidth'] += 2
            else:                
                grppviz.add_edge(n1, n2, linewidth=2.5)
                grppviz[n1][n2]['cnt'] = 1
                grppviz.add_node(n1, y=complete_g.nodes[n1]['y'], x=complete_g.nodes[n1]['x'])
                grppviz.add_node(n2, y=complete_g.nodes[n2]['y'], x=complete_g.nodes[n2]['x']) 
    for e in grppviz.edges(data=True):
        if e[2]['cnt'] < 5:
            e[2]['color_cnt'] = color_seq[e[2]['cnt']]
        else:
            e[2]['color_cnt'] = color_seq[4]
    return grppviz

def circuit_parser(circuit_rpp, complete_g, print_file, print_directory):
    """
    Creates dataframe of lat/lon coordinates corresponding to each node in the route
    Prints the dataframe to a csv file
    Args:
        circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp and
            edited by circuit_path_string_to_int
        complete_g (NetworkX MultiDiGraph): complete street network graph
            generated by initialize_rpp.InnerAndOuterToEdgeListFile
            lat/lon data contained her
        print_file (str): filename to print to
        print_directory (str): directory name for print_file
    Returns:
        rppdf (Pandas DataFrame): table of lat/lon coordinates of nodes in order they appear in route
        prints table to csv file
       
    """
    rpplist = []
    for ee in circuit_rpp:
        path = ee[3]['path'].copy()
        if ee == circuit_rpp[0]:
            if path[-1] in circuit_rpp[1][3]['path']:
                path = path
            else:
                path.reverse()
                path = path
            for n in path:
                rpplist.append({
                    'node' : n,
                    'lat' : complete_g.nodes[n]['y'],
                    'lon' : complete_g.nodes[n]['x']
                })
        else:
            if path[0] == rpplist[-1]['node']:
                path = path[1:]
            elif path[-1] == rpplist[-1]['node']:
                path.reverse()
                path = path[1:]
            for n in path:
                rpplist.append({
                    'node' : n,
                    'lat' : complete_g.nodes[n]['y'],
                    'lon' : complete_g.nodes[n]['x']
                })
    rppdf = pd.DataFrame(rpplist)
    rppdf.to_csv(print_directory + print_file + '.csv', index=False)
    print('\nRoute written to .csv as ' + print_directory + print_file + '.csv')
    return rppdf

def gpx_writer(rppdf, print_file, print_directory):
    """
    Writes route dataframe into gpx file
    Args:
        rppdf (Pandas DataFrame): route dataframe generated by circuit_parser
        print_file (str): filename to print to
        print_directory (str): directory name for print_file
    Returns:
        gpx file containing route
    """
    gpx = gpxpy.gpx.GPX()
    gpx_track = gpxpy.gpx.GPXTrack()
    gpx.tracks.append(gpx_track)
    gpx_segment = gpxpy.gpx.GPXTrackSegment()
    gpx_track.segments.append(gpx_segment)

    for index, row in rppdf.iterrows():
        gpx_segment.points.append(gpxpy.gpx.GPXTrackPoint(row['lat'], row['lon']))
    f = open(print_directory + print_file + '.gpx', 'w+')
    print(gpx.to_xml(), file = f)
    f.close()
    print('Route written to .gpx as ' + print_directory + print_file + '.gpx')

def plot_req_and_opt_graph(req_and_opt_graph):
    """
    Plots all edges of route and colors by required and optional edges
    Args:
        req_and_opt_graph (NetworkX Graph): graph of required and optional route edges
        generated by create_req_and_opt_graph
    Returns:
        plot of req_and_opt_graph
    """
    fig, ax = plt.subplots()

    pos = {k: (req_and_opt_graph.nodes[k].get('x'), req_and_opt_graph.nodes[k].get('y')) for k in req_and_opt_graph.nodes()}    

    el_opt = [e for e in req_and_opt_graph.edges(data=True) if e[2].get('granular_type') == 'optional'] 
    nx.draw_networkx_edges(req_and_opt_graph, pos, edgelist=el_opt, ax = ax, width=6.0, edge_color='blue', alpha=1.0)

    el_rd = [e for e in req_and_opt_graph.edges(data=True) if e[2].get('granular_type') in ['req street', 'connector', 'req street and connector']]
    nx.draw_networkx_edges(req_and_opt_graph, pos, edgelist=el_rd, ax = ax, width=3.0, edge_color='black', alpha=0.8)

    ax.set_xlim(left = min([pos[node][0] for node in req_and_opt_graph.nodes()]), right = max([pos[node][0] for node in req_and_opt_graph.nodes()]))
    ax.set_ylim(bottom = min([pos[node][1] for node in req_and_opt_graph.nodes()]), top = max([pos[node][1] for node in req_and_opt_graph.nodes()]))

    print('\nRequired and Optional Edge Graph\nRequired Edges in Black\nOptional Edges in Blue')

    plt.show()

def plot_number_of_passes_graph(grppviz):
    """
    Plots route graph with color and thickness of edge determined by number of passes
    Args:
        grppviz (NetworkX Graph): route graph with number of pass data
        generated by create_number_of_passes_graph
    Returns:
        Plot of route with number of passes depicted by color and edgewidth
    """
    fig, ax = plt.subplots()

    pos = {k: (grppviz.nodes[k]['x'], grppviz.nodes[k]['y']) for k in grppviz.nodes()}    

    e_width = [e[2]['linewidth'] for e in grppviz.edges(data=True)]
    e_color = [e[2]['color_cnt'] for e in grppviz.edges(data=True)]

    nx.draw_networkx_edges(grppviz, pos, ax = ax, width=e_width, edge_color=e_color, alpha=0.7)

    ax.set_xlim(left = min([pos[node][0] for node in grppviz.nodes()]), right = max([pos[node][0] for node in grppviz.nodes()]))
    ax.set_ylim(bottom = min([pos[node][1] for node in grppviz.nodes()]), top = max([pos[node][1] for node in grppviz.nodes()]))

    print('\nNumber of Passes Graph\nBlack\t1 Pass\nMagenta\t2 Passes\nOrange\t3 Passes\nYellow\t4 or More Passes')
    plt.show()


#	Modifications and additions to code written by brooksandrew

import osmnx as ox
import networkx as nx
import pandas as pd
import matplotlib.pyplot as plt
import gpxpy
from collections import Counter

def circuit_path_string_to_int(circuit_rpp):
    """
    Converts nodes in path lists from strings to integers
    Args:
        circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp
    Returns:
        circuit_rpp (list): modified circuit
    """
    for e in circuit_rpp:
        if type(e[3]['path']) == str:
            exec('e[3]["path"]=' + e[3]["path"])
    return(circuit_rpp)

def create_req_and_opt_graph(req_comp_g_contracted, complete_g, circuit_rpp, GranularConnector_EdgeList):
    """
    Creates a graph with required and optional edges delineated for visualization
    Args:
        req_comp_g_contracted (NetworkX MultiDiGraph): required component graph containing contracted edges
            generated by initialize_rpp.InnerAndOuterToEdgeListFile
        complete_g (NetworkX MultiDiGraph): complete street network graph
            generated by initialize_rpp.InnerAndOuterToEdgeListFile
        circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp and
            edited by circuit_path_string_to_int
    Returns:
        final_graph (NetworkX DiGraph): graph of route with optional and required edges delineated
        print statements with required and optional edge breakdown
    """
    final_graph = req_comp_g_contracted.copy()
    unexpanded_edges = 0
    unexpanded_edges_list = []
    granular_connector_edges = 0
    granular_connector_edges_list = []
    granular_req_edges = 0
    granular_req_edges_list = []
    optional_edges = 0
    optional_edges_list = []
    for e in circuit_rpp:
        if [e[0], e[1]] not in unexpanded_edges_list:
            unexpanded_edges+=1
            unexpanded_edges_list+=[[e[0], e[1]]]
        # add granular optional edges to final_graph
        path = e[3]['path']
        for pair in list(zip(path[:-1], path[1:])):
            if (req_comp_g_contracted.has_edge(pair[0], pair[1])):
                edge = req_comp_g_contracted[pair[0]][pair[1]][0]
                if [pair[0], pair[1]] in GranularConnector_EdgeList:
                    final_graph[pair[0]][pair[1]][0]['granular_type'] = 'req street and connector'
                    if [pair[0], pair[1]] not in granular_connector_edges_list:
                        granular_connector_edges+=1
                        granular_connector_edges_list+=[[pair[0], pair[1]]]
                    elif 1 in req_comp_g_contracted[pair[0]][pair[1]]:
                            granular_connector_edges+=1
                else:
                    if [pair[0], pair[1]] not in granular_req_edges_list:
                        final_graph[pair[0]][pair[1]][0]['granular_type'] = 'req street'
                        granular_req_edges+=1
                        granular_req_edges_list+=[[pair[0], pair[1]]]
            else:
                if [pair[0], pair[1]] in GranularConnector_EdgeList:
                    if [pair[0], pair[1]] not in granular_connector_edges_list:
                        final_graph.add_edge(pair[0], pair[1], granular=True, granular_type='connector')
                        granular_connector_edges+=1
                        granular_connector_edges_list+=[[pair[0], pair[1]]]
                elif [pair[0], pair[1]] not in optional_edges_list:
                    final_graph.add_edge(pair[0], pair[1], granular=True, granular_type='optional')
                    optional_edges+=1
                    optional_edges_list+=[[pair[0], pair[1]]]
                else:
                    print(pair)
        for n in path:
            final_graph.add_node(n, y=complete_g.nodes[n]['y'], x=complete_g.nodes[n]['x'])
    print('Edges in Circuit')
    print('\tTotal Unexpanded Edges: {}'.format(unexpanded_edges))
    print('\tTotal Edges (All Contracted Edges Granularized): {}'.format(granular_connector_edges+granular_req_edges+optional_edges))
    print('\t\tGranular Connector Edges: {}'.format(granular_connector_edges))
    print('\t\tGranular Required Edges: {}'.format(granular_req_edges))
    print('\t\tGranular Optional Edges: {}'.format(optional_edges))
    return final_graph

def create_number_of_passes_graph(circuit_rpp, complete_g):
    """
    labels each edge with number of passes for visualization
    Args:
        circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp and
            edited by circuit_path_string_to_int
         complete_g (NetworkX MultiDiGraph): complete street network graph
            generated by initialize_rpp.InnerAndOuterToEdgeListFile
    Returns:
        grppviz(NetworkX Graph): route graph with number of pass attribute added to each edge
    """
    color_seq = [None, 'black', 'magenta', 'orange', 'yellow']
    grppviz = nx.Graph()
    for e in circuit_rpp:
        for n1, n2 in zip(e[3]['path'][:-1], e[3]['path'][1:]):
            if grppviz.has_edge(n1, n2):
                grppviz[n1][n2]['cnt'] += 1
                if grppviz[n1][n2]['cnt'] < 5:
                    grppviz[n1][n2]['linewidth'] += 2
            else:                
                grppviz.add_edge(n1, n2, linewidth=2.5)
                grppviz[n1][n2]['cnt'] = 1
                grppviz.add_node(n1, y=complete_g.nodes[n1]['y'], x=complete_g.nodes[n1]['x'])
                grppviz.add_node(n2, y=complete_g.nodes[n2]['y'], x=complete_g.nodes[n2]['x']) 
    for e in grppviz.edges(data=True):
        if e[2]['cnt'] < 5:
            e[2]['color_cnt'] = color_seq[e[2]['cnt']]
        else:
            e[2]['color_cnt'] = color_seq[4]
    return grppviz

def circuit_parser(circuit_rpp, complete_g, print_file, print_directory):
    """
    Creates dataframe of lat/lon coordinates corresponding to each node in the route
    Prints the dataframe to a csv file
    Args:
        circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp and
            edited by circuit_path_string_to_int
        complete_g (NetworkX MultiDiGraph): complete street network graph
            generated by initialize_rpp.InnerAndOuterToEdgeListFile
            lat/lon data contained her
        print_file (str): filename to print to
        print_directory (str): directory name for print_file
    Returns:
        rppdf (Pandas DataFrame): table of lat/lon coordinates of nodes in order they appear in route
        prints table to csv file
       
    """
    rpplist = []
    for ee in circuit_rpp:
        path = ee[3]['path'].copy()
        if ee == circuit_rpp[0]:
            if path[-1] in circuit_rpp[1][3]['path']:
                path = path
            else:
                path.reverse()
                path = path
            for n in path:
                rpplist.append({
                    'node' : n,
                    'lat' : complete_g.nodes[n]['y'],
                    'lon' : complete_g.nodes[n]['x']
                })
        else:
            if path[0] == rpplist[-1]['node']:
                path = path[1:]
            elif path[-1] == rpplist[-1]['node']:
                path.reverse()
                path = path[1:]
            for n in path:
                rpplist.append({
                    'node' : n,
                    'lat' : complete_g.nodes[n]['y'],
                    'lon' : complete_g.nodes[n]['x']
                })
    rppdf = pd.DataFrame(rpplist)
    rppdf.to_csv(print_directory + print_file + '.csv', index=False)
    print('\nRoute written to .csv as ' + print_directory + print_file + '.csv')
    return rppdf

def gpx_writer(rppdf, print_file, print_directory):
    """
    Writes route dataframe into gpx file
    Args:
        rppdf (Pandas DataFrame): route dataframe generated by circuit_parser
        print_file (str): filename to print to
        print_directory (str): directory name for print_file
    Returns:
        gpx file containing route
    """
    gpx = gpxpy.gpx.GPX()
    gpx_track = gpxpy.gpx.GPXTrack()
    gpx.tracks.append(gpx_track)
    gpx_segment = gpxpy.gpx.GPXTrackSegment()
    gpx_track.segments.append(gpx_segment)

    for index, row in rppdf.iterrows():
        gpx_segment.points.append(gpxpy.gpx.GPXTrackPoint(row['lat'], row['lon']))
    f = open(print_directory + print_file + '.gpx', 'w+')
    print(gpx.to_xml(), file = f)
    f.close()
    print('Route written to .gpx as ' + print_directory + print_file + '.gpx')

def plot_req_and_opt_graph(req_and_opt_graph):
    """
    Plots all edges of route and colors by required and optional edges
    Args:
        req_and_opt_graph (NetworkX Graph): graph of required and optional route edges
        generated by create_req_and_opt_graph
    Returns:
        plot of req_and_opt_graph
    """
    fig, ax = plt.subplots()

    pos = {k: (req_and_opt_graph.nodes[k].get('x'), req_and_opt_graph.nodes[k].get('y')) for k in req_and_opt_graph.nodes()}    

    el_opt = [e for e in req_and_opt_graph.edges(data=True) if e[2].get('granular_type') == 'optional'] 
    nx.draw_networkx_edges(req_and_opt_graph, pos, edgelist=el_opt, ax = ax, width=6.0, edge_color='blue', alpha=1.0)

    el_rd = [e for e in req_and_opt_graph.edges(data=True) if e[2].get('granular_type') in ['req street', 'connector', 'req street and connector']]
    nx.draw_networkx_edges(req_and_opt_graph, pos, edgelist=el_rd, ax = ax, width=3.0, edge_color='black', alpha=0.8)

    ax.set_xlim(left = min([pos[node][0] for node in req_and_opt_graph.nodes()]), right = max([pos[node][0] for node in req_and_opt_graph.nodes()]))
    ax.set_ylim(bottom = min([pos[node][1] for node in req_and_opt_graph.nodes()]), top = max([pos[node][1] for node in req_and_opt_graph.nodes()]))

    print('\nRequired and Optional Edge Graph\nRequired Edges in Black\nOptional Edges in Blue')

    plt.show()

def plot_number_of_passes_graph(grppviz):
    """
    Plots route graph with color and thickness of edge determined by number of passes
    Args:
        grppviz (NetworkX Graph): route graph with number of pass data
        generated by create_number_of_passes_graph
    Returns:
        Plot of route with number of passes depicted by color and edgewidth
    """
    fig, ax = plt.subplots()

    pos = {k: (grppviz.nodes[k]['x'], grppviz.nodes[k]['y']) for k in grppviz.nodes()}    

    e_width = [e[2]['linewidth'] for e in grppviz.edges(data=True)]
    e_color = [e[2]['color_cnt'] for e in grppviz.edges(data=True)]

    nx.draw_networkx_edges(grppviz, pos, ax = ax, width=e_width, edge_color=e_color, alpha=0.7)

    ax.set_xlim(left = min([pos[node][0] for node in grppviz.nodes()]), right = max([pos[node][0] for node in grppviz.nodes()]))
    ax.set_ylim(bottom = min([pos[node][1] for node in grppviz.nodes()]), top = max([pos[node][1] for node in grppviz.nodes()]))

    print('\nNumber of Passes Graph\nBlack\t1 Pass\nMagenta\t2 Passes\nOrange\t3 Passes\nYellow\t4 or More Passes')
    plt.show()


#	Modifications and additions to code written by brooksandrew Converts nodes in path lists from strings to integers     Args:         circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp     Returns:         circuit_rpp (list): modified circuit Creates a graph with required and optional edges delineated for visualization     Args:         req_comp_g_contracted (NetworkX MultiDiGraph): required component graph containing contracted edges             generated by initialize_rpp.InnerAndOuterToEdgeListFile         complete_g (NetworkX MultiDiGraph): complete street network graph             generated by initialize_rpp.InnerAndOuterToEdgeListFile         circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp and             edited by circuit_path_string_to_int     Returns:         final_graph (NetworkX DiGraph): graph of route with optional and required edges delineated         print statements with required and optional edge breakdown # add granular optional edges to final_graph labels each edge with number of passes for visualization     Args:         circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp and             edited by circuit_path_string_to_int          complete_g (NetworkX MultiDiGraph): complete street network graph             generated by initialize_rpp.InnerAndOuterToEdgeListFile     Returns:         grppviz(NetworkX Graph): route graph with number of pass attribute added to each edge Creates dataframe of lat/lon coordinates corresponding to each node in the route     Prints the dataframe to a csv file     Args:         circuit_rpp (list): rpp circuit generated by postman_problems.solver.rpp and             edited by circuit_path_string_to_int         complete_g (NetworkX MultiDiGraph): complete street network graph             generated by initialize_rpp.InnerAndOuterToEdgeListFile             lat/lon data contained her         print_file (str): filename to print to         print_directory (str): directory name for print_file     Returns:         rppdf (Pandas DataFrame): table of lat/lon coordinates of nodes in order they appear in route         prints table to csv file Writes route dataframe into gpx file     Args:         rppdf (Pandas DataFrame): route dataframe generated by circuit_parser         print_file (str): filename to print to         print_directory (str): directory name for print_file     Returns:         gpx file containing route Plots all edges of route and colors by required and optional edges     Args:         req_and_opt_graph (NetworkX Graph): graph of required and optional route edges         generated by create_req_and_opt_graph     Returns:         plot of req_and_opt_graph Plots route graph with color and thickness of edge determined by number of passes     Args:         grppviz (NetworkX Graph): route graph with number of pass data         generated by create_number_of_passes_graph     Returns:         Plot of route with number of passes depicted by color and edgewidth