import geopandas as gpd
import pandas as pd
import matplotlib.pyplot as plt
import os
from shapely.geometry import MultiPolygon
os.chdir('C:/Users/jtrum/Desktop/hydroatlas_africa')

# Create an empty dictionary to store the geodataframes
gdfs = {} #empty dictionary of gdfs

for level in range(1, 13): #for loop to load all shapefiles
    filename = f'hybas_af_lev{level:02d}_v1c.shp'  #:02d ensures two-digit format
    gdf = gpd.read_file(filename) #read the shapefile
    key = f'L{level:02d}' #add the geodataframe to the dictionary with the corresponding key
    gdfs[key] = gdf

# this script takes data from HydroBASINS at the continent level, and uploads the
# 12 different levels of basins. it also loads the geojson of the AOI and then dependent
# on which level of basin is selected, it will determine which basins are within the AOI
# and complete a union of those basins and save the catchment as a new geojson file.

# load the AOI
aoi = gpd.read_file('luan.geojson')

# write a for loop to load all the basin levels and name them L1 - L12
for i in range(1,13):
    globals()['L{}'.format(i)] = gpd.read_file('hybas_af_lev0{}_v1c.shp'.format(i))
    

# to load in one level of basin, uncomment the level you want to load
# L1 = gpd.read_file('hybas_af_lev01_v1c.shp')
# L2 = gpd.read_file('hybas_af_lev02_v1c.shp')
# L3 = gpd.read_file('hybas_af_lev03_v1c.shp')
# L4 = gpd.read_file('hybas_af_lev04_v1c.shp')
# L5 = gpd.read_file('hybas_af_lev05_v1c.shp')
L6 = gpd.read_file('hybas_af_lev06_v1c.shp')
# L7 = gpd.read_file('hybas_af_lev07_v1c.shp')
# L8 = gpd.read_file('hybas_af_lev08_v1c.shp')
# L9 = gpd.read_file('hybas_af_lev09_v1c.shp')
# L10 = gpd.read_file('hybas_af_lev10_v1c.shp')
# L11 = gpd.read_file('hybas_af_lev11_v1c.shp')
# L12 = gpd.read_file('hybas_af_lev12_v1c.shp')

# create a function to determine which basins are within the AOI
def basin_in_aoi(basin, aoi):
    basin['in_aoi'] = basin.geometry.within(aoi.geometry.iloc[0])  # Check if basins are within the AOI geometry

    basins_within_aoi = basin[basin['in_aoi'] == 1]  # Create a new DataFrame with only the basins within the AOI
    if len(basins_within_aoi) > 0:
        catchment = basins_within_aoi.unary_union  # Make a union of the basins within AOI
        catchment = MultiPolygon([catchment])  # Wrap it in a MultiPolygon in case it's a single geometry
        catchment.to_file('catchment.geojson')  # Save as a new GeoJSON file
        return catchment
    else:
        return None  # No basins within AOI

catchment = basin_in_aoi(L6, aoi)
if catchment:
    gpd.GeoSeries([catchment]).plot()  # Plot the catchment if it exists

# create a function to determine which basins are within the AOI
def basin_in_aoi(basin, aoi):
    basin['in_aoi'] = 0
    for i in range(len(aoi)):
        if basin.geometry.within(aoi.loc[i, 'geometry']): # if the basin is within the AOI, set the in_aoi column to 1
            basin.loc[:, 'in_aoi'] = 1
        basin = basin[basin['in_aoi'] == 1]  # create a new dataframe with only the basins within the AOI, and make a union of them to return one catchment
        basin = basin.unary_union
        basin.to_file('catchment.geojson')         # save as a new geojson file
    return basin

catchment = basin_in_aoi(L6, aoi)
catchment.plot()

#set working directory


#read in shapefiles of different levels
# L1 = gpd.read_file('hybas_af_lev01_v1c.shp')
# L2 = gpd.read_file('hybas_af_lev02_v1c.shp')
# L3 = gpd.read_file('hybas_af_lev03_v1c.shp')
L4 = gpd.read_file('hybas_af_lev04_v1c.shp')
L5 = gpd.read_file('hybas_af_lev05_v1c.shp')
L6 = gpd.read_file('hybas_af_lev06_v1c.shp')
L7 = gpd.read_file('hybas_af_lev07_v1c.shp')
L8 = gpd.read_file('hybas_af_lev08_v1c.shp')
L9 = gpd.read_file('hybas_af_lev09_v1c.shp')
L10 = gpd.read_file('hybas_af_lev10_v1c.shp')
L11 = gpd.read_file('hybas_af_lev11_v1c.shp')
L12 = gpd.read_file('hybas_af_lev12_v1c.shp')

#load in luanda geojson
luanda = gpd.read_file('luan.geojson')