src/data/loaders/ades.py

"""
ADES groundwater level data loader.
Single Responsibility: Load groundwater (karst) connectivity features.
"""
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from pathlib import Path
from typing import Optional, List
from .base import BaseDataLoader


class ADESLoader(BaseDataLoader):
    """
    Loads ADES groundwater level data for karst connectivity analysis.
    """

    def __init__(
        self,
        data_path: Path,
        file_format: str = "csv",
        station_ids: Optional[List[str]] = None
    ):
        """
        Initialize ADES loader.

        Args:
            data_path: Path to ADES data directory
            file_format: Data format (csv, json, etc.)
            station_ids: Optional list of groundwater station IDs
        """
        super().__init__(data_path)
        self.file_format = file_format.lower()
        self.station_ids = station_ids

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

        Returns:
            DataFrame with groundwater level measurements
        """
        if self.file_format == "csv":
            df = self._load_csv()
        else:
            raise ValueError(f"Unsupported format: {self.file_format}")

        # Filter by station IDs if provided
        if self.station_ids:
            if "station_id" in df.columns:
                df = df[df["station_id"].isin(self.station_ids)]

        # Ensure date column is datetime if present
        if "date" in df.columns:
            df["date"] = pd.to_datetime(df["date"])

        return df

    def _load_csv(self) -> pd.DataFrame:
        """Load from CSV file(s)."""
        if self.data_path.is_file():
            return pd.read_csv(self.data_path)
        elif self.data_path.is_dir():
            # Load groundwater levels and stations, then merge
            levels_path = self.data_path / "groundwater_levels_watershed.csv"
            stations_path = self.data_path / "groundwater_stations.csv"

            if levels_path.exists() and stations_path.exists():
                levels = pd.read_csv(levels_path)
                stations = pd.read_csv(stations_path)

                # Merge to get coordinates with each measurement
                df = pd.merge(
                    levels[['code_bss', 'date_mesure', 'niveau_nappe_eau', 'profondeur_nappe']],
                    stations[['code_bss', 'x', 'y']],
                    on='code_bss',
                    how='left'
                )
                df = df.rename(columns={
                    'date_mesure': 'date',
                    'niveau_nappe_eau': 'groundwater_level_m',
                    'profondeur_nappe': 'groundwater_depth_m',
                    'x': 'lon',
                    'y': 'lat'
                })
                return df
            else:
                csv_files = list(self.data_path.glob("*.csv"))
                if not csv_files:
                    raise FileNotFoundError(f"No CSV files found in {self.data_path}")
                dfs = [pd.read_csv(f) for f in csv_files]
                return pd.concat(dfs, ignore_index=True)
        else:
            raise ValueError(f"Invalid path: {self.data_path}")

    def aggregate_to_stations(
        self,
        groundwater_df: pd.DataFrame,
        station_coords: pd.DataFrame,
        max_distance_km: float = 20.0
    ) -> pd.DataFrame:
        """
        Aggregate groundwater levels spatially to hydrometric stations.

        Args:
            groundwater_df: ADES data with [date, code_bss, groundwater_level_m, lat, lon]
            station_coords: Station coordinates [station_code, lat, lon]
            max_distance_km: Max distance to consider wells (km)

        Returns:
            DataFrame with [date, station_code, avg_groundwater_level_m, n_wells]
        """
        # Haversine distance
        def haversine(lat1, lon1, lat2, lon2):
            R = 6371  # Earth radius in km
            lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
            dlat = lat2 - lat1
            dlon = lon2 - lon1
            a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
            return R * 2 * np.arcsin(np.sqrt(a))

        results = []

        for _, station in station_coords.iterrows():
            station_code = station['station_code']
            station_lat = station['lat']
            station_lon = station['lon']

            # Find wells within radius
            groundwater_df['distance_km'] = groundwater_df.apply(
                lambda row: haversine(station_lat, station_lon, row['lat'], row['lon']),
                axis=1
            )
            nearby_wells = groundwater_df[groundwater_df['distance_km'] <= max_distance_km].copy()

            if len(nearby_wells) == 0:
                continue

            # Aggregate by date
            daily_agg = nearby_wells.groupby('date').agg({
                'groundwater_level_m': 'mean',
                'groundwater_depth_m': 'mean',
                'code_bss': 'nunique'
            }).reset_index()

            daily_agg['station_code'] = station_code
            daily_agg = daily_agg.rename(columns={
                'groundwater_level_m': 'avg_groundwater_level_m',
                'groundwater_depth_m': 'avg_groundwater_depth_m',
                'code_bss': 'n_wells'
            })

            results.append(daily_agg)

        if not results:
            return pd.DataFrame()

        return pd.concat(results, ignore_index=True)

    def get_metadata(self) -> dict:
        """Get ADES data metadata."""
        meta = super().get_metadata()
        meta.update({
            "data_type": "ades_groundwater",
            "format": self.file_format
        })
        return meta

    def plot_levels(
        self,
        df: Optional[pd.DataFrame] = None,
        wells: Optional[List[str]] = None,
        figsize: tuple = (12, 5),
        save_path: Optional[Path] = None,
    ) -> plt.Axes:
        """
        Plot groundwater level time series, one line per well/station.
        Works with either raw well-level data (code_bss, groundwater_level_m)
        or data already aggregated to hydrometric stations via
        `aggregate_to_stations` (station_code, avg_groundwater_level_m).

        Args:
            df: Optional pre-loaded DataFrame. If not provided, data is loaded from disk.
            wells: Optional list of well/station IDs to plot. Defaults to all.
            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()

        if "station_code" in df.columns and "avg_groundwater_level_m" in df.columns:
            id_col, value_col = "station_code", "avg_groundwater_level_m"
        elif "code_bss" in df.columns and "groundwater_level_m" in df.columns:
            id_col, value_col = "code_bss", "groundwater_level_m"
        else:
            raise ValueError(
                "Could not find a recognizable (id, level) column pair in the data"
            )

        plot_df = df.dropna(subset=[value_col])
        if wells:
            plot_df = plot_df[plot_df[id_col].isin(wells)]
        if plot_df.empty:
            raise ValueError("No groundwater level data available to plot")

        fig, ax = plt.subplots(figsize=figsize)
        for well_id, group in plot_df.groupby(id_col):
            group = group.sort_values("date")
            ax.plot(group["date"], group[value_col], linewidth=1, marker="o",
                     markersize=2, label=well_id)

        ax.set_title("Groundwater Level")
        ax.set_xlabel("Date")
        ax.set_ylabel("Groundwater level (m)")
        ax.xaxis.set_major_locator(mdates.AutoDateLocator())
        ax.xaxis.set_major_formatter(mdates.ConciseDateFormatter(ax.xaxis.get_major_locator()))
        ax.legend(title=id_col, 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_well_locations(
        self,
        df: Optional[pd.DataFrame] = None,
        figsize: tuple = (8, 8),
        save_path: Optional[Path] = None,
    ) -> plt.Axes:
        """
        Scatter map of well locations, colored by mean groundwater level.
        Requires raw well-level data with lat/lon (i.e. before aggregation
        to hydrometric stations).

        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()

        required = {"code_bss", "lat", "lon", "groundwater_level_m"}
        if not required.issubset(df.columns):
            raise ValueError(
                f"plot_well_locations requires columns {required}; "
                "this looks like already-aggregated station data."
            )

        summary = df.dropna(subset=["lat", "lon"]).groupby(
            ["code_bss", "lat", "lon"]
        )["groundwater_level_m"].mean().reset_index()

        fig, ax = plt.subplots(figsize=figsize)
        scatter = ax.scatter(
            summary["lon"], summary["lat"], c=summary["groundwater_level_m"],
            cmap="viridis_r", s=40, edgecolor="black", linewidth=0.3,
        )
        cbar = plt.colorbar(scatter, ax=ax)
        cbar.set_label("Mean groundwater level (m)")

        ax.set_title(f"Groundwater Well Locations (n={len(summary)})")
        ax.set_xlabel("Longitude")
        ax.set_ylabel("Latitude")
        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