generator.py

import simplejson as json
import argparse
import signal
import sys
import os

from rich.table import Table
from rich.markup import escape
from rich.console import Console
from rich.progress  import Progress, BarColumn, TextColumn
console = Console()

parser = argparse.ArgumentParser(description="Webduino source builder")

parser.add_argument("-v", "--verbose", action='store_true', dest='verbose',
                    help="Displays verbose output")

parser.add_argument("-o", "--output", metavar="file", type=str,
                    default="./stations.json", dest='output',
                    help="location of the output file (default: ./stations.json)")

parser.add_argument("-m", "--matrix", action='store_true',
                    dest='matrix', default=False,
                    help="Create transition Matrices")

args = parser.parse_args()

def warn(text):
    """Creates warning in stdout"""

    console.print("[bold yellow]WARNING:", "[bold white]" + text)

def query_yes_no(question, default="yes"):
    """Queries user for confimration"""

    valid = {"yes": True, "y": True, "ye": True,
             "no": False, "n": False}
    if default is None:
        prompt = " [y/n] "
    elif default == "yes":
        prompt = " [Y/n] "
    elif default == "no":
        prompt = " [y/N] "
    else:
        raise ValueError("invalid default answer: '%s'" % default)

    while True:
        console.print(question + escape(prompt))
        choice = input().lower()
        if default is not None and choice == '':
            return valid[default]
        elif choice in valid:
            return valid[choice]
        else:
            console.print("Please respond with 'yes' or 'no' "
                             "(or 'y' or 'n').\n")

# Dump matricies
matrix_dist_file = "./matrix_dist.csv"
matrix_time_file = "./matrix_time.csv"

ask_continue = False

# Ask for user confirmation if files exist
if os.path.exists(args.output):
    warn(args.output + " will be overridden!")
    ask_continue = True
if args.matrix and os.path.exists(matrix_dist_file):
    warn(matrix_dist_file + " will be overridden!")
    ask_continue = True
if args.matrix and os.path.exists(matrix_time_file):
    warn(matrix_time_file + " will be overridden!")
    ask_continue = True

if ask_continue:
    try:
        if not query_yes_no("Would you like to continue?"):
            sys.exit(0)
    except KeyboardInterrupt:
        sys.exit(0)

# Abort handler
def signal_handler(sig, frame):
    print('Aborting!')
    sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)

