src/data_processing/load_population.py

"""
Load population data from WorldPop raster and compute per-union statistics.
Falls back to estimated population figures if raster processing fails.
"""

import geopandas as gpd
import pandas as pd
import numpy as np
import logging
from rasterstats import zonal_stats
from src.config import (
    POPULATION_RASTER_FILE, PROCESSED_UNIONS,
    POPULATION_BY_UNION, PROCESSED_DIR,
)
import rasterio

logger = logging.getLogger(__name__)

# Approximate upazila populations for Cox's Bazar (2022 census estimates)
FALLBACK_POPULATION = {
    "Cox's Bazar Sadar": 620_000,
    "Cox'S Bazar Sadar": 620_000,
    "Sadar": 620_000,
    "Ramu": 320_000,
    "Chakaria": 520_000,
    "Kutubdia": 125_000,
    "Maheshkhali": 330_000,
    "Pekua": 220_000,
    "Teknaf": 280_000,
    "Ukhia": 210_000,
}


def load_population():
    """
    Extract population per union from WorldPop raster, or estimate.

    Returns:
        DataFrame with columns: union_name, upazila_name, population,
        pop_density_per_cell, cells_count.
    """
    logger.info("Processing population data...")
    PROCESSED_DIR.mkdir(parents=True, exist_ok=True)

    if not PROCESSED_UNIONS.exists():
        raise FileNotFoundError(
            f"Union boundaries not found at {PROCESSED_UNIONS}. "
            "Run load_boundaries.py first."
        )

    unions = gpd.read_file(PROCESSED_UNIONS)

    if POPULATION_RASTER_FILE.exists():
        try:
            pop_df = _extract_from_raster(unions)
            if pop_df is not None:
                zero_frac = (pop_df["population"] == 0).mean()
                total_pop = pop_df["population"].sum()
                if total_pop == 0 or zero_frac > 0.8:
                    logger.warning(
                        f"Raster extraction produced invalid data "
                        f"(total={total_pop:,}, zero_frac={zero_frac:.1%}). "
                        "Falling back to estimator."
                    )
                    pop_df = _estimate_population(unions)
                pop_df.to_csv(POPULATION_BY_UNION, index=False)
                logger.info(f"Saved population data → {POPULATION_BY_UNION}")
                return pop_df
        except Exception as e:
            logger.warning(f"Raster extraction failed: {e}")
            logger.info("Falling back to estimated population...")
    else:
        logger.info("WorldPop raster not found — using estimates...")

    pop_df = _estimate_population(unions)
    pop_df.to_csv(POPULATION_BY_UNION, index=False)
    logger.info(f"Saved estimated population → {POPULATION_BY_UNION}")
    return pop_df


def _extract_from_raster(unions):
    """Zonal statistics on WorldPop raster per union polygon."""
    logger.info(f"  Running zonal statistics on {POPULATION_RASTER_FILE.name}...")

    logger.info(f"  Unions CRS: {unions.crs}")
    if unions.crs and unions.crs.to_epsg() != 4326:
        unions = unions.to_crs(epsg=4326)

    with rasterio.open(str(POPULATION_RASTER_FILE)) as src:
        nodata_val = src.nodata
        logger.info(f"  Raster CRS: {src.crs}")

    # Force evaluating to list in case it's a generator
    stats = list(zonal_stats(
        unions, str(POPULATION_RASTER_FILE),
        stats=["sum", "mean", "count"],
        nodata=nodata_val,
        all_touched=True
    ))

    result = pd.DataFrame()
    if "union_name" in unions.columns:
        result["union_name"] = unions["union_name"].values
    if "upazila_name" in unions.columns:
        result["upazila_name"] = unions["upazila_name"].values

    result["population"] = [
        int(round(s["sum"])) if s["sum"] else 0 for s in stats
    ]
    result["pop_density_per_cell"] = [
        round(s["mean"], 2) if s["mean"] else 0 for s in stats
    ]
    result["cells_count"] = [s["count"] if s["count"] else 0 for s in stats]

    for col in ("union_name", "upazila_name", "population", "pop_density_per_cell", "cells_count"):
        if col not in result.columns:
            result[col] = 0 if col != "union_name" and col != "upazila_name" else "Unknown"

    total = result["population"].sum()
    logger.info(f"  Total extracted population: {total:,}")
    logger.info(f"  Unions: {len(result)}")
    logger.info(
        f"  Range: {result['population'].min():,} – {result['population'].max():,}"
    )
    return result


def _estimate_population(unions):
    """Generate estimated population from known upazila totals."""
    logger.info("  Generating estimated population figures...")
    np.random.seed(42)

    result = pd.DataFrame()
    result["union_name"] = (
        unions["union_name"].values
        if "union_name" in unions.columns
        else [f"Union_{i}" for i in range(len(unions))]
    )
    result["upazila_name"] = (
        unions["upazila_name"].values
        if "upazila_name" in unions.columns
        else "Unknown"
    )

    populations = []
    for _, row in unions.iterrows():
        upazila = str(row.get("upazila_name", ""))

        # Look up known population
        upazila_pop = 300_000  # default
        for key, pop in FALLBACK_POPULATION.items():
            if key.lower() in upazila.lower():
                upazila_pop = pop
                break

        # Count sibling unions for distribution
        n_unions = max(
            1,
            len(unions[unions["upazila_name"] == upazila])
            if "upazila_name" in unions.columns
            else 10,
        )
        base = upazila_pop / n_unions
        populations.append(int(base * (0.7 + 0.6 * np.random.random())))

    result["population"] = populations
    result["pop_density_per_cell"] = 0
    result["cells_count"] = 0
    logger.info(f"  Estimated total population: {sum(populations):,}")
    return result


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
    load_population()