Py.Cafe

import os
import sys
import subprocess
import importlib.util
import pandas as pd
import geopandas as gpd
import vizro.plotly.express as px
from vizro import Vizro
import vizro.models as vm
from vizro.models.types import capture
from vizro.tables import dash_ag_grid
from vizro.actions import export_data
from datetime import datetime, timedelta

import json



csv_head = "https://py.cafe/files/JayCampanell/kenya-child-malnutrition/"
geojson_url = "https://data.humdata.org/dataset/e66dbc70-17fe-4230-b9d6-855d192fc05c/resource/51939d78-35aa-4591-9831-11e61e555130/download/kenya.geojson"
@capture("graph")
def choropleth(data_frame, geojson, location, color):

    # Comment Out Acute vs Moderate vs Severe Options
    """
    if outcome == "Moderate" or outcome == "Acute":
        color_map = {'>=30%':"#bd0026", '[15.0%,30.0%)': "#f03b20", "[10.0%,15.0%)": "#fd8d3c",  "[3.0%,10.0%)": "#fecc5c", '[0%,3.0%)': "#ffffb2"}
        category_orders = {color: ['[0%,3.0%)', "[3.0%,10.0%)", "[10.0%,15.0%)", '[15.0%,30.0%)', '>=30.0%']}
    elif outcome == "Severe":
        color_map = {'>=5%':"#bd0026", '[3.0%,5.0%)': "#f03b20", "[1.0%,3.0%)": "#fd8d3c",  "[0.5%,1.0%)": "#fecc5c", '[0%,0.5%)': "#ffffb2"}
        category_orders = {color: ['[0%,0.5%)', "[0.5%,1.0%)", "[1.0%,3.0%)", '[3.0%,5.0%)', '>=5%']}
    """

    color_map = {'>=30%':"#bd0026", '[15.0%,30.0%)': "#f03b20", "[10.0%,15.0%)": "#fd8d3c",  "[3.0%,10.0%)": "#fecc5c", '[0%,3.0%)': "#ffffb2"}
    category_orders = {color: ['[0%,3.0%)', "[3.0%,10.0%)", "[10.0%,15.0%)", '[15.0%,30.0%)', '>=30.0%']}

    fig = px.choropleth_map(
        data_frame,
        geojson=geojson,
        locations=location,
        featureidkey="properties.shapeName",
        color=color,
        center={"lat": 0.0236, "lon": 37.9062},
        zoom=5,
        color_discrete_map=color_map,
        category_orders=category_orders,
        height=600
    )
    return fig


# Load and preprocess data
paths = os.listdir("results/")
df_concat = pd.DataFrame()



for file in paths:
    file_path = os.path.join(csv_head, "results/", file)
    model = file.split('_G', 1)[0].replace('.', '')
    df = pd.read_csv(file_path)
    df['model'] = model
    df_concat = pd.concat([df_concat, df], ignore_index=True)

with open("kenya.geojson", "r") as f:
    geojson = json.load(f)

df_concat['id'] = df_concat['subcounty'].str.replace(" Sub County", '')

valid_ids = {f["properties"]["shapeName"] for f in geojson["features"]}
df_concat = df_concat[df_concat['id'].isin(valid_ids)].copy()

df_concat['month_dt'] = pd.to_datetime(df_concat['month'], format="%B %Y", errors='coerce')

idx = (
    df_concat
    .groupby(['model', 'horizon', 'id'])['month_dt']
    .idxmax()
)

df_grouped = (
    df_concat.loc[idx]
    .drop(columns=['true', 'subcounty'], errors='ignore')
    .reset_index(drop=True)
)

gdf = gpd.GeoDataFrame(df_grouped)

base_date = gdf['month_dt'].max()
original_date = base_date.strftime("%B %Y")

one_month = (base_date + pd.DateOffset(months=1)).strftime("%B %Y")
three_month = (base_date + pd.DateOffset(months=3)).strftime("%B %Y")
six_month = (base_date + pd.DateOffset(months=6)).strftime("%B %Y")

gdf['Forecasted Month'] = gdf['horizon'].replace(
    to_replace={
        1: one_month,
        3: three_month,
        6: six_month
    }
)

view_table = gdf.loc[:, ["id", "month", "Forecasted Month", "IPC AMN", "model"]]
view_table.columns = ["Sub-county", "Base Month", "Forecasted Month", "IPC AMN", 'Model']

radio_options = [{'label': f"1 Month Forecast ({one_month})", 'value': one_month},
                    {'label': f"3 Month Forecast ({three_month})", 'value': three_month},
                    {'label': f"6 Month Forecast ({six_month})", 'value': six_month}]

model_options = [{'label': "Outcomes + Clinical + Satellite", 'value': "Modis+Prev_o+clinical"},
                    {'label': "Previous Outcomes + Clinical", "value": "Prev_o+clinical"},
                    {'label': "Previous Outcomes + Satellite", "value": "Modis+prev_outcome"}]

table = vm.Page(
    title="Kenya Malnutrition Table" + " (" + original_date + ")",
    components=[
        vm.Button(text="Download Table", actions=[vm.Action(function=export_data())]),
        vm.AgGrid(figure=dash_ag_grid(data_frame=view_table, dashGridOptions={"pagination": True}))
    ],
    controls=[
        vm.Filter(column="Model", selector=vm.RadioItems(value="Modis+Prev_o+clinical", options=model_options)),
        vm.Filter(column="Forecasted Month", selector=vm.RadioItems(value=three_month, options=radio_options))#,
        #vm.Filter(column="Outcome", selector=vm.RadioItems(value="Acute", options=["Acute", "Moderate", "Severe"]))
    ],
    layout=vm.Flex())

map_page = vm.Page(
    title="Kenya Malnutrition Map" + " (" + original_date + ")",
    components=[
        vm.Graph(figure=choropleth(
            data_frame=gdf,
            geojson=geojson,
            location="id",
            color="IPC AMN"
        ),
            id='graph'),
        vm.Card(text="""This visualization tool presents sub-county level data on malnutrition, including historical prevalence and future predictions 
                        based on Gradient Boosting across 1-, 3-, and 6-month time horizons. The malnutrition scale used is from the IPC defintion to categorize the 
                        malnutrition rate into one of 5 categories: [0, 3%), [3, 10%), [10, 15%), [15, 30%), and ≥ 30%. This allows for easy comparisons to be made between the subcounties.
                        It also includes Gross Primary Productivity data obtained from NASA's MODIS satellite to depict agricultural activity in each sub-county. 
                        The clinical data for the models is sourced from aggregate DHIS2 data.""" ,
                extra={'style': {'height': '150px'}}
                )
    ],
    controls=[
        vm.Filter(column="model", selector=vm.RadioItems(value="Modis+Prev_o+clinical", options=model_options)),
        vm.Filter(column="Forecasted Month", selector=vm.RadioItems(value=three_month, options=radio_options))#,
        #vm.Parameter(targets=['graph.outcome'], selector=vm.RadioItems(value="Acute", options=["Acute", "Moderate", "Severe"]))
    ],
    layout=vm.Flex())

dashboard = vm.Dashboard(pages=[map_page, table])

# Launch browser (only if not running as PyInstaller bundle)
#if not getattr(sys, 'frozen', False):

# Build and run the Vizro dashboard
app = Vizro().build(dashboard)  

if __name__ == "__main__":  
    app.run()