import pandas as pd
import matplotlib.pyplot as plt
import geopandas as gpd
from shapely import wkb

def function(rows):
    # turn the RiskScape input into a dataframe
    df = pd.DataFrame(rows)
    # turn the dataframe into a choroplath map
    choropleth_map(df, model_output('regional-collapsed-buildings.png'))

def choropleth_map(df, filename):  
    # deserialize the WKB and turn it back into geometry
    geometry = [ wkb.loads(row['the_geom'][0]) for row in df['Region'] ]
    gdf = gpd.GeoDataFrame(df, crs="EPSG:4326", geometry=geometry)

    ax = gdf.plot(column='Collapsed', cmap='Reds', legend=True)
    # labels
    gdf.apply(lambda x: ax.annotate(text=x['Region']['Region'], xy=x.geometry.centroid.coords[0], ha='center', size=5), axis=1)
    ax.set_axis_off()
    ax.set_title('Collapsed buildings in Upolu by region')
    plt.savefig(filename)
