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

def function(rows, parameters={}):
    # turn the RiskScape input into a dataframe
    df = pd.DataFrame(rows)
    # work out the attribute to plot
    attribute = parameters.get('plot_attribute', 'Collapsed')
    # turn the dataframe into a choroplath map
    filename = 'regional-%s-buildings.png' % attribute.lower()
    choropleth_map(df, model_output(filename), column=attribute)

def choropleth_map(df, filename, column='Collapsed'):
    # 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=column, 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(column + ' buildings in Upolu by region')
    plt.savefig(filename)
