"""
Instance Vectorizer

Converts Mask R-CNN instance predictions to high-quality vector polygons.

Each instance becomes one feature with complete attributes including:
- instance_id, class_id, class_name
- confidence score
- geometric properties (area, perimeter, compactness)
- bounding box coordinates

Pipeline (per instance):
1. Mathematical morphology (closing) to eliminate micro-holes
2. Region filtering (remove noise pixels)
3. Polygonization with rasterio.features.shapes() (topology-aware)
4. Topology-aware simplification with Shapely (preserve_topology=True)
5. Optional Chaikin smoothing for organic shapes

Note: YOLO Segment uses PolygonVectorizer instead (direct polygon handling).
"""

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


class InstanceVectorizer:
    """
    Convert Mask R-CNN instance segmentation predictions to professional-quality vector polygons.

    Processes binary masks through a quality pipeline including:
    - Morphological cleaning
    - Region filtering
    - Topology-aware polygonization
    - Simplification with topology preservation
    - Optional Chaikin smoothing

    Note: This vectorizer is designed for raster-based instance masks (Mask R-CNN).
    YOLO Segment predictions use PolygonVectorizer for direct polygon handling.

    Args:
        instances: List of instance dicts with 'mask_local', 'mask_offset', 'class_id', 'score', 'bbox'
        output_path: Path for output GeoPackage (.gpkg)
        model_info: Model configuration with class definitions
        geo_transform: GDAL geotransform (6-tuple)
        projection: WKT projection string
        simplify_tolerance: Simplification tolerance in pixels (default: 1.0)
        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 (default: 1)
        
    Example:
        >>> vectorizer = InstanceVectorizer(
        ...     instances=final_instances,
        ...     output_path='instances.gpkg',
        ...     model_info=model_info,
        ...     geo_transform=processor.geo_transform,
        ...     projection=processor.projection,
        ...     simplify_tolerance=1.0,
        ...     morph_kernel_size=3,
        ...     chaikin_iterations=1
        ... )
        >>> vectorizer.vectorize(progress_callback=callback)
    """
    
    def __init__(
        self,
        instances: List[Dict[str, Any]],
        output_path: str,
        model_info: Dict[str, Any],
        geo_transform: tuple,
        projection: str,
        simplify_tolerance: float = 1.0,
        morph_kernel_size: int = 3,
        min_region_pixels: int = 9,
        chaikin_iterations: int = 1,
        contour_approximation: str = 'simple'
    ):
        """
        Initialize instance vectorizer.
        
        Args:
            instances: List of instance dicts
            output_path: Output GeoPackage path
            model_info: Model configuration
            geo_transform: GDAL geotransform (6-tuple)
            projection: WKT projection string
            simplify_tolerance: Simplification tolerance in pixels
            morph_kernel_size: Morphological closing kernel size
            min_region_pixels: Minimum region size for filtering
            chaikin_iterations: Chaikin smoothing iterations
            contour_approximation: Legacy parameter (ignored, kept for compatibility)
        """
        self.instances = instances
        self.output_path = Path(output_path)
        self.model_info = model_info
        self.geo_transform = geo_transform
        self.projection = projection
        self.simplify_tolerance = simplify_tolerance
        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']
        
        # Build class lookup
        self.class_lookup = {c['id']: c for c in self.classes}
        
        # Get pixel size for area/perimeter calculations
        self.pixel_width = abs(geo_transform[1])
        self.pixel_height = abs(geo_transform[5])
        self.pixel_area_m2 = self.pixel_width * self.pixel_height
        
        print(f"[InstanceVectorizer] 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):
        """
        Convert all instances to vector polygons.
        
        Args:
            progress_callback: Optional callback(current, total)
            
        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')
            ds = driver.CreateDataSource(str(self.output_path))
            
            if ds is None:
                raise RuntimeError(f"Could not create output file: {self.output_path}")
            
            # Create spatial reference
            srs = osr.SpatialReference()
            srs.ImportFromWkt(self.projection)
            
            # Create layer
            layer = ds.CreateLayer(
                'instances',
                srs=srs,
                geom_type=ogr.wkbMultiPolygon
            )
            
            # Create attribute fields
            self._create_fields(layer)
            
            # Process each instance
            total_instances = len(self.instances)
            
            print(f"\n[InstanceVectorizer] Processing {total_instances} instances...")
            
            for idx, instance in enumerate(self.instances):
                self._process_instance(instance, idx, layer)
                
                # Progress callback
                if progress_callback:
                    progress_callback(idx + 1, total_instances)
            
            print(f"[InstanceVectorizer] ✅ Created {total_instances} instance polygons")
            
            # Cleanup
            ds = None
            
        except Exception as e:
            raise RuntimeError(f"Instance vectorization failed: {str(e)}")
    
    def _create_fields(self, layer: ogr.Layer):
        """
        Create attribute fields for instance features.
        
        Args:
            layer: OGR layer to add fields to
        """
        layer.CreateField(ogr.FieldDefn('instance_id', ogr.OFTInteger))
        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('confidence', ogr.OFTReal))
        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('bbox_xmin', ogr.OFTReal))
        layer.CreateField(ogr.FieldDefn('bbox_ymin', ogr.OFTReal))
        layer.CreateField(ogr.FieldDefn('bbox_xmax', ogr.OFTReal))
        layer.CreateField(ogr.FieldDefn('bbox_ymax', ogr.OFTReal))
    
    def _process_instance(self, instance: Dict[str, Any], instance_id: int, layer: ogr.Layer):
        """
        Process a single instance through the professional pipeline.
        
        Pipeline:
        1. Extract mask and position
        2. Apply morphological closing
        3. Filter small regions
        4. Polygonize with rasterio (topology-aware)
        5. Transform to global coordinates
        6. Simplify with topology preservation
        7. Optional Chaikin smoothing
        8. Add to output layer with attributes
        
        Args:
            instance: Instance dict with 'mask_local', 'mask_offset', 'class_id', 'score', 'bbox'
            instance_id: Unique instance ID (0-indexed, will be 1-indexed in output)
            layer: OGR layer to add feature to
        """
        # Extract instance data
        mask_local = instance['mask_local']  
        mask_offset = instance['mask_offset']
        class_id = instance['class_id']
        score = instance['score']
        bbox = instance['bbox']
        
        # Get class name
        class_info = self.class_lookup.get(class_id)
        class_name = class_info['name'] if class_info else f"Class_{class_id}"

        # Ensure mask is binary uint8
        mask_binary = (mask_local > 0).astype(np.uint8)
        
        # Step 1: Morphological closing (eliminate small holes)
        if self.morph_kernel_size > 0:
            structure = ndimage.generate_binary_structure(2, 2)
            mask_binary = ndimage.binary_closing(
                mask_binary,
                structure=structure,
                iterations=self.morph_kernel_size
            ).astype(np.uint8)
        
        # Step 2: Filter small regions (remove noise)
        if self.min_region_pixels > 0:
            mask_binary = self._filter_small_regions(mask_binary, self.min_region_pixels)
        
        # Check if anything remains after filtering
        if not np.any(mask_binary):
            print(f"[InstanceVectorizer] Instance {instance_id} removed by filtering, skipping.")
            return
        
        # Step 3: Create local transform for this mask
        # The mask is in local coordinates (offset from global origin)
        offset_y, offset_x = mask_offset
        
        # Create transform: maps mask pixel coords to global pixel coords
        local_transform = Affine(
            1.0, 0.0, offset_x,  # a, b, c
            0.0, 1.0, offset_y   # d, e, f
        )
        
        # Step 4: Polygonize (in global pixel coordinates)
        try:
            polygons = self._polygonize_mask(mask_binary, local_transform)
        except Exception as e:
            print(f"[InstanceVectorizer] Warning: Could not polygonize instance {instance_id}: {e}")
            return
        
        if not polygons:
            print(f"[InstanceVectorizer] Warning: Instance {instance_id} produced no valid polygons")
            return
        
        # Step 5: Merge all polygons for this instance (if multiple)
        # This can happen if the mask has multiple disconnected regions
        if len(polygons) == 1:
            polygon = polygons[0]
        else:
            # Merge multiple polygons into MultiPolygon
            polygon = unary_union(polygons)
        
        # Step 6: Transform from pixel coordinates to map coordinates
        polygon = self._transform_polygon_to_map_coords(polygon)
        
        # Step 7: 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)
        
        # Step 8: Optional Chaikin smoothing
        if self.chaikin_iterations > 0:
            polygon = self._chaikin_smooth(polygon, self.chaikin_iterations)
        
        # Validate geometry
        if not polygon.is_valid or polygon.is_empty:
            print(f"[InstanceVectorizer] Warning: Instance {instance_id} produced invalid geometry")
            return
        
        # Calculate geometric attributes
        area_m2 = polygon.area
        
        # Skip very small polygons (artifacts)
        if area_m2 < self.pixel_area_m2 * 0.5:
            return
        
        perimeter_m = polygon.length
        
        # Compactness
        if perimeter_m > 0:
            compactness = (4 * np.pi * area_m2) / (perimeter_m ** 2)
            compactness = min(compactness, 1.0)
        else:
            compactness = 0.0
        
        # Convert bounding box to map coordinates
        bbox_map = self._pixel_to_map_bbox(bbox)
        
        # Convert Shapely to OGR geometry
        ogr_geom = self._shapely_to_ogr(polygon)
        
        if ogr_geom is None:
            return
        
        # Create feature
        feature = ogr.Feature(layer.GetLayerDefn())
        feature.SetGeometry(ogr_geom)
        feature.SetField('instance_id', instance_id + 1)
        feature.SetField('class_id', int(class_id))
        feature.SetField('class_name', class_name)
        feature.SetField('confidence', float(score))
        feature.SetField('area_m2', area_m2)
        feature.SetField('perimeter_m', perimeter_m)
        feature.SetField('compactness', compactness)
        feature.SetField('bbox_xmin', bbox_map[0])
        feature.SetField('bbox_ymin', bbox_map[1])
        feature.SetField('bbox_xmax', bbox_map[2])
        feature.SetField('bbox_ymax', bbox_map[3])
        
        layer.CreateFeature(feature)
    
    def _filter_small_regions(self, mask: np.ndarray, min_pixels: int) -> np.ndarray:
        """
        Remove small connected regions from binary mask.
        
        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
        keep_regions = region_sizes >= min_pixels
        keep_regions[0] = False 
        
        # Filter mask
        filtered_mask = keep_regions[labeled]
        
        return filtered_mask.astype(np.uint8)
    
    def _polygonize_mask(self, mask: np.ndarray, transform: Affine) -> List[Polygon]:
        """
        Convert binary mask to Shapely polygons using rasterio.
        
        Uses rasterio.features.shapes() which natively handles:
        - Coordinate transformation
        - Topologically correct boundaries
        - Holes (interior rings)
        
        Args:
            mask: Binary mask (0/1) in local coordinates
            transform: Affine transform (mask pixels → global pixels)
            
        Returns:
            List of Shapely Polygon/MultiPolygon objects in global pixel coordinates
        """
        polygons = []
        
        # Extract shapes (polygons) from mask
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            
            for geom, value in shapes(mask, mask=(mask == 1), transform=transform):
                if value == 1:
                    poly = shape(geom)
                    
                    if poly.is_valid and not poly.is_empty:
                        polygons.append(poly)
        
        return polygons
    
    def _transform_polygon_to_map_coords(self, geom: Polygon) -> Polygon:
        """
        Transform polygon from global pixel coordinates to map coordinates.
        
        The rasterio polygonization gives us coordinates in global pixel space.
        We need to transform these to actual map coordinates using geo_transform.
        
        Args:
            geom: Shapely geometry in global pixel coordinates
            
        Returns:
            Shapely geometry in map coordinates
        """
        # Extract geotransform components

        origin_x = self.geo_transform[0]
        pixel_width = self.geo_transform[1]
        origin_y = self.geo_transform[3]
        pixel_height = self.geo_transform[5]
        
        transformed = affinity.affine_transform(
            geom,
            [pixel_width, 0, 0, pixel_height, origin_x, origin_y]
        )
        
        return transformed
    
    def _chaikin_smooth(self, geom: Polygon, iterations: int = 1) -> Polygon:
        """
        Apply Chaikin corner-cutting algorithm for smooth curves.
        
        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)
            
            if not smoothed_poly.is_valid:
                return geom
            
            return smoothed_poly
        
        except Exception:
            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)):
                p1 = coords[i]
                p2 = coords[(i + 1) % len(coords)]
                
                # Quarter points
                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 _pixel_to_map_bbox(self, bbox: np.ndarray) -> tuple:
        """
        Convert pixel bounding box to map coordinates.
        
        Args:
            bbox: [xmin, ymin, xmax, ymax] in pixels
            
        Returns:
            Tuple of (map_xmin, map_ymin, map_xmax, map_ymax)
        """
        xmin_map = self.geo_transform[0] + bbox[0] * self.geo_transform[1]
        ymin_map = self.geo_transform[3] + bbox[1] * self.geo_transform[5]
        xmax_map = self.geo_transform[0] + bbox[2] * self.geo_transform[1]
        ymax_map = self.geo_transform[3] + bbox[3] * self.geo_transform[5]
        
        return (xmin_map, ymin_map, xmax_map, ymax_map)
    
    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
            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)