src/data/loaders/catchment.py

"""
Catchment (drainage basin) area loader.
Single Responsibility: Load per-station catchment/drainage-basin area,
as produced by download_catchment_area.py from Hub'Eau's own
hydrometrie referentiel (referentiel/sites, field `surface_bv`).
"""
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Optional
from .base import BaseDataLoader


class CatchmentAreaLoader(BaseDataLoader):
    """
    Loads per-station catchment (drainage basin) area in km².

    Hub'Eau doesn't publish surface_bv for every site — some rows will
    have a missing catchment_area_km2 (e.g. observer/manual stations,
    partner-network sites). That's expected, not a loading error; see
    `coverage()` to check how much of your station set has real data
    before relying on it (e.g. for drainage-area-ratio discharge scaling).
    """

    def __init__(self, data_path: Path):
        """
        Initialize catchment area loader.

        Args:
            data_path: Path to catchment_area.csv
                (columns: station_code, code_site, catchment_area_km2)
        """
        super().__init__(data_path)

    def load(self) -> pd.DataFrame:
        """
        Load catchment area data.

        Returns:
            DataFrame with [station_code, code_site, catchment_area_km2]
        """
        df = pd.read_csv(self.data_path)
        required = {"station_code", "catchment_area_km2"}
        missing = required - set(df.columns)
        if missing:
            raise ValueError(f"catchment_area.csv is missing expected column(s): {missing}")
        return df

    def coverage(self, df: Optional[pd.DataFrame] = None) -> dict:
        """
        Quick summary of how many stations actually have a catchment
        area value, since Hub'Eau leaves this blank for a meaningful
        fraction of sites.

        Returns:
            {"n_stations": int, "n_with_area": int, "fraction": float}
        """
        if df is None:
            df = self.load()
        n_total = len(df)
        n_with = int(df["catchment_area_km2"].notna().sum())
        return {
            "n_stations": n_total,
            "n_with_area": n_with,
            "fraction": n_with / n_total if n_total else 0.0,
        }

    def get_metadata(self) -> dict:
        """Get catchment area metadata."""
        meta = super().get_metadata()
        meta.update({
            "data_type": "catchment_area",
            "source_organization": "Hub'Eau (hydrometrie referentiel/sites, surface_bv)",
            "units": "km2",
        })
        return meta

    def plot_coverage(
        self,
        df: Optional[pd.DataFrame] = None,
        figsize: tuple = (5, 5),
        save_path: Optional[Path] = None,
    ) -> plt.Axes:
        """
        Simple pie/bar of how many stations have a catchment area value
        vs. how many don't — a quick sanity check before trusting this
        feature for drainage-area scaling.

        Args:
            df: Optional pre-loaded DataFrame. If not provided, data is loaded from disk.
            figsize: Figure size in inches.
            save_path: If provided, saves the figure to this path.

        Returns:
            The matplotlib Axes object.
        """
        if df is None:
            df = self.load()
        cov = self.coverage(df)

        fig, ax = plt.subplots(figsize=figsize)
        counts = [cov["n_with_area"], cov["n_stations"] - cov["n_with_area"]]
        labels = [f"Has area\n({counts[0]})", f"No area\n({counts[1]})"]
        ax.pie(counts, labels=labels, colors=["#2E6F95", "#C7CDC3"],
               autopct="%1.0f%%", startangle=90,
               textprops={"fontsize": 10})
        ax.set_title(f"Catchment Area Coverage ({cov['n_stations']} stations)",
                      fontsize=12, fontweight="bold")
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches="tight")
            print(f"Plot saved to {save_path}")

        return ax

    def plot_by_station(
        self,
        df: Optional[pd.DataFrame] = None,
        stations: Optional[list] = None,
        figsize: tuple = (10, 8),
        save_path: Optional[Path] = None,
    ) -> plt.Axes:
        """
        Horizontal bar chart of catchment area per station, sorted
        smallest to largest. Stations without data are omitted (see
        `plot_coverage` for how many that is).

        Args:
            df: Optional pre-loaded DataFrame. If not provided, data is loaded from disk.
            stations: Optional list of station codes to include. Defaults to all with data.
            figsize: Figure size in inches.
            save_path: If provided, saves the figure to this path.

        Returns:
            The matplotlib Axes object.
        """
        if df is None:
            df = self.load()

        plot_df = df.dropna(subset=["catchment_area_km2"]).copy()
        if stations:
            plot_df = plot_df[plot_df["station_code"].isin(stations)]
        if plot_df.empty:
            raise ValueError("No stations with catchment area data to plot")

        plot_df = plot_df.sort_values("catchment_area_km2")

        fig, ax = plt.subplots(figsize=figsize)
        ax.barh(plot_df["station_code"], plot_df["catchment_area_km2"], color="#2E6F95")

        ax.set_title("Catchment Area by Station", fontsize=12, fontweight="bold")
        ax.set_xlabel("Catchment area (km²)")
        ax.set_ylabel("Station")
        ax.grid(True, alpha=0.3, axis="x")
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches="tight")
            print(f"Plot saved to {save_path}")

        return ax