ihk-dijkstra/graph.py

"""
Module, which creates a graph
"""

from classes import DataSet, LocationType, TransportMethod, TransportKind
from args import Settings


def calc_co2(distance: float, kind: TransportKind) -> float:
    """
    Return co2 emission in kg
    """
    match kind:
        case TransportKind.INDIVIDUAL:
            if distance <= 1:
                return 0
            else:
                # TODO: CITY == AUTOBAHN
                return (distance - 1) * TransportMethod.CITY.value.co2
        case TransportKind.TRAIN:
            # TODO: OPENV == ICE
            return distance * TransportMethod.ICE.value.co2
        case TransportKind.FLYING:
            return distance * TransportMethod.AIRPLANE.value.co2


def calc_time(
        distance: float,
        kind: TransportKind,
        flights: dict = None,
        frm: dict = None,
        ) -> float:
    """
    distance: distance without penalty
    kind: kind of transport
    flights: all flights (optional)
    frm: start location (need for flights, optional)
    """
    match kind:
        case TransportKind.INDIVIDUAL:
            if distance <= 1:
                return distance / TransportMethod.WALK.value.speed
            elif distance <= 10 + 1:
                speed = 1 / TransportMethod.WALK.value.speed
                new_dist = (distance - 1)
                speed += (new_dist / TransportMethod.CITY.value.speed)
                return speed
            elif distance < 2000 + 10 + 1:
                speed = 1 / TransportMethod.WALK.value.speed
                speed += 10 / TransportMethod.CITY.value.speed
                new_dist = (distance - 1 - 10)
                speed += (new_dist / TransportMethod.AUTOBAHN.value.speed)
                return speed
        case TransportKind.TRAIN:
            if distance <= 25:
                return distance / TransportMethod.OEPNV.value.speed
            else:
                return (distance / TransportMethod.ICE.value.speed)
        case TransportKind.FLYING:
            assert(flights)
            assert(frm)
            stops = flights[frm]["stops"]
            domestic = flights[frm]["domestic"]
            if domestic:
                time = 2
            else:
                time = 3
            waittime = (stops + 1) * time
            return waittime + (distance / TransportMethod.AIRPLANE.value.speed)


def create_graph(dataset: DataSet, settings: Settings) -> dict:
    """
    Creates the initial graph, with all edges
    """
    locations = dataset.locations
    graph: dict = {}

    if settings.debug:
        print("Creating graph...")

    # add nodes with no edges
    for start in locations:
        graph.update({start: []})

    # add edges
    for start in locations:
        for dest in locations:
            # skip, if we wouldnt go anywhere
            if start == dest:
                continue
            if settings.debug:
                print(f"Searching for nodes from {start} to {dest}...")
            distance = locations[start].distance(locations[dest])

            # Individualtransport
            if distance <= 2000 + 10 + 1:
                dist = distance + distance * TransportKind.INDIVIDUAL.value
                time = calc_time(dist, TransportKind.INDIVIDUAL)
                co2 = calc_co2(dist, TransportKind.INDIVIDUAL)
                new_connection: dict = {dest:
                                        {"kind": TransportKind.INDIVIDUAL,
                                         "co2": co2,
                                         "time": time}}
                graph[start].append(new_connection)
                if settings.debug:
                    print(f"Added new INDIVIDUAL connection from {start} to {dest}")

            # Train
            is_haltestelle = locations[start].type == LocationType.HALTESTELLE
            is_airport = locations[start].type == LocationType.FLUGHAFEN
            is_same_continent = locations[start].is_same_continent(locations[dest])
            # trains only drive from haltestelle and same continent
            if (is_haltestelle or is_airport) and is_same_continent:
                dest_type = locations[dest].type
                # there are only trains between haltestellen or airports
                is_dest_haltestelle = dest_type == LocationType.HALTESTELLE
                is_dest_flughafen = dest_type == LocationType.FLUGHAFEN
                if is_dest_haltestelle or is_dest_flughafen:
                    dist = distance + distance * TransportKind.TRAIN.value
                    time = calc_time(dist, TransportKind.TRAIN)
                    co2 = calc_co2(dist, TransportKind.TRAIN)
                    new_connection: dict = {dest:
                                            {"kind": TransportKind.TRAIN,
                                             "co2": co2,
                                             "time": time}}
                    graph[start].append(new_connection)
                    if settings.debug:
                        print(f"Added new TRAIN connection from {start} to {dest}")

    # Flying
    if settings.debug:
        print("Adding flights...")
    flights = dataset.flights
    for flight in flights:
        start = flight
        dest = flights[flight]["to"]
        distance = locations[start].distance(locations[dest])
        dist = distance + distance * TransportKind.FLYING.value
        time = calc_time(dist, TransportKind.FLYING, flights, start)
        co2 = calc_co2(dist, TransportKind.FLYING)
        new_connection: dict = {dest:
                                {"kind": TransportKind.FLYING,
                                 "co2": co2,
                                 "time": time}}
        graph[start].append(new_connection)
        if settings.debug:
            print(f"Added new FLYING connection from {start} to {dest}")

    return graph