import geopandas as gpd
import rasterio as rio
from rasterio.plot import show
# import the datasets
countries = gpd.read_file('countries.geojson')
stations = gpd.read_file('ozone_stations.csv')
cities = gpd.read_file('populated_places.geojson')
rain = rio.open('rain.tif')
# preprocess stations
stations['geometry'] = gpd.points_from_xy(x=stations['X'], y=stations['Y'])
stations = stations.set_crs('epsg:4326')
# select the cities that are in south american
continents = countries.dissolve(by='continent')
south_america = continents.query("continent == 'South America'")
# filter major cities in south america
sa_cities = cities.clip(south_america)
major_cities = sa_cities.sort_values(by='pop_max', ascending=False).iloc[:10]
# perform the spatial nearest
major_cities = major_cities.to_crs('EPSG:32662')[['name', 'pop_max', 'geometry']]
stations = stations.to_crs('EPSG:32662')[['platform_name', 'country', 'X', 'Y', 'geometry']]
combined = major_cities.sjoin_nearest(stations, how='left')
# create a list with the coordinates to be sampled
combined = combined.to_crs('EPSG:4326')
cities_coords = [(x, y) for x, y in zip(combined['geometry'].x, combined['geometry'].y)]
stations_coords = [(x, y) for x, y in zip(combined['X'].astype('float'), combined['Y'].astype('float'))]
# perform the raster sampling
combined['cities_rain'] = list(rain.sample(cities_coords))
combined['stations_rain'] = list(rain.sample(stations_coords))
### display the results
# plot South America
ax = south_america.plot(facecolor='none')
show(rain, ax=ax)
combined.plot(ax=ax, color='orange')
selected_stations = gpd.gpd.points_from_xy(combined['X'], combined['Y'])
gpd.GeoSeries(selected_stations).plot(ax=ax, color='red', marker='^')
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')
ax.figure.savefig('output/proximity_sampling.png', dpi=300)
print(combined.to_html())