# Generate Data
with Progress("[progress.percentage]{task.percentage:>3.0f}%", BarColumn(), "[progress.description]{task.description}") as progress:
    task1 = progress.add_task("Generating Data", total=1)

    step = 0
    # Provide funtion to update progress bar
    def update_progress(text):
        global step
        progress.tasks[task1].description = text
        progress.update(task1, completed=(step / (9 if args.matrix else 8)))
        progress.refresh()
        step = step + 1

    # Parse CSV file and return a list of lists
    def parse_csv(text, sep=",", name="unknown file"):
        lines = list(map(lambda t: t.strip(), text.split("\n")))
        lines = list(filter(lambda l: l != "" and l[0] != "#", lines))
        ret = list(map(lambda l: l.split(sep), lines))
        if args.verbose:
            padding = 0 if max(map(lambda l: len(l), ret)) > 10 else 1
            table = Table(title=name, show_header=False, padding=(0,padding,0,padding))
            for row in ret:
                table.add_row(*list(map(lambda l: str(l), row)))
            progress.stop()
            console.print(table, "\n\n\n")
            progress.start()
        return ret

    # Makes station names uniform. For example M,1 and M,01 are turned to M01
    def station_id(char, number):
        return char + ("0" if number < 10 else "") + str(number)

    # Load station names (english and japanese for use in add_connections)
    update_progress("Loading stations (1/2 english)")
    stations_eng = {}
    with open("data/stations_eng.csv", encoding="utf8") as c_file:
        for line in parse_csv(c_file.read(), name="stations_eng.csv"):
            stations_eng[station_id(line[0], int(line[1]))] = line[2]
            
    update_progress("Loading stations (2/2 japanese)")
    stations_jap = {}
    with open("data/stations_jap.csv", encoding="utf8") as c_file:
        for line in parse_csv(c_file.read(), name="stations_jap.csv"):
            stations_jap[station_id(line[0], int(line[1]))] = line[2]

    # Load transition types (for use in add_connections)
    update_progress("Loading transition types")
    default_type = 1
    types = {}
    with open("data/types.csv", encoding="utf8") as c_file:
        for line in parse_csv(c_file.read(), name="types.csv"):
            types[int(line[0])] = line[1]


    # Add connection to stations structure
    def add_connection(container):
        def inner(origin, destination, typ, dist, dur):
            entry = {
                "name_en": str(stations_eng.get(origin)),
                "name_jp": str(stations_jap.get(origin)),
                "connections": []
            }
            try:
                entry = container[origin]
            except KeyError:
                pass

            conn = {
                "target_id": destination,
                "type_id": typ,
                "type": types.get(typ),
                "duration": dur,
                "distance": dist
            }
            entry["connections"] += [conn]
            container[origin] = entry
            return container

        return inner
    # Structure to hold the read information
    stations = {}
    add_station = add_connection(stations)

    # Add line to line structure
    def add_line(container):
        def inner(station, l_key, name):
            entry = {
                "name_en": "",
                "name_jp": ""
            }
            try:
                entry = container[station]
            except KeyError:
                pass

            entry[l_key] = name
            container[station] = entry
            return container

        return inner
    # Structure to hold the line information
    lines = {}
    add_line = add_line(lines)

    # Load line names (english)
    update_progress("Loading line names (1/2 english)")
    with open("data/lines_eng.csv", encoding="utf8") as c_file:
        for line in parse_csv(c_file.read(), name="lines_eng.csv"):
            add_line(line[0], "name_en", line[1])

    # Load line names (japanese)
    update_progress("Loading line names (2/2 japanese)")
    with open("data/lines_jap.csv", encoding="utf8") as c_file:
        for line in parse_csv(c_file.read(), name="lines_jap.csv"):
            add_line(line[0], "name_jp", line[1])

    # Read lines.csv and append all transitions to metro_map
    update_progress("Loading station transitions (1/2 lines.csv)")
    with open("data/lines.csv", encoding="utf8") as c_file:
        for line in parse_csv(c_file.read(), name="lines.csv"):
            char = line[0]
            first = int(line[1])
            last = int(line[2])

            distances = list(map(float, line[3:3+last-first])) # first n entries are distances in km
            durations = list(map(int  , line[3+last-first:]))  # next  n entries are durations in s

            for station_number in range(first, last):
                distance = distances[station_number - first]
                duration = durations[station_number - first]
                station_a = station_id(char, station_number)
                station_b = station_id(char, station_number + 1)
                stations = add_station(station_a, station_b, default_type, distance, duration)
                stations = add_station(station_b, station_a, default_type, distance, duration)

    # Read transitions.csv and append all transitions to metro_map
    update_progress("Loading station transitions (1/2 transitions.csv)")
    with open("data/transitions.csv", encoding="utf8") as c_file:
        for transition in parse_csv(c_file.read(), name="transitions.csv"):
            station1 = station_id(transition[0], int(transition[1]))
            station2 = station_id(transition[2], int(transition[3]))
            type = int(transition[4])
            distance = float(transition[5])
            duration =   int(transition[6])

            stations = add_station(station1, station2, type, distance, duration)
            stations = add_station(station2, station1, type, distance, duration)

    # Dump data as JSON
    update_progress("Dumping data as JSON")

    dump = {}
    dump["lines"] = lines
    dump["stations"] = stations
    dump["transition_types"] = types
    with open(args.output, "w+", encoding="utf8") as c_file:
        c_file.write(json.dumps(dump, indent=4, sort_keys=True, ensure_ascii=False))

    # Dump data in Matrix format
    if args.matrix:
        update_progress("Dumping data in matrix format... ")

        import pandas as pd
        import numpy as np
        
        # Get names of Stations for indexing (rows and columns)
        names = stations.keys()

        # Create empty matrix for transitions
        matrix_dist = np.full((len(names), len(names)), -1.0)
        matrix_time = np.full((len(names), len(names)), -1.0)

        # Index matrix
        matrix_dist = pd.DataFrame(matrix_dist, columns=names, index=names)
        matrix_time = pd.DataFrame(matrix_time, columns=names, index=names)
        
        # Fill matrix with data
        for station in stations:
            for connection in stations[station]["connections"]:
                matrix_dist[station][connection["target_id"]] = connection["distance"]
                matrix_time[station][connection["target_id"]] = connection["duration"]

        # Verbose output matrix
        if args.verbose:
            progress.stop()
            console.print("\n\nTransition distance matrix:\n")
            console.print(matrix_dist)
            console.print("\n\nTransition time matrix:\n")
            console.print(matrix_time)
            progress.start()

        # Save matrix files
        with open(matrix_dist_file, "w+", encoding="utf8") as c_file:
            c_file.write(matrix_dist.to_csv(index=True))

        with open(matrix_time_file, "w+", encoding="utf8") as c_file:
            c_file.write(matrix_time.to_csv(index=True))
        
    ###
    ### Done
    ###
    update_progress("Done")