"""
Vectorizer

Converts raster predictions to high-quality vector polygons using rasterio and Shapely.
Optimized pipeline for semantic segmentation with topology preservation.

Pipeline:
1. Mathematical morphology (closing) to eliminate micro-holes and isolated pixels
2. Region filtering (remove small noise regions before polygonization)
3. Polygonization with rasterio.features.shapes() (native topological correctness)
4. Topology-aware simplification with Shapely (preserve_topology=True)
5. Optional Chaikin smoothing for organic shapes (1-2 iterations)
"""

import numpy as np
import rasterio
from rasterio.features import shapes
from osgeo import ogr, osr
from pathlib import Path
from typing import Dict, Any, Optional, Callable, List, Tuple
from shapely.geometry import shape, Polygon, MultiPolygon
from scipy import ndimage
import warnings


class Vectorizer:
    """
    Convert raster predictions to professional-quality vector polygons.
    
    Uses rasterio + Shapely for topologically correct vectorization with:
    - Mathematical morphology for noise reduction
    - Region filtering before polygonization
    - Topology-aware simplification
    - Optional Chaikin smoothing
    
    Args:
        raster_path: Path to input class raster (GeoTIFF uint8)
        output_path: Path for output vector file (.gpkg)
        model_info: Model configuration with class definitions
        simplify_tolerance: Simplification tolerance in pixels (0 = no simplification)
        exclude_background: Skip class 0 (background)
        morph_kernel_size: Kernel size for morphological closing (default: 3)
        min_region_pixels: Minimum region size in pixels (default: 9)
        chaikin_iterations: Chaikin smoothing iterations, 0 to disable (default: 1)
        
    Example:
        >>> vectorizer = Vectorizer(
        ...     raster_path='prediction.tif',
        ...     output_path='polygons.gpkg',
        ...     model_info=model_info,
        ...     simplify_tolerance=1.0,
        ...     morph_kernel_size=3,
        ...     min_region_pixels=9,
        ...     chaikin_iterations=1
        ... )
        >>> vectorizer.vectorize(progress_callback=my_callback)
    """
    
    def __init__(
        self,
        raster_path: str,
        output_path: str,
        model_info: Dict[str, Any],
        simplify_tolerance: float = 1.0,
        exclude_background: bool = True,
        morph_kernel_size: int = 3,
        min_region_pixels: int = 9,
        chaikin_iterations: int = 1
    ):
        """
        Initialize vectorizer.
        
        Args:
            raster_path: Input raster path
            output_path: Output vector path (.gpkg)
            model_info: Model configuration
            simplify_tolerance: Simplification tolerance (pixels)
            exclude_background: Skip background class
            morph_kernel_size: Morphological closing kernel size (3 recommended)
            min_region_pixels: Minimum region size for filtering (9 = 3x3 pixels)
            chaikin_iterations: Chaikin smoothing iterations (0-2, 1 recommended)
        """
        self.raster_path = Path(raster_path)
        self.output_path = Path(output_path)
        self.model_info = model_info
        self.simplify_tolerance = simplify_tolerance
        self.exclude_background = exclude_background
        self.morph_kernel_size = morph_kernel_size
        self.min_region_pixels = min_region_pixels
        self.chaikin_iterations = chaikin_iterations
        
        # Get class information
        self.classes = model_info['classes']
        self.num_classes = model_info['num_classes']
        
        # Open raster with rasterio
        try:
            self.raster_dataset = rasterio.open(self.raster_path, 'r')
        except Exception as e:
            raise ValueError(f"Could not open raster: {self.raster_path}\nError: {e}")
        
        # Get raster metadata
        self.transform = self.raster_dataset.transform
        self.crs = self.raster_dataset.crs
        self.pixel_width = abs(self.transform.a)
        self.pixel_height = abs(self.transform.e)
        
        # Calculate pixel area in m² (assuming projected CRS in meters)
        self.pixel_area_m2 = self.pixel_width * self.pixel_height
        
        print(f"[Vectorizer] Initialized with professional pipeline")
        print(f"   Morphology kernel: {morph_kernel_size}x{morph_kernel_size}")
        print(f"   Min region: {min_region_pixels} pixels ({min_region_pixels * self.pixel_area_m2:.1f} m²)")
        print(f"   Simplification: {simplify_tolerance} pixels ({simplify_tolerance * self.pixel_width:.2f} m)")
        print(f"   Chaikin smoothing: {chaikin_iterations} iteration(s)")
    
    def vectorize(self, progress_callback: Optional[Callable[[int, int], None]] = None):
        """
        Run vectorization pipeline.
        
        Steps:
        1. Create output GeoPackage
        2. For each class:
           a. Extract binary mask
           b. Apply morphology + region filtering
           c. Polygonize with rasterio
           d. Simplify + smooth with Shapely
           e. Add to output layer
        3. Close datasets
        
        Args:
            progress_callback: Optional callback(current_class, total_classes)
        
        Raises:
            RuntimeError: If vectorization fails
        """
        try:
            # Remove existing output file
            if self.output_path.exists():
                self.output_path.unlink()
            
            # Create output GeoPackage
            driver = ogr.GetDriverByName('GPKG')
            vector_dataset = driver.CreateDataSource(str(self.output_path))
            
            if vector_dataset is None:
                raise RuntimeError(f"Could not create output file: {self.output_path}")
            
            # Create spatial reference from rasterio CRS
            srs = osr.SpatialReference()
            if self.crs is not None:
                srs.ImportFromWkt(self.crs.to_wkt())
            
            # Create layer
            layer = vector_dataset.CreateLayer(
                'predictions',
                srs=srs,
                geom_type=ogr.wkbMultiPolygon
            )
            
            # Create attribute fields
            self._create_fields(layer)
            
            # Determine which classes to process
            classes_to_process = []
            for class_info in self.classes:
                class_id = class_info['id']
                
                # Skip background if requested
                if self.exclude_background and class_id == 0:
                    continue
                
                classes_to_process.append(class_info)
            
            total_classes = len(classes_to_process)
            
            print(f"\n[Vectorizer] Processing {total_classes} classes...")
            
            # Process each class
            for idx, class_info in enumerate(classes_to_process):
                self._process_class(class_info, layer)
                
                # Progress callback
                if progress_callback:
                    progress_callback(idx + 1, total_classes)
            
            # Cleanup
            vector_dataset = None
            self.raster_dataset.close()
            
            print(f"[Vectorizer] ✅ Vectorization complete: {self.output_path}")
            
        except Exception as e:
            raise RuntimeError(f"Vectorization failed: {str(e)}")
    
    def _create_fields(self, layer: ogr.Layer):
        """
        Create attribute fields for output layer.
        
        Fields:
        - class_id: Integer class identifier
        - class_name: String class name
        - area_m2: Real area in square meters
        - perimeter_m: Real perimeter in meters
        - compactness: Real shape compactness index [0, 1]
        - hole_count: Integer number of interior holes
        
        Args:
            layer: OGR layer to add fields to
        """
        layer.CreateField(ogr.FieldDefn('class_id', ogr.OFTInteger))
        
        class_name_field = ogr.FieldDefn('class_name', ogr.OFTString)
        class_name_field.SetWidth(50)
        layer.CreateField(class_name_field)
        
        layer.CreateField(ogr.FieldDefn('area_m2', ogr.OFTReal))
        layer.CreateField(ogr.FieldDefn('perimeter_m', ogr.OFTReal))
        layer.CreateField(ogr.FieldDefn('compactness', ogr.OFTReal))
        layer.CreateField(ogr.FieldDefn('hole_count', ogr.OFTInteger))
    
    def _process_class(self, class_info: Dict[str, Any], layer: ogr.Layer):
        """
        Process a single class through the complete pipeline.
        
        Pipeline:
        1. Extract binary mask for this class
        2. Apply morphological closing (remove micro-holes)
        3. Filter small regions (noise removal)
        4. Polygonize with rasterio (topology-aware)
        5. Simplify with Shapely (preserve_topology=True)
        6. Optional Chaikin smoothing
        7. Add to output layer with attributes
        
        Args:
            class_info: Class configuration dict with 'id' and 'name'
            layer: OGR layer to add features to
        """
        class_id = class_info['id']
        class_name = class_info['name']
        
        # Read raster array
        raster_array = self.raster_dataset.read(1)
        
        # Create binary mask for this class
        class_mask = (raster_array == class_id).astype(np.uint8)
        
        # Check if class exists in raster
        if not np.any(class_mask):
            print(f"[Vectorizer] Class {class_id} ({class_name}) not found in raster, skipping.")
            return
        
        # Morphological closing (eliminate small holes and isolated pixels)
        if self.morph_kernel_size > 0:
            structure = ndimage.generate_binary_structure(2, 2)
            class_mask = ndimage.binary_closing(
                class_mask,
                structure=structure,
                iterations=self.morph_kernel_size
            ).astype(np.uint8)
        
        # Filter small regions (remove noise before polygonization)
        if self.min_region_pixels > 0:
            class_mask = self._filter_small_regions(class_mask, self.min_region_pixels)
        
        # Check if anything remains after filtering
        if not np.any(class_mask):
            print(f"[Vectorizer] Class {class_id} ({class_name}) removed by filtering, skipping.")
            return
        
        # Polygonize with rasterio (native geospatial coordinates)
        polygons = self._polygonize_mask(class_mask)
        
        if not polygons:
            print(f"[Vectorizer] Class {class_id} ({class_name}) produced no valid polygons, skipping.")
            return
        
        # Process each polygon (simplify + smooth)
        feature_count = 0
        
        for polygon in polygons:
            # Simplify with topology preservation
            if self.simplify_tolerance > 0:
                tolerance_map = self.simplify_tolerance * self.pixel_width
                polygon = polygon.simplify(tolerance_map, preserve_topology=True)
            
            # Optional Chaikin smoothing for organic shapes
            if self.chaikin_iterations > 0:
                polygon = self._chaikin_smooth(polygon, self.chaikin_iterations)
            
            # Validate geometry
            if not polygon.is_valid or polygon.is_empty:
                continue
            
            # Calculate geometric attributes
            area_m2 = polygon.area
            
            # Skip very small polygons (artifacts)
            if area_m2 < self.pixel_area_m2 * 0.5:
                continue
            
            perimeter_m = polygon.length
            
            # Compactness (shape index): 4πA/P² → [0, 1]
            # Circle = 1.0, complex shape → 0
            if perimeter_m > 0:
                compactness = (4 * np.pi * area_m2) / (perimeter_m ** 2)
                compactness = min(compactness, 1.0)
            else:
                compactness = 0.0
            
            # Count holes (interior rings)
            hole_count = 0
            if isinstance(polygon, Polygon):
                hole_count = len(polygon.interiors)
            elif isinstance(polygon, MultiPolygon):
                hole_count = sum(len(p.interiors) for p in polygon.geoms)
            
            # Convert Shapely to OGR geometry
            ogr_geom = self._shapely_to_ogr(polygon)
            
            if ogr_geom is None:
                continue
            
            # Create feature in output layer
            feature = ogr.Feature(layer.GetLayerDefn())
            feature.SetGeometry(ogr_geom)
            feature.SetField('class_id', class_id)
            feature.SetField('class_name', class_name)
            feature.SetField('area_m2', area_m2)
            feature.SetField('perimeter_m', perimeter_m)
            feature.SetField('compactness', compactness)
            feature.SetField('hole_count', hole_count)
            
            layer.CreateFeature(feature)
            feature_count += 1
        
        print(f"[Vectorizer] Class {class_id} ({class_name}): {feature_count} polygons created")
    
    def _filter_small_regions(self, mask: np.ndarray, min_pixels: int) -> np.ndarray:
        """
        Remove small connected regions from binary mask.
        
        Uses scipy.ndimage.label to identify connected components
        and filters those smaller than threshold.
        
        Args:
            mask: Binary mask (0/1)
            min_pixels: Minimum region size in pixels
            
        Returns:
            Filtered binary mask
        """
        # Label connected components
        labeled, num_features = ndimage.label(mask, structure=ndimage.generate_binary_structure(2, 2))
        
        if num_features == 0:
            return mask
        
        # Count pixels per region
        region_sizes = np.bincount(labeled.ravel())
        
        # Create mask of regions to keep (size >= min_pixels)
        # Note: region_sizes[0] is background, skip it
        keep_regions = region_sizes >= min_pixels
        keep_regions[0] = False  # Never keep background
        
        # Filter mask
        filtered_mask = keep_regions[labeled]
        
        return filtered_mask.astype(np.uint8)
    
    def _polygonize_mask(self, mask: np.ndarray) -> List[Polygon]:
        """
        Convert binary mask to Shapely polygons using rasterio.
        
        Uses rasterio.features.shapes() which natively handles:
        - Geospatial coordinate transformation
        - Topologically correct boundaries
        - Holes (interior rings)
        
        Args:
            mask: Binary mask (0/1) in raster coordinates
            
        Returns:
            List of Shapely Polygon/MultiPolygon objects in map coordinates
        """
        polygons = []
        
        # Extract shapes (polygons) from mask
        # shapes() returns (geometry, value) tuples
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")  # Ignore topology warnings from rasterio
            
            for geom, value in shapes(mask, mask=(mask == 1), transform=self.transform):
                if value == 1:  
                    # Convert GeoJSON-like dict to Shapely geometry
                    poly = shape(geom)
                    
                    # Validate and add
                    if poly.is_valid and not poly.is_empty:
                        polygons.append(poly)
        
        return polygons
    
    def _chaikin_smooth(self, geom: Polygon, iterations: int = 1) -> Polygon:
        """
        Apply Chaikin corner-cutting algorithm for smooth curves.
        
        Chaikin smoothing iteratively replaces each edge with two edges
        by cutting corners at 1/4 and 3/4 of the original edge length.
        
        This creates smooth, organic shapes without losing topology.
        Recommended: 1-2 iterations max to avoid over-smoothing.
        
        Args:
            geom: Input Shapely Polygon or MultiPolygon
            iterations: Number of smoothing iterations (1-2 recommended)
            
        Returns:
            Smoothed Shapely geometry
        """
        if iterations <= 0:
            return geom
        
        if isinstance(geom, MultiPolygon):
            # Smooth each polygon in MultiPolygon
            smoothed_polys = [self._chaikin_smooth(p, iterations) for p in geom.geoms]
            return MultiPolygon(smoothed_polys)
        
        if not isinstance(geom, Polygon):
            return geom
        
        # Smooth exterior ring
        exterior_coords = list(geom.exterior.coords)
        smoothed_exterior = self._chaikin_smooth_coords(exterior_coords, iterations)
        
        # Smooth interior rings (holes)
        smoothed_interiors = []
        for interior in geom.interiors:
            interior_coords = list(interior.coords)
            smoothed_interior = self._chaikin_smooth_coords(interior_coords, iterations)
            smoothed_interiors.append(smoothed_interior)
        
        # Create new polygon
        try:
            smoothed_poly = Polygon(smoothed_exterior, smoothed_interiors)
            
            # Validate result
            if not smoothed_poly.is_valid:
                # If invalid, return original (shouldn't happen with topology preservation)
                return geom
            
            return smoothed_poly
        
        except Exception:
            # If any error, return original
            return geom
    
    def _chaikin_smooth_coords(
        self,
        coords: List[Tuple[float, float]],
        iterations: int
    ) -> List[Tuple[float, float]]:
        """
        Apply Chaikin smoothing to coordinate list.
        
        Args:
            coords: List of (x, y) coordinate tuples
            iterations: Number of smoothing passes
            
        Returns:
            Smoothed coordinate list
        """
        if len(coords) < 3:
            return coords
        
        # Remove duplicate last coordinate if closed ring
        is_closed = coords[0] == coords[-1]
        if is_closed:
            coords = coords[:-1]
        
        # Apply Chaikin iterations
        for _ in range(iterations):
            new_coords = []
            
            for i in range(len(coords)):
                # Current point and next point (wrap around)
                p1 = coords[i]
                p2 = coords[(i + 1) % len(coords)]
                
                # Calculate quarter points (1/4 and 3/4 along edge)
                q = (
                    0.75 * p1[0] + 0.25 * p2[0],
                    0.75 * p1[1] + 0.25 * p2[1]
                )
                r = (
                    0.25 * p1[0] + 0.75 * p2[0],
                    0.25 * p1[1] + 0.75 * p2[1]
                )
                
                new_coords.extend([q, r])
            
            coords = new_coords
        
        # Close ring if needed
        if is_closed:
            coords.append(coords[0])
        
        return coords
    
    def _shapely_to_ogr(self, geom: Polygon) -> Optional[ogr.Geometry]:
        """
        Convert Shapely geometry to OGR geometry.
        
        Ensures output is always MultiPolygon for layer consistency.
        
        Args:
            geom: Shapely Polygon or MultiPolygon
            
        Returns:
            OGR MultiPolygon geometry, or None if conversion fails
        """
        try:
            # Convert to WKT then to OGR
            wkt = geom.wkt
            ogr_geom = ogr.CreateGeometryFromWkt(wkt)
            
            if ogr_geom is None:
                return None
            
            # Convert to MultiPolygon if needed (for layer consistency)
            if ogr_geom.GetGeometryName() == 'POLYGON':
                multi_geom = ogr.Geometry(ogr.wkbMultiPolygon)
                multi_geom.AddGeometry(ogr_geom)
                ogr_geom = multi_geom
            
            return ogr_geom
        
        except Exception:
            return None
    
    def get_output_path(self) -> str:
        """
        Get path to output vector file.
        
        Returns:
            Path to created GeoPackage
        """
        return str(self.output_path)


