"""
Polygon Vectorizer

Converts YOLO Segment polygon detections directly to vector format.
Optimized pipeline: polygon → geo transform → simplify → GeoPackage

No rasterization step (unlike InstanceVectorizer which is designed for CNN masks).
YOLO already provides clean polygons.

Key features:
- Direct polygon handling (no polygon→mask→polygon conversion)
- Coordinate transformation (pixel → geographic)
- Optional Douglas-Peucker simplification
- Attribute enrichment (class, confidence, geometric properties)
- GeoPackage export (.gpkg)

Used by: YOLO Segment
"""

import numpy as np
from osgeo import ogr, osr
from pathlib import Path
from typing import List, Dict, Any, Tuple, Optional, Callable
from shapely.geometry import Polygon, MultiPolygon
from shapely.validation import make_valid


class PolygonVectorizer:
    """
    Convert YOLO Segment polygon detections to vector polygons.
    
    Optimized for pre-polygonized data (YOLO masks.xy).
    
    Args:
        detections: List of detection dicts with 'polygon', 'bbox', 'class_id', 'score'
        output_path: Path for output GeoPackage (.gpkg)
        model_info: Model configuration with class definitions
        geo_transform: GDAL geotransform (6 parameters)
        projection: WKT projection string
        simplify_tolerance: Simplification tolerance in map units (default: None = no simplification)
        
    Example:
        >>> vectorizer = PolygonVectorizer(
        ...     detections=yolo_detections,
        ...     output_path='segments.gpkg',
        ...     model_info=model_info,
        ...     geo_transform=geotransform,
        ...     projection=crs_wkt,
        ...     simplify_tolerance=0.2
        ... )
        >>> vectorizer.vectorize()
    """
    
    def __init__(
        self,
        detections: List[Dict[str, Any]],
        output_path: str,
        model_info: Dict[str, Any],
        geo_transform: Tuple[float, ...],
        projection: str,
        simplify_tolerance: Optional[float] = None,
        chaikin_iterations: int = 0
    ):
        """Initialize polygon vectorizer."""
        self.detections = detections
        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.chaikin_iterations = chaikin_iterations
        
        # Get class information
        self.classes = {c['id']: c['name'] for c in model_info['classes']}
        
        # Get pixel size for area 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
    
    def vectorize(self, progress_callback: Optional[Callable[[int, int], None]] = None):
        """
        Convert detections to GeoPackage.
        
        Args:
            progress_callback: Optional callback(current, total)
            
        Raises:
            RuntimeError: If vectorization fails
        """
        try:
            # Create output directory if needed
            self.output_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Remove existing file
            if self.output_path.exists():
                self.output_path.unlink()
            
            # Create GeoPackage
            driver = ogr.GetDriverByName('GPKG')
            ds = driver.CreateDataSource(str(self.output_path))
            
            if ds is None:
                raise RuntimeError(f"Could not create GeoPackage: {self.output_path}")
            
            # Create spatial reference
            srs = osr.SpatialReference()
            srs.ImportFromWkt(self.projection)
            
            # Create layer (MultiPolygon for consistency)
            layer = ds.CreateLayer('segments', srs, ogr.wkbMultiPolygon)
            
            # Create attribute fields
            self._create_fields(layer)
            
            # Process each detection
            total = len(self.detections)
            feature_count = 0
            
            for idx, detection in enumerate(self.detections):
                feature = self._create_feature(detection, idx, layer)
                if feature is not None:
                    layer.CreateFeature(feature)
                    feature_count += 1
                    feature = None
                
                # Progress callback
                if progress_callback and (idx + 1) % max(1, total // 10) == 0:
                    progress_callback(idx + 1, total)
            
            # Final progress
            if progress_callback:
                progress_callback(total, total)
            
            # Close dataset
            ds = None
            
            print(f"[PolygonVectorizer] Created {feature_count} polygon features")
            
        except Exception as e:
            raise RuntimeError(f"Polygon vectorization failed: {str(e)}")
    
    def _create_fields(self, layer: ogr.Layer):
        """Create attribute fields in layer."""
        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))
    
    def _create_feature(
        self,
        detection: Dict[str, Any],
        instance_id: int,
        layer: ogr.Layer
    ) -> Optional[ogr.Feature]:
        """
        Create OGR feature from detection.
        
        Args:
            detection: Detection dict with 'polygon', 'bbox', 'class_id', 'score'
            instance_id: Unique instance ID
            layer: OGR layer for feature definition
            
        Returns:
            OGR feature or None if invalid
        """
        # Extract detection info
        polygon_coords = detection['polygon']
        class_id = detection['class_id']
        confidence = detection['score']
        
        # Get class name
        class_name = self.classes.get(class_id, f'class_{class_id}')
        
        # Convert to Shapely polygon in map coordinates
        shapely_poly = self._create_shapely_polygon(polygon_coords)
        
        if shapely_poly is None or shapely_poly.is_empty:
            return None
        
        # Ensure valid geometry
        if not shapely_poly.is_valid:
            shapely_poly = make_valid(shapely_poly)
            if shapely_poly.is_empty:
                return None
        
        # Simplify if requested
        if self.simplify_tolerance and self.simplify_tolerance > 0:
            shapely_poly = shapely_poly.simplify(
                self.simplify_tolerance,
                preserve_topology=True
            )

        # Smooth with Chaikin if requested
        if self.chaikin_iterations > 0:
            shapely_poly = self._chaikin_smooth(shapely_poly, self.chaikin_iterations)
        
        # Calculate geometric properties
        area_m2 = shapely_poly.area
        
        # Skip very small polygons
        if area_m2 < self.pixel_area_m2 * 0.5:
            return None
        
        perimeter_m = shapely_poly.length
        
        # Compactness: 4πA/P²
        if perimeter_m > 0:
            compactness = (4 * np.pi * area_m2) / (perimeter_m ** 2)
            compactness = min(compactness, 1.0)
        else:
            compactness = 0.0
        
        # Convert Shapely to OGR
        ogr_geom = self._shapely_to_ogr(shapely_poly)
        
        if ogr_geom is None:
            return None
        
        # 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(confidence))
        feature.SetField('area_m2', area_m2)
        feature.SetField('perimeter_m', perimeter_m)
        feature.SetField('compactness', compactness)
        
        return feature
    
    def _chaikin_smooth(self, geom: Polygon, iterations: int = 1) -> Polygon:
        """Apply Chaikin corner-cutting for smooth curves."""
        if iterations <= 0:
            return geom
        
        if isinstance(geom, 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
        exterior_coords = list(geom.exterior.coords)
        smoothed_exterior = self._chaikin_smooth_coords(exterior_coords, iterations)
        
        # Smooth interiors
        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)
        
        try:
            smoothed_poly = Polygon(smoothed_exterior, smoothed_interiors)
            return smoothed_poly if smoothed_poly.is_valid else geom
        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."""
        if len(coords) < 3:
            return coords
        
        # Remove duplicate last coordinate if closed
        is_closed = coords[0] == coords[-1]
        if is_closed:
            coords = coords[:-1]
        
        # Apply 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 _create_shapely_polygon(self, coords: List[List[float]]) -> Optional[Polygon]:
        """
        Create Shapely polygon from pixel coordinates.
        
        Transforms coordinates from pixel space to map space.
        
        Args:
            coords: List of [x, y] in pixel coordinates
            
        Returns:
            Shapely Polygon in map coordinates or None if invalid
        """
        if len(coords) < 3:
            return None
        
        # Transform each coordinate
        map_coords = []
        for x_pixel, y_pixel in coords:
            x_map, y_map = self._pixel_to_map(x_pixel, y_pixel)
            map_coords.append((x_map, y_map))
        
        try:
            polygon = Polygon(map_coords)
            return polygon if polygon.is_valid else None
        except Exception:
            return None
    
    def _pixel_to_map(
        self,
        x_pixel: float,
        y_pixel: float
    ) -> Tuple[float, float]:
        """
        Convert pixel coordinates to map coordinates.
        
        Args:
            x_pixel: X coordinate in pixels
            y_pixel: Y coordinate in pixels
            
        Returns:
            Tuple of (x_map, y_map) in CRS coordinates
        """
        gt = self.geo_transform
        
        x_map = gt[0] + x_pixel * gt[1] + y_pixel * gt[2]
        y_map = gt[3] + x_pixel * gt[4] + y_pixel * gt[5]
        
        return x_map, y_map
    
    def _shapely_to_ogr(self, geom: Polygon) -> Optional[ogr.Geometry]:
        """
        Convert Shapely geometry to OGR MultiPolygon.
        
        Args:
            geom: Shapely Polygon or MultiPolygon
            
        Returns:
            OGR MultiPolygon or None if conversion fails
        """
        try:
            wkt = geom.wkt
            ogr_geom = ogr.CreateGeometryFromWkt(wkt)
            
            if ogr_geom is None:
                return None
            
            # Convert to MultiPolygon 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 GeoPackage."""
        return str(self.output_path)