src/generate_plots.py

"""
Generate a curated, high-quality set of plots for every data loader in the
project: hydrometric, ADES groundwater, IDPR, SAFRAN/ERA5, watershed
shapefile, and station elevations.

Each loader's plots are generated independently and skipped gracefully
(with a clear message) if its data files aren't present or an optional
dependency (geopandas, xarray) isn't installed — one missing dataset
never stops the others from running.

Usage:
    python generate_plots.py --data-root DATASETS_DIR --output-root OUTPUT_DIR

Expected layout under DATASETS_DIR (any missing piece is just skipped):
    hydrometric/discharge_observations.csv
    hydrometric/waterlevel_observations.csv
    ades/groundwater_levels_watershed.csv
    ades/groundwater_stations.csv
    idpr.csv
    safran/era5_*.nc
    shapefile/*.shp
    station_elevations.csv
"""
import argparse
from pathlib import Path

import pandas as pd
import matplotlib.pyplot as plt

try:
    from .data.loaders.base import BaseDataLoader  # noqa: F401 (import check only)
    from .data.loaders.hydrometric import HydrometricLoader
    from .data.loaders.ades import ADESLoader
    from .data.loaders.idpr import IDPRLoader
    from .data.loaders.station_elevations import StationElevationsLoader
except ImportError:
    import sys
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
    from src.data.loaders.hydrometric import HydrometricLoader
    from src.data.loaders.ades import ADESLoader
    from src.data.loaders.idpr import IDPRLoader
    from src.data.loaders.station_elevations import StationElevationsLoader


# ----------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------

def _section(title: str) -> None:
    print(f"\n{'=' * 60}\n{title}\n{'=' * 60}")


def select_complete_stations(
    df, n: int, recent_years: int = 5, require_both_variables: bool = True
) -> list:
    """Pick the n stations with the most complete data in the most recent window."""
    cutoff = df["date"].max() - pd.DateOffset(years=recent_years)
    recent = df[df["date"] >= cutoff]
    expected_days = (recent["date"].max() - recent["date"].min()).days + 1

    scores = []
    for station_code, group in recent.groupby("station_code"):
        disc_frac = group["discharge_m3s"].notna().sum() / expected_days
        wl_frac = group["waterlevel_mm"].notna().sum() / expected_days
        if require_both_variables and (disc_frac == 0 or wl_frac == 0):
            continue
        scores.append((station_code, min(disc_frac, wl_frac)))

    scores.sort(key=lambda x: x[1], reverse=True)
    selected = [s for s, _ in scores[:n]]

    if len(selected) < n:
        counts = df.dropna(subset=["discharge_m3s"]).groupby("station_code").size()
        fallback = counts.sort_values(ascending=False).head(n).index.tolist()
        selected = list(dict.fromkeys(selected + fallback))[:n]

    return selected


def filter_recent(df, years: int = 5):
    cutoff = df["date"].max() - pd.DateOffset(years=years)
    return df[df["date"] >= cutoff]


# ----------------------------------------------------------------------
# Per-loader plot generators — each returns True if it ran, False if skipped
# ----------------------------------------------------------------------

