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

from classes import DataSet, LocationType, TransportMethod, TransportKind


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 * 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:
            distance += distance * TransportKind.INDIVIDUAL.value
            if distance <= 1:
                return distance / TransportMethod.WALK.value.speed
            elif distance <= 10 + 1:
                speed = 1 / TransportMethod.WALK.value.speed
                return speed + ((distance - 1) / TransportMethod.CITY.value.speed)
            elif distance < 2000 + 10 + 1:
                speed = 1 / TransportMethod.WALK.value.speed
                speed += 10 / TransportMethod.WALK.value.speed
                return speed + ((distance - 1 - 10) / TransportMethod.AUTOBAHN.value.speed)
        case TransportKind.TRAIN:
            distance += distance * TransportKind.TRAIN.value
            if distance <= 25:
                return distance / TransportMethod.OEPNV.value.speed
            else:
                speed = 25 / TransportMethod.OEPNV.value.speed
                # NOTE: should we subtract here?
                return speed + ((distance - 25) / TransportMethod.ICE.value.speed)
        case TransportKind.FLYING:
            assert(flights != None)
            assert(frm != None)
            distance += distance * TransportKind.FLYING.value
            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) -> dict:
    """
    Creates the initial graph, with all edges
    """
    locations = dataset.locations
    graph: dict = {}

    # 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
            print(f"Searching for nodes from {start} to {dest}...")
            distance = locations[start].distance(locations[dest])

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

            # Train
            if locations[start].type == LocationType.HALTESTELLE:
                dest_type = locations[dest].type
                # there are only trains between haltestellen or airports
                if dest_type == LocationType.HALTESTELLE or dest_type == LocationType.FLUGHAFEN:
                    time = calc_time(distance, TransportKind.TRAIN)
                    co2 = calc_co2(distance, TransportKind.TRAIN)
                    new_connection: dict = {dest:
                                            {"kind": TransportKind.TRAIN,
                                             "co2": co2,
                                             "time": time}}
                    graph[start].append(new_connection)
                    print(f"Added new TRAIN connection from {start} to {dest}")

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

    return graph