src/data/loaders/shapefile.py

"""
Shapefile loader for watershed boundaries.
Single Responsibility: Load geospatial watershed data.
"""
import geopandas as gpd
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Optional
from .base import BaseDataLoader


class ShapefileLoader(BaseDataLoader):
    """
    Loads watershed shapefiles for La Risle and La Eure basins.
    """

    def __init__(self, data_path: Path, crs: Optional[str] = None):
        """
        Initialize shapefile loader.

        Args:
            data_path: Path to shapefile (.shp) or directory containing shapefiles
            crs: Coordinate reference system to reproject to (optional)
        """
        super().__init__(data_path)
        self.crs = crs

    def load(self) -> gpd.GeoDataFrame:
        """
        Load shapefile(s) into GeoDataFrame.

        Returns:
            GeoDataFrame with watershed geometries and attributes
        """
        if self.data_path.suffix == ".shp":
            gdf = gpd.read_file(self.data_path)
        elif self.data_path.is_dir():
            # Find .shp files in directory
            shp_files = list(self.data_path.glob("*.shp"))
            if not shp_files:
                raise FileNotFoundError(f"No shapefiles found in {self.data_path}")
            # Load first shapefile found (or merge if multiple)
            gdf = gpd.read_file(shp_files[0])
        else:
            raise ValueError(f"Invalid shapefile path: {self.data_path}")

        # Reproject if CRS specified
        if self.crs and gdf.crs != self.crs:
            gdf = gdf.to_crs(self.crs)

        return gdf

    def get_metadata(self) -> dict:
        """Get shapefile metadata."""
        meta = super().get_metadata()
        meta.update({
            "data_type": "watershed_shapefile",
            "target_crs": self.crs
        })
        return meta

    def print_summary(self, gdf: Optional[gpd.GeoDataFrame] = None) -> None:
        """
        Print the shapefile's data and metadata in a readable format.

        Args:
            gdf: Optional pre-loaded GeoDataFrame. If not provided, the
                 shapefile will be loaded from disk.
        """
        if gdf is None:
            gdf = self.load()

        meta = self.get_metadata()

        print("=" * 60)
        print("SHAPEFILE METADATA")
        print("=" * 60)
        for key, value in meta.items():
            print(f"{key:>15}: {value}")

        print()
        print("=" * 60)
        print("SHAPEFILE DATA SUMMARY")
        print("=" * 60)
        print(f"{'Feature count':>15}: {len(gdf)}")
        print(f"{'CRS':>15}: {gdf.crs}")
        print(f"{'Columns':>15}: {list(gdf.columns)}")
        print(f"{'Bounds':>15}: {tuple(gdf.total_bounds)}")
        print(f"{'Geometry types':>15}: {gdf.geom_type.unique().tolist()}")

        print()
        print("-" * 60)
        print("Attribute preview:")
        print("-" * 60)
        print(gdf.drop(columns="geometry").head())
        print("=" * 60)

    def plot(
        self,
        gdf: Optional[gpd.GeoDataFrame] = None,
        column: Optional[str] = None,
        title: Optional[str] = None,
        figsize: tuple = (10, 10),
        save_path: Optional[Path] = None,
        cmap: str = "Blues",
        edgecolor: str = "black",
    ) -> plt.Axes:
        """
        Plot the watershed polygon(s).

        Args:
            gdf: Optional pre-loaded GeoDataFrame. If not provided, the
                 shapefile will be loaded from disk.
            column: Optional column name to color/shade features by
                    (e.g. a category or numeric attribute). If None,
                    all features are drawn with a single fill color.
            title: Optional plot title. Defaults to the data path stem.
            figsize: Figure size in inches (width, height).
            save_path: If provided, saves the figure to this path
                       instead of (or in addition to) displaying it.
            cmap: Matplotlib colormap used when `column` is set.
            edgecolor: Outline color for the polygons.

        Returns:
            The matplotlib Axes object, for further customization.
        """
        if gdf is None:
            gdf = self.load()

        fig, ax = plt.subplots(figsize=figsize)

        if column and column in gdf.columns:
            gdf.plot(column=column, ax=ax, cmap=cmap, edgecolor=edgecolor, legend=True)
        else:
            gdf.plot(ax=ax, color="steelblue", edgecolor=edgecolor, alpha=0.6)

        ax.set_title(title or f"Watershed: {self.data_path.stem}")
        ax.set_xlabel("Easting")
        ax.set_ylabel("Northing")
        ax.set_aspect("equal")

        crs_label = str(gdf.crs) if gdf.crs else "Unknown CRS"
        ax.annotate(
            crs_label,
            xy=(0.01, 0.01),
            xycoords="axes fraction",
            fontsize=8,
            color="gray",
        )

        plt.tight_layout()

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

        return ax