import sys
import os
import rasterio

import geopandas as gpd
import pandas as pd

import GOSTurban.UrbanRaster as urban
import GOSTRocks.mapMisc as mapMisc

# Define input population raster
tutorial_folder = "../../data/tutorial_data"
aoi_file = os.path.join(tutorial_folder, "AOI.geojson")
pop_file = os.path.join(tutorial_folder, "wp_2020_1k_AOI.tif")

# Define output files
out_urban = os.path.join(tutorial_folder, "urban_extents.geojson")
out_hd_urban = os.path.join(tutorial_folder, "hd_urban_extents.geojson")

inAOI = gpd.read_file(aoi_file)

# Shouldn't need to execute this unless you change your AOI;
#    you will need to find a global population to extract from
if not os.path.exists(pop_file):
    sys.path.append("../../../gostrocks/src")
    import GOSTRocks.rasterMisc as rMisc

    global_population = "/path/to/global/pop_layer/ppp_2020_1km_Aggregated.tif"
    inR = rasterio.open(global_population)
    rMisc.clipRaster(inR, inAOI, pop_file)

inR = rasterio.open(pop_file)
# Initiate the urban calculator
urban_calculator = urban.urbanGriddedPop(inR)

# Extract the urban extents (minimum density 300/km2, minimum total population 5000)
urban_extents = urban_calculator.calculateUrban(
    densVal=300, totalPopThresh=5000, smooth=False, queen=False, verbose=True
)
urban_extents["Type"] = 1
urban_extents.head()

# Map of urban extents
mapMisc.static_map_vector(urban_extents, "Type")

# Extract the high density urban extents (minimum density 1500/km2, minimum total population 50000)
hd_urban_extents = urban_calculator.calculateUrban(
    densVal=1500,
    totalPopThresh=50000,
    smooth=True,
    queen=True,  # high density extents use queen's case contiguity, and
    verbose=True,
)  # High density extents have hole smoothing applied.
hd_urban_extents["Type"] = 2
hd_urban_extents.head()

mapMisc.static_map_vector(hd_urban_extents, "Type", colormap="Blues")