def generate_hydrometric_plots(data_dir: Path, output_dir: Path, recent_years: int) -> bool:
    _section("HYDROMETRIC")
    if not data_dir.exists():
        print(f"Skipped: {data_dir} not found")
        return False

    loader = HydrometricLoader(data_path=data_dir)
    try:
        full_df = loader.load()
    except FileNotFoundError as e:
        print(f"Skipped: {e}")
        return False

    df = filter_recent(full_df, years=recent_years)
    print(f"{len(full_df)} rows, {full_df['station_code'].nunique()} stations. "
          f"Using last {recent_years}y: {df['date'].min().date()}–{df['date'].max().date()}")

    comparison_stations = select_complete_stations(full_df, n=3, recent_years=recent_years)
    deepdive_station = select_complete_stations(full_df, n=1, recent_years=recent_years)[0]
    print(f"Comparison stations: {comparison_stations} | Deep dive: {deepdive_station}")

    output_dir.mkdir(parents=True, exist_ok=True)

    loader.plot_waterlevel(df=df, stations=comparison_stations,
                            save_path=output_dir / "waterlevel_comparison.png")
    plt.close("all")
    loader.plot_discharge(df=df, stations=comparison_stations,
                           save_path=output_dir / "discharge_comparison.png")
    plt.close("all")
    loader.plot_seasonal_climatology("discharge_m3s", df=df, stations=comparison_stations,
                                      save_path=output_dir / "climatology_discharge.png")
    plt.close("all")

    station_df = df[df["station_code"] == deepdive_station]
    loader.plot_station(deepdive_station, df=station_df,
                         save_path=output_dir / "station_deepdive.png")
    plt.close("all")
    try:
        loader.plot_rating_curve(deepdive_station, df=station_df,
                                  save_path=output_dir / "rating_curve.png")
        plt.close("all")
    except ValueError as e:
        print(f"  Skipped rating curve: {e}")

    print(f"Saved plots to {output_dir}")
    return True


def generate_ades_plots(data_dir: Path, output_dir: Path, n_wells: int = 3) -> bool:
    _section("ADES GROUNDWATER")
    if not data_dir.exists():
        print(f"Skipped: {data_dir} not found")
        return False

    loader = ADESLoader(data_path=data_dir)
    try:
        df = loader.load()
    except FileNotFoundError as e:
        print(f"Skipped: {e}")
        return False

    if df.empty:
        print("Skipped: no groundwater data loaded")
        return False

    output_dir.mkdir(parents=True, exist_ok=True)

    # Pick the n_wells with the most observations, for a readable comparison.
    id_col = "code_bss" if "code_bss" in df.columns else "station_code"
    value_col = "groundwater_level_m" if "groundwater_level_m" in df.columns else "avg_groundwater_level_m"
    top_wells = (df.dropna(subset=[value_col]).groupby(id_col).size()
                 .sort_values(ascending=False).head(n_wells).index.tolist())
    print(f"Wells: {top_wells}")

    loader.plot_levels(df=df, wells=top_wells, save_path=output_dir / "groundwater_levels.png")
    plt.close("all")
    try:
        loader.plot_well_locations(df=df, save_path=output_dir / "well_locations.png")
        plt.close("all")
    except ValueError as e:
        print(f"  Skipped well map: {e}")

    print(f"Saved plots to {output_dir}")
    return True


def generate_idpr_plots(data_path: Path, output_dir: Path) -> bool:
    _section("IDPR")
    if not data_path.exists():
        print(f"Skipped: {data_path} not found")
        return False

    loader = IDPRLoader(data_path=data_path)
    df = loader.load()

    output_dir.mkdir(parents=True, exist_ok=True)
    loader.plot_distribution(df=df, save_path=output_dir / "idpr_distribution.png")
    plt.close("all")
    try:
        loader.plot_spatial(df=df, save_path=output_dir / "idpr_spatial.png")
        plt.close("all")
    except ValueError as e:
        print(f"  Skipped spatial plot: {e}")

    print(f"Saved plots to {output_dir}")
    return True


def generate_safran_plots(
    data_dir: Path, output_dir: Path, station_coords: "pd.DataFrame | None"
) -> bool:
    _section("SAFRAN / ERA5")
    if not data_dir.exists():
        print(f"Skipped: {data_dir} not found")
        return False
    if station_coords is None or station_coords.empty:
        print("Skipped: no station coordinates available (needs station_elevations.csv "
              "or another source of station lat/lon)")
        return False

    try:
        from .data.loaders.safran import SAFRANLoader
    except ImportError:
        try:
            from src.data.loaders.safran import SAFRANLoader
        except ImportError as e:
            print(f"Skipped: {e}")
            return False
    try:
        import xarray  # noqa: F401
    except ImportError:
        print("Skipped: xarray is not installed in this environment")
        return False

    loader = SAFRANLoader(data_path=data_dir, station_coords=station_coords)
    try:
        df = loader.load()
    except FileNotFoundError as e:
        print(f"Skipped: {e}")
        return False

    output_dir.mkdir(parents=True, exist_ok=True)
    loader.plot_temperature(df=df, save_path=output_dir / "temperature.png")
    plt.close("all")
    loader.plot_precipitation(df=df, save_path=output_dir / "precipitation_cumulative.png")
    plt.close("all")

    print(f"Saved plots to {output_dir}")
    return True


