import geopandas as gpd
import matplotlib.pyplot as plt
import pandas as pd

world = gpd.read_file("../../data/110m_cultural.zip", layer="ne_110m_admin_0_countries")
cities = gpd.read_file("../../data/110m_cultural.zip",
                       layer="ne_110m_populated_places")


def find_neighbouring_countries(df, country_name):
    target = df[df["NAME"] == country_name]

    if len(target) == 0:
        print(f"{country_name} not found")
        return gpd.GeoDataFrame()

    # sjoin  spatial-join
    neighbours = gpd.sjoin(target, df, how="inner", predicate="touches")
    return neighbours[["NAME_right", "CONTINENT_right", "geometry"]]

def get_close_cities(df, city_name):
    target = df[df["NAMEASCII"] == city_name]

    if len(target) == 0:
        print(f"{city_name} not found")
        return gpd.GeoDataFrame()

    # sjoin  spatial-join
    #combined = gpd.sjoin(target, df, how="inner")
    #combined["distance"] = combined.distance(target.iloc[0])
    neighbours = gpd.sjoin_nearest(target, df, max_distance=10e12)
    return neighbours


#intersects
#contains
#within
#touches
#crosses
#overlaps
german_neighbours = find_neighbouring_countries(world, "Germany")
print(german_neighbours[["NAME_right", "CONTINENT_right"]])
print(len(german_neighbours))
print("-"*100)
w2 = world.to_crs("EPSG:3857")
c2 = cities.to_crs("EPSG:3857")
print(list(c2.columns))
german_neighbours = get_close_cities(c2, "Paris")
print(len(german_neighbours))
print(german_neighbours)

w2.plot(column="CONTINENT", legend=True)
plt.show()

def get_cities_in_country(city_df, world_df, country):
    country_match = world_df[world_df["NAME"] == country]

    if len(country_match) == 0:
        print(f"Returning empty for {country}")
        return gpd.GeoDataFrame()

    cities = gpd.sjoin(country_match, city_df,
                       how="inner",
                       predicate="contains")
    cities = cities[["NAME_left", "NAMEASCII"]]
    return cities

def get_cities_in_countries(city_df, world_df, countries):
    res = get_cities_in_country(city_df, world_df, countries[0])
    for country in countries[1:]:
        tmp = get_cities_in_country(city_df, world_df, country)
        res = pd.concat((res, tmp))
    return res


# 1) get_cities_in_country    (sjoin)  # NAME_right, NAME_left
# 2) get_cities_in_countries  (concat)
us_cities = get_cities_in_country(cities, world, "United States of America")
print(us_cities)
french_cities = get_cities_in_country(cities, world, "Switzerland")
print(french_cities)

central_eurpean_cities = get_cities_in_countries(
    cities, world,
    ["France", "Switzerland", "Narnia", "Germany", "Denmark"]
)
print(central_eurpean_cities)

# United States = 9
# France = 4
# Italy: 3
# Germany: 1
