import cartopy
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import csv


class Flight:
    def __init__(self, name, city, country, IATA, latitude, longitude):
        self.longitude = longitude
        self.latitude = latitude
        self.IATA = IATA
        self.country = country
        self.city = city
        self.name = name

    def __repr__(self):
        return f"Airport {self.name} in {self.city}, {self.country}, with IATA: {self.IATA}"


def collect_flights_to_list(flights, airports) -> list:
    with open(flights, "r", encoding="UTF-8") as csv_file, \
            open(airports, "r", encoding="UTF-8") as dat_file:
        flights_reader = csv.reader(csv_file, delimiter=";")
        data_dict = {}
        for line in dat_file.readlines():
            line = line.split(",")
            name = line[1].replace('"', '')
            city = line[2].replace('"', '')
            country = line[3].replace('"', '')
            IATA = line[4].replace('"', '')
            latitude = line[6]
            longitude = line[7]
            data_dict[IATA] = Flight(name, city, country, IATA, latitude, longitude)

        next(flights_reader)

        airports_data = []
        for line in flights_reader:
            city, IATA = line
            airports_data.append(data_dict.get(IATA))

        return airports_data


def add_flight(start_airport: Flight, target_airport: Flight, color):
    start_longitude = float(start_airport.longitude)
    start_latitude = float(start_airport.latitude)
    end_longitude = float(target_airport.longitude)
    end_latitude = float(target_airport.latitude)

    plt.plot([start_longitude, end_longitude], [start_latitude, end_latitude],
             color=color, linewidth=1.5, marker='o',
             transform=ccrs.Geodetic())


def add_airport_name(airport: Flight):
    longitude = float(airport.longitude)
    latitude = float(airport.latitude)
    plt.text(longitude - 0.75, latitude + 0.25, airport.IATA,
             horizontalalignment='right',
             transform=ccrs.PlateCarree())


if __name__ == '__main__':
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.add_feature(cartopy.feature.LAND)
    ax.add_feature(cartopy.feature.OCEAN)
    ax.add_feature(cartopy.feature.COASTLINE)
    ax.add_feature(cartopy.feature.BORDERS, linestyle='-', alpha=.5)
    plt.title("Dmitri Härm")

    airports_fly = collect_flights_to_list("data/otselennud.csv", "data/airports.dat")
    airports_covid_time = collect_flights_to_list("data/flights21.csv", "data/airports.dat")
    tll = airports_fly[0]
    add_airport_name(tll)

    for airport in airports_fly:
        add_flight(tll, airport, "blue")
        add_airport_name(airport)

    for airport in airports_covid_time:
        add_flight(tll, airport, "red")

    plt.show()