def generate_shapefile_plots(data_path: Path, output_dir: Path) -> bool:
    _section("WATERSHED SHAPEFILE")
    if not data_path.exists():
        print(f"Skipped: {data_path} not found")
        return False

    try:
        from .data.loaders.shapefile import ShapefileLoader
    except ImportError:
        try:
            from src.data.loaders.shapefile import ShapefileLoader
        except ImportError as e:
            print(f"Skipped: {e}")
            return False
    try:
        import geopandas  # noqa: F401
    except ImportError:
        print("Skipped: geopandas is not installed in this environment")
        return False

    loader = ShapefileLoader(data_path=data_path)
    try:
        gdf = loader.load()
    except (FileNotFoundError, ValueError) as e:
        print(f"Skipped: {e}")
        return False

    output_dir.mkdir(parents=True, exist_ok=True)
    loader.plot(gdf=gdf, save_path=output_dir / "watershed_boundary.png")
    plt.close("all")

    print(f"Saved plots to {output_dir}")
    return True


def generate_station_elevation_plots(data_path: Path, output_dir: Path) -> "pd.DataFrame | None":
    """Returns the coordinates DataFrame (for reuse by SAFRAN) or None if skipped."""
    _section("STATION ELEVATIONS")
    if not data_path.exists():
        print(f"Skipped: {data_path} not found")
        return None

    loader = StationElevationsLoader(data_path=data_path)
    df = loader.load()

    output_dir.mkdir(parents=True, exist_ok=True)
    loader.plot_elevation_bar(df=df, save_path=output_dir / "elevation_bar.png")
    plt.close("all")
    loader.plot_station_map(df=df, save_path=output_dir / "station_map.png")
    plt.close("all")

    print(f"Saved plots to {output_dir}")
    return loader.get_coordinates_df()


# ----------------------------------------------------------------------
# Orchestrator
# ----------------------------------------------------------------------

def generate_all(data_root: Path, output_root: Path, recent_years: int = 5) -> None:
    results = {}

    # Station elevations first: SAFRAN needs its coordinates output.
    station_coords = generate_station_elevation_plots(
        data_root / "station_elevations.csv", output_root / "station_elevations"
    )
    results["station_elevations"] = station_coords is not None

    results["hydrometric"] = generate_hydrometric_plots(
        data_root / "hydrometric", output_root / "hydrometric", recent_years
    )
    results["ades"] = generate_ades_plots(
        data_root / "ades", output_root / "ades"
    )
    results["idpr"] = generate_idpr_plots(
        data_root / "idpr.csv", output_root / "idpr"
    )
    results["safran"] = generate_safran_plots(
        data_root / "safran", output_root / "safran", station_coords
    )
    results["shapefile"] = generate_shapefile_plots(
        data_root / "shapefile", output_root / "shapefile"
    )

    _section("SUMMARY")
    for name, ok in results.items():
        print(f"  {'✓' if ok else '—'} {name}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate plots for all data loaders")
    parser.add_argument("--data-root", type=Path, default=Path("datasets"),
                         help="Root directory containing each loader's data subfolder/file")
    parser.add_argument("--output-root", type=Path, default=Path("plots"),
                         help="Root directory to save generated plots (one subfolder per loader)")
    parser.add_argument("--recent-years", type=int, default=5,
                         help="Recent-years window for hydrometric station selection/plots")
    args = parser.parse_args()

    generate_all(args.data_root, args.output_root, recent_years=args.recent_years)


if __name__ == "__main__":
    main()