Harmonize climate and health data into a Chap-ready CSV file

from pathlib import Path
import re

import pandas as pd

# ---------------------------------------------------------------------
# Parameters
# ---------------------------------------------------------------------
FREQ = "monthly"

# This notebook consumes harmonized outputs produced by earlier workflow notebooks.
DATA_DIR = Path("../data").resolve()

# Harmonized inputs (already aligned to the same orgunit and monthly time_period)
DENGUE_CSV = DATA_DIR / "nepal-dengue-monthly-admin.csv" # columns: location, time_period, disease_cases
POP_CSV    = DATA_DIR / "nepal-worldpop-monthly-admin.csv"    # columns: location, time_period, population
PRCP_CSV   = DATA_DIR / "nepal-era5-prcp-monthly-admin.csv"   # columns: location, time_period, precip_mm
T2M_CSV    = DATA_DIR / "nepal-era5-t2m-monthly-admin.csv"    # columns: location, time_period, t2m_c

# Output
OUTPUT_CSV = DATA_DIR / "nepal_dengue_pop_climate_chap.csv"

# ---------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------

PERIOD_YYYYMM = re.compile(r"^\d{6}$")

def read_csv(path: Path, required: list[str]) -> pd.DataFrame:
    if not path.exists():
        raise FileNotFoundError(f"Missing input: {path}")

    df = pd.read_csv(path)

    # check for missing columns
    missing = [c for c in required if c not in df.columns]
    if missing:
        raise KeyError(f"{path.name} missing {missing}. Found: {list(df.columns)}")

    # normalize keys
    df["location"] = df["location"].astype(str).str.strip()
    df["time_period"] = (
        df["time_period"].astype(str).str.replace(r"\D", "", regex=True).str.zfill(6)
    )

    # check period format
    bad = ~df["time_period"].str.match(PERIOD_YYYYMM)
    if bad.any():
        raise ValueError(f"Invalid time_period (YYYYMM). Examples: {df.loc[bad, 'time_period'].head(5).tolist()}")

    # drop missing location and periods
    df = df.dropna(subset=["location", "time_period"]).copy()

    # check for duplicates
    if df.duplicated(["location", "time_period"]).any():
        raise ValueError(f"{path.name} has duplicate (location, time_period) rows")

    return df

def month_index(start_yyyymm: str, end_yyyymm: str) -> pd.Index:
    return pd.period_range(start_yyyymm, end_yyyymm, freq="M").astype(str).str.replace("-", "", regex=False)

# Load harmonized inputs (already monthly, already on the admin-unit spine)
dengue = read_csv(DENGUE_CSV, ["location", "time_period", "disease_cases"])
pop    = read_csv(POP_CSV,    ["location", "time_period", "population"])
prcp   = read_csv(PRCP_CSV,   ["location", "time_period", "tp"])
t2m    = read_csv(T2M_CSV,    ["location", "time_period", "t2m_c"])

# Coerce numeric columns (preserve missing values as NaN)
dengue["disease_cases"] = pd.to_numeric(dengue["disease_cases"], errors="coerce")
pop["population"]       = pd.to_numeric(pop["population"], errors="coerce")
prcp["tp"]              = pd.to_numeric(prcp["tp"], errors="coerce")
t2m["t2m_c"]            = pd.to_numeric(t2m["t2m_c"], errors="coerce")

# Quick shape summary
summary = pd.DataFrame({
    "rows": [len(dengue), len(pop), len(prcp), len(t2m)],
    "locations": [dengue["location"].nunique(), pop["location"].nunique(), prcp["location"].nunique(), t2m["location"].nunique()],
    "months": [dengue["time_period"].nunique(), pop["time_period"].nunique(), prcp["time_period"].nunique(), t2m["time_period"].nunique()],
}, index=["dengue", "population", "precip", "t2m"])

summary

# Determine global modelling window from dengue (target outcome)
start_tp = dengue["time_period"].min()
end_tp = dengue["time_period"].max()

months = month_index(start_tp, end_tp)

# Spatial spine: union of all locations present across harmonized inputs
locations = pd.Index(
    pd.unique(pd.concat([
        dengue["location"],
        pop["location"],
        prcp["location"],
        t2m["location"],
    ], ignore_index=True)).astype(str),
    name="location"
).sort_values()

# Full modelling grid: all locations × all months
grid = pd.MultiIndex.from_product([locations, months], names=["location", "time_period"]).to_frame(index=False)

grid

# Start from the grid
df = grid.copy()

# Outcome + covariates (all keyed by location + time_period)
df = df.merge(dengue, on=["location", "time_period"], how="left")
df = df.merge(pop,    on=["location", "time_period"], how="left")
df = df.merge(prcp,   on=["location", "time_period"], how="left")
df = df.merge(t2m,    on=["location", "time_period"], how="left")

df

# One row per (location, time_period)
assert not df.duplicated(["location", "time_period"]).any(), "Duplicate (location, time_period) rows found."

# Required Chap fields present
required_cols = {"location", "time_period", "disease_cases"}
missing = required_cols - set(df.columns)
assert not missing, f"Missing required columns: {missing}"

# Missingness summary (expected; no thresholds enforced here)
print('Missing values (share of total, 0-1):')
df.isna().mean().sort_values(ascending=False).head(20)

from dhis2eo.integrations.chap import dataframe_to_chap_csv

column_map = {
    "time_period": "time_period",
    "location": "location",
    "disease_cases": "disease_cases",
    "population": "population",   # optional but recommended
}

dataframe_to_chap_csv(
    df=df,
    output_path=OUTPUT_CSV,
    freq="monthly",
    column_map=column_map,
)

OUTPUT_CSV

df_chap = pd.read_csv(OUTPUT_CSV)
df_chap