src/data/loaders/idpr.py

"""
IDPR data loader for infiltration vs runoff tendency.
Single Responsibility: Load BRGM IDPR values.
"""
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Optional
from .base import BaseDataLoader


class IDPRLoader(BaseDataLoader):
    """
    Loads IDPR (infiltration vs runoff tendency) values from BRGM.
    """

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

        Args:
            data_path: Path to IDPR CSV file
        """
        super().__init__(data_path)

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

        Returns:
            DataFrame with IDPR values (spatial points or basin aggregates)
        """
        df = pd.read_csv(self.data_path)
        return df

    def get_metadata(self) -> dict:
        """Get IDPR data metadata."""
        meta = super().get_metadata()
        meta.update({
            "data_type": "idpr_infiltration_runoff",
            "source_organization": "BRGM"
        })
        return meta

    def _find_value_column(self, df: pd.DataFrame) -> str:
        """Best-effort detection of the IDPR value column by name."""
        candidates = [c for c in df.columns if "idpr" in c.lower()]
        if candidates:
            return candidates[0]
        numeric_cols = [c for c in df.select_dtypes(include="number").columns
                         if c.lower() not in ("x", "y", "lat", "lon", "latitude", "longitude")]
        if numeric_cols:
            return numeric_cols[0]
        raise ValueError(
            "Could not auto-detect the IDPR value column; pass value_col explicitly."
        )

    def plot_distribution(
        self,
        df: Optional[pd.DataFrame] = None,
        value_col: Optional[str] = None,
        bins: int = 40,
        figsize: tuple = (8, 5),
        save_path: Optional[Path] = None,
    ) -> plt.Axes:
        """
        Histogram of IDPR values. IDPR is centered around 0: negative values
        indicate infiltration-dominated areas, positive values runoff-dominated.

        Args:
            df: Optional pre-loaded DataFrame. If not provided, data is loaded from disk.
            value_col: Column holding the IDPR value. Auto-detected if not given.
            bins: Number of histogram bins.
            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()
        value_col = value_col or self._find_value_column(df)

        fig, ax = plt.subplots(figsize=figsize)
        ax.hist(df[value_col].dropna(), bins=bins, color="teal", edgecolor="white")
        ax.axvline(0, color="black", linewidth=1, linestyle="--",
                    label="0 (infiltration ↔ runoff)")

        ax.set_title("IDPR Value Distribution")
        ax.set_xlabel(value_col)
        ax.set_ylabel("Count")
        ax.legend(fontsize=8)
        ax.grid(True, alpha=0.3)
        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_spatial(
        self,
        df: Optional[pd.DataFrame] = None,
        value_col: Optional[str] = None,
        x_col: Optional[str] = None,
        y_col: Optional[str] = None,
        figsize: tuple = (9, 8),
        save_path: Optional[Path] = None,
    ) -> plt.Axes:
        """
        Spatial scatter of IDPR values, colored on a diverging scale
        centered at 0 (infiltration vs. runoff tendency).

        Args:
            df: Optional pre-loaded DataFrame. If not provided, data is loaded from disk.
            value_col: Column holding the IDPR value. Auto-detected if not given.
            x_col: Longitude/easting column. Auto-detected from common names if not given.
            y_col: Latitude/northing column. Auto-detected from common names if not given.
            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()
        value_col = value_col or self._find_value_column(df)

        x_col = x_col or next((c for c in ["lon", "longitude", "x"] if c in df.columns), None)
        y_col = y_col or next((c for c in ["lat", "latitude", "y"] if c in df.columns), None)
        if not x_col or not y_col:
            raise ValueError(
                "Could not auto-detect coordinate columns; pass x_col/y_col explicitly."
            )

        plot_df = df.dropna(subset=[x_col, y_col, value_col])
        vmax = plot_df[value_col].abs().max()

        fig, ax = plt.subplots(figsize=figsize)
        scatter = ax.scatter(
            plot_df[x_col], plot_df[y_col], c=plot_df[value_col],
            cmap="RdBu_r", vmin=-vmax, vmax=vmax, s=10, alpha=0.8,
        )
        cbar = plt.colorbar(scatter, ax=ax)
        cbar.set_label(f"{value_col} (infiltration ← 0 → runoff)")

        ax.set_title("IDPR Spatial Distribution")
        ax.set_xlabel(x_col)
        ax.set_ylabel(y_col)
        ax.set_aspect("equal")
        plt.tight_layout()

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

        return ax