# Utility functions for advanced users

def optimize_vectorization_params(
    raster_path: str,
    sample_class_id: int = 1,
    target_quality: str = 'balanced'
) -> Dict[str, Any]:
    """
    Analyze raster and suggest optimal vectorization parameters.
    
    This utility function inspects the input raster and recommends
    parameters based on pixel size, class distribution, and target quality.
    
    Args:
        raster_path: Path to class raster
        sample_class_id: Class ID to analyze (default: 1, first non-background)
        target_quality: 'fast', 'balanced', or 'maximum' (default: 'balanced')
        
    Returns:
        Dictionary with recommended parameters for Vectorizer
        
    Example:
        >>> params = optimize_vectorization_params('prediction.tif', target_quality='maximum')
        >>> vectorizer = Vectorizer(raster_path, output_path, model_info, **params)
    """
    with rasterio.open(raster_path, 'r') as src:
        # Get pixel size
        transform = src.transform
        pixel_width = abs(transform.a)
        
        # Read sample
        data = src.read(1)
        class_mask = (data == sample_class_id)
        
        # Analyze mask characteristics
        class_pixel_count = np.sum(class_mask)
        
        # Labeled regions
        labeled, num_regions = ndimage.label(class_mask)
        
        if num_regions > 0:
            region_sizes = np.bincount(labeled.ravel())[1:]
            avg_region_size = np.mean(region_sizes)
        else:
            avg_region_size = 0
    
    # Determine parameters based on target quality
    if target_quality == 'fast':
        params = {
            'simplify_tolerance': 2.0,
            'morph_kernel_size': 2,
            'min_region_pixels': 16,
            'chaikin_iterations': 0
        }
    elif target_quality == 'maximum':
        params = {
            'simplify_tolerance': 0.5,
            'morph_kernel_size': 5,
            'min_region_pixels': 4,
            'chaikin_iterations': 2
        }
    else:  # balanced (default)
        params = {
            'simplify_tolerance': 1.0,
            'morph_kernel_size': 3,
            'min_region_pixels': 9,
            'chaikin_iterations': 1
        }
    
    print(f"\n[Vectorizer] Recommended parameters ({target_quality} quality):")
    print(f"   Pixel size: {pixel_width:.2f} m")
    print(f"   Class {sample_class_id} pixels: {class_pixel_count}")
    print(f"   Average region size: {avg_region_size:.1f} pixels")
    print(f"   → simplify_tolerance: {params['simplify_tolerance']}")
    print(f"   → morph_kernel_size: {params['morph_kernel_size']}")
    print(f"   → min_region_pixels: {params['min_region_pixels']}")
    print(f"   → chaikin_iterations: {params['chaikin_iterations']}")
    
    return params