src/data_processing/load_shelters.py

"""
Load and process shelter data from the pipeline-ready CSV.

Schema: shelter_name, union, upazila, district, shelter_type, capacity
Source: data/raw/shelters/shelters_coxs_bazar_pipeline_ready.csv (610 records)
"""

import pandas as pd
import logging
from src.config import SHELTER_DATA_FILE, PROCESSED_SHELTERS, PROCESSED_DIR

logger = logging.getLogger(__name__)

EXPECTED_COLUMNS = [
    "shelter_name", "union", "upazila", "district", "shelter_type", "capacity"
]


def load_shelters() -> pd.DataFrame:
    """
    Load shelter CSV, validate schema, and return a clean DataFrame.
    Also saves a processed GeoJSON stub (tabular, no geometry yet) for
    downstream spatial joins with union boundaries.

    Returns:
        pd.DataFrame with 610 shelter records.
    """
    logger.info("Loading shelter data...")

    if not SHELTER_DATA_FILE.exists():
        raise FileNotFoundError(
            f"Shelter data file not found: {SHELTER_DATA_FILE}\n"
            "Expected the pipeline-ready CSV at this path."
        )

    df = pd.read_csv(SHELTER_DATA_FILE)
    logger.info(f"Loaded {len(df)} shelter records from {SHELTER_DATA_FILE.name}")

    # ── Schema validation ─────────────────────────────────────────────────
    missing = [c for c in EXPECTED_COLUMNS if c not in df.columns]
    if missing:
        raise ValueError(
            f"Shelter CSV missing required columns: {missing}\n"
            f"Found columns: {list(df.columns)}"
        )

    # ── Type enforcement ──────────────────────────────────────────────────
    df["capacity"] = pd.to_numeric(df["capacity"], errors="coerce").fillna(0).astype(int)

    # ── Strip whitespace from string columns ──────────────────────────────
    for col in ["shelter_name", "union", "upazila", "district", "shelter_type"]:
        df[col] = df[col].astype(str).str.strip()

    # ── Summary stats ─────────────────────────────────────────────────────
    total_capacity = df["capacity"].sum()
    types = df["shelter_type"].value_counts().to_dict()
    upazilas = sorted(df["upazila"].unique())

    logger.info(f"  Total capacity: {total_capacity:,}")
    logger.info(f"  Shelter types: {types}")
    logger.info(f"  Upazilas ({len(upazilas)}): {upazilas}")

    # ── Save processed CSV ────────────────────────────────────────────────
    PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
    processed_csv = PROCESSED_DIR / "shelters_processed.csv"
    df.to_csv(processed_csv, index=False)
    logger.info(f"Saved processed shelters → {processed_csv}")

    return df


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
    shelters = load_shelters()
    print(f"\n✅ {len(shelters)} shelters loaded successfully")
    print(f"   Total capacity: {shelters['capacity'].sum():,}")
    print(f"   Types: {shelters['shelter_type'].nunique()}")
    print(f"   Upazilas: {shelters['upazila'].nunique()}")