"""
Topological feature extractor for basin characteristics.
Single Responsibility: Extract elevation, slope, drainage density, TWI, area.
"""
import numpy as np
import geopandas as gpd
from typing import Dict, Any, Optional
from pathlib import Path
from .base import BaseFeatureExtractor


class TopologicalExtractor(BaseFeatureExtractor):
    """
    Extracts topological/terrain features from basin geometries and DEM data.

    Features extracted:
    - Basin area (km²)
    - Mean/min/max elevation (m)
    - Mean basin slope (degrees)
    - Drainage density (km/km²)
    - Topographic Wetness Index (TWI)
    """

    def __init__(self, dem_path: Optional[Path] = None, **kwargs):
        """
        Initialize topological extractor.

        Args:
            dem_path: Path to DEM raster (optional, for elevation/slope/TWI)
            **kwargs: Additional parameters
        """
        super().__init__(**kwargs)
        self.dem_path = Path(dem_path) if dem_path else None

    def extract(self, basin_gdf: gpd.GeoDataFrame) -> Dict[str, Any]:
        """
        Extract topological features from basin.

        Args:
            basin_gdf: GeoDataFrame with basin geometry

        Returns:
            Dictionary with topological features
        """
        if not self.validate_inputs(basin_gdf):
            raise ValueError("Invalid basin GeoDataFrame")

        features = {}

        # Basic geometric features
        features.update(self._extract_basic_geometry(basin_gdf))

        # DEM-based features (if DEM provided)
        if self.dem_path and self.dem_path.exists():
            features.update(self._extract_dem_features(basin_gdf))

        return features

    def _extract_basic_geometry(self, basin_gdf: gpd.GeoDataFrame) -> Dict[str, float]:
        """Extract basic geometric features."""
        # Ensure CRS is projected (for area calculation)
        if basin_gdf.crs and basin_gdf.crs.is_geographic:
            # Reproject to appropriate UTM zone
            basin_gdf = basin_gdf.to_crs(basin_gdf.estimate_utm_crs())

        # Basin area in km²
        area_km2 = basin_gdf.geometry.area.sum() / 1e6

        # Basin centroid
        centroid = basin_gdf.geometry.centroid.iloc[0]

        return {
            "basin_area_km2": area_km2,
            "centroid_lon": centroid.x,
            "centroid_lat": centroid.y
        }

    def _extract_dem_features(self, basin_gdf: gpd.GeoDataFrame) -> Dict[str, float]:
        """
        Extract DEM-based features: elevation, slope, TWI.

        Note: This is a placeholder for DEM processing.
        Full implementation would use rasterio to:
        1. Clip DEM to basin boundary
        2. Calculate zonal statistics (mean, min, max elevation)
        3. Derive slope from DEM
        4. Calculate TWI using flow accumulation and slope
        """
        try:
            import rasterio
            from rasterio.mask import mask
            import rasterio.features

            with rasterio.open(self.dem_path) as src:
                # Ensure basin is in same CRS as DEM
                basin_gdf_reproj = basin_gdf.to_crs(src.crs)

                # Clip DEM to basin
                geoms = [mapping for mapping in basin_gdf_reproj.geometry]
                out_image, out_transform = mask(src, geoms, crop=True)
                elevation = out_image[0]

                # Remove nodata values
                nodata = src.nodata
                if nodata is not None:
                    elevation = elevation[elevation != nodata]

                # Calculate statistics
                features = {
                    "elevation_mean_m": float(np.mean(elevation)),
                    "elevation_min_m": float(np.min(elevation)),
                    "elevation_max_m": float(np.max(elevation)),
                    "elevation_range_m": float(np.max(elevation) - np.min(elevation))
                }

                # Slope calculation (simplified - would need proper implementation)
                # This is a placeholder
                features["mean_slope_deg"] = self._estimate_slope(elevation)

                return features

        except ImportError:
            # If rasterio not available, return placeholder values
            return {
                "elevation_mean_m": None,
                "elevation_min_m": None,
                "elevation_max_m": None,
                "elevation_range_m": None,
                "mean_slope_deg": None
            }
        except Exception as e:
            print(f"Warning: Could not extract DEM features: {e}")
            return {}

    def _estimate_slope(self, elevation: np.ndarray) -> float:
        """
        Estimate mean slope from elevation data.
        Simplified calculation - full version would use proper gradient computation.
        """
        if len(elevation) < 2:
            return 0.0

        # Simple approximation using elevation variability
        slope_proxy = np.std(elevation) / np.mean(elevation) * 100 if np.mean(elevation) > 0 else 0
        return float(slope_proxy)

    def get_feature_names(self) -> list:
        """Get list of feature names."""
        base_features = ["basin_area_km2", "centroid_lon", "centroid_lat"]
        dem_features = [
            "elevation_mean_m", "elevation_min_m", "elevation_max_m",
            "elevation_range_m", "mean_slope_deg"
        ]
        return base_features + (dem_features if self.dem_path else [])