"""
Prediction Manager for QPredict

This module orchestrates the complete prediction pipeline in a QThread to
prevent blocking the QGIS UI. It coordinates RasterProcessor and Predictor,
handles progress updates, and manages the output raster layer.

Key features:
- Non-blocking execution in QThread
- Progress and log signals for UI updates
- Graceful cancellation support
- Automatic output layer addition to QGIS
- Error handling and cleanup
"""

from qgis.PyQt.QtCore import QThread, pyqtSignal
from qgis.PyQt.QtGui import QColor 
from qgis.core import (
    QgsProject, QgsRasterLayer, QgsMessageLog, Qgis, QgsVectorLayer,
    QgsSymbol, QgsRendererCategory, QgsCategorizedSymbolRenderer, 
    QgsMarkerSymbol, QgsLineSymbol, QgsFillSymbol  
)

import requests
import numpy as np
import base64
import io
import json
from typing import Dict, Any, Optional
from pathlib import Path
from osgeo import gdal, osr  

# Vectorizers (still needed for creating geo files from server predictions)
from ..core.vectorizer import Vectorizer
from ..core.box_vectorizer import BoxVectorizer
from ..core.instance_vectorizer import InstanceVectorizer

# Import client HTTP
from ..clients.qpredict_client import QPredictClient


# ============================================================================
# Helper functions for HTTP communication with server
# ============================================================================

def encode_tile(tile_array: np.ndarray) -> str:
    """Encode numpy array to base64 string for HTTP transfer."""
    buffer = io.BytesIO()
    np.save(buffer, tile_array)
    buffer.seek(0)
    return base64.b64encode(buffer.read()).decode('utf-8')


def decode_prediction(encoded_pred: str) -> np.ndarray:
    """Decode base64 string to numpy array from server response."""
    pred_bytes = base64.b64decode(encoded_pred)
    buffer = io.BytesIO(pred_bytes)
    return np.load(buffer, allow_pickle=False)

class PredictionManager(QThread):
    """
    Manage prediction pipeline in a separate thread.
    
    This class handles:
    - Pipeline orchestration (RasterProcessor + Predictor)
    - Progress tracking and UI updates
    - Graceful cancellation
    - Output layer management
    - Error handling
    
    Signals:
        progress_updated(int, int): Emitted with (current, total) tiles processed
        log_message(str): Emitted with log messages for UI display
        prediction_finished(str): Emitted when done with output path
        prediction_failed(str): Emitted on error with error message
    
    Args:
        model: PyTorch model (already loaded and on device)
        model_info: Model information from QMTPLoader
        config: Prediction configuration dictionary from UI
        
    Example:
        >>> manager = PredictionManager(model, model_info, config)
        >>> manager.progress_updated.connect(update_progress)
        >>> manager.log_message.connect(append_log)
        >>> manager.start()
    """
    
    # Signals for UI communication
    progress_updated = pyqtSignal(int, int) 
    log_message = pyqtSignal(str)
    prediction_finished = pyqtSignal(str)  
    prediction_failed = pyqtSignal(str)
    cleanup_finished = pyqtSignal() 
    
    def __init__(
        self,
        qmtp_path: str,
        model_info: Dict[str, Any],
        config: Dict[str, Any],
        parent=None
    ):
        """
        Initialize prediction manager.
        
        Args:
            qmtp_path: Path to .qmtp file (will be sent to server)
            model_info: Model information from QMTPLoader
            config: Prediction configuration
            parent: Parent QObject
        """
        super().__init__(parent)
        
        self.qmtp_path = qmtp_path
        self.model_info = model_info
        self.config = config
        
        # Create HTTP client (gets URL/token dynamically from QGeoAIClient)
        self.client = QPredictClient()
        
        # Stop flag for graceful cancellation
        self._should_stop = False
        
        # Store processor for cleanup
        self.raster_processor = None
        
        # Server model ID (obtained after loading)
        self.model_id = None
        self.divisibility = 1
    
 
    def run(self):
        """
        Main thread execution method - HTTP version.
        
        This method:
        1. Loads model on server via HTTP
        2. Creates RasterProcessor for tiles
        3. Sends tiles to server for inference
        4. Receives predictions and creates geo files
        5. Adds output to QGIS
        
        Runs in separate thread (non-blocking).
        """
        try:
            self.log_message.emit("\n🚀 Starting prediction pipeline (HTTP mode)...")
            
            # ================================================================
            # 1. CHECK SERVER CONNECTION
            # ================================================================
            self.log_message.emit(f"   Server: {self.client._server_url}")
            
            # ================================================================
            # 2. LOAD MODEL ON SERVER
            # ================================================================
            self.log_message.emit("\n📦 Loading model on server...")
            
            success, error, model_info_server = self.client.load_model(
                qmtp_path=self.qmtp_path,
                device=self.config.get('device', 'cuda'),
                cuda_device_id=self.config.get('cuda_device_id', 0)
            )
            
            if not success:
                raise RuntimeError(error)
            
            self.model_id = model_info_server['model_id']
            self.log_message.emit(f"   ✅ Model loaded: {model_info_server['architecture']}")
            self.divisibility = model_info_server.get('divisibility', 1)
            self.log_message.emit(f"   Task: {model_info_server['task']}")
            self.log_message.emit(f"   Device: {model_info_server['device']}")
            self.log_message.emit(f"   Model ID: {self.model_id}")
            if self.divisibility > 1:
                self.log_message.emit(f"   Divisibility constraint: {self.divisibility}x")
            
            # Check if stopped
            if self._should_stop:
                self._unload_model_from_server()
                self.log_message.emit("\n⚠️ Prediction cancelled by user")
                return
            
            # ================================================================
            # 3. CREATE RASTER PROCESSOR (tiles generation)
            # ================================================================
            self.log_message.emit("\n📊 Initializing raster processor...")
            
            from ..core.raster_processor import RasterProcessor
            
            self.raster_processor = RasterProcessor(
                raster_layer=self.config['input_layer'],
                tile_size=self.config['tile_size'],
                overlap=self.config['overlap'],
                padding=self.config['padding'],
                normalization=self.model_info['normalization'],
                extent_mode=self.config['extent_mode'],
                selection_extent=self.config.get('selection_extent'),
                divisibility=self.divisibility
            )
            
            # Log processor info
            info = self.raster_processor.get_raster_info()
            self.log_message.emit(f"   Raster: {info['width']}x{info['height']} pixels")
            self.log_message.emit(f"   Tiles: {info['num_tiles_x']}x{info['num_tiles_y']} = {info['total_tiles']} total")
            
            # ================================================================
            # 4. PREPARE TILES AND SEND TO SERVER
            # ================================================================
            self.log_message.emit("\n🔄 Running inference on server...")
            
            total_tiles = info['total_tiles']
            tiles_encoded = []
            tiles_metadata = []
            
            # Generate and encode tiles
            self.log_message.emit("   Encoding tiles...")
            for i, tile_data in enumerate(self.raster_processor.generate_tiles()):
                if self._should_stop:
                    self._unload_model_from_server()
                    self.log_message.emit("\n⚠️ Prediction cancelled by user")
                    return
                
                # Get tile tensor and convert to numpy
                tile_tensor = tile_data['tensor']  # (C, H, W)
                tile_np = tile_tensor.cpu().numpy()
                
                # Encode to base64
                tiles_encoded.append(encode_tile(tile_np))
                
                # Prepare metadata
                metadata = {
                    'tile_width': tile_data['tile_width'],
                    'tile_height': tile_data['tile_height'],
                    'pad_left': tile_data['pad_left'],
                    'pad_top': tile_data['pad_top'],
                    'out_x': tile_data['x_offset'] - self.raster_processor.extent_x_min,
                    'out_y': tile_data['y_offset'] - self.raster_processor.extent_y_min,
                }
                
                # YOLO-specific metadata
                if model_info_server['task'].startswith('yolo_'):
                    metadata['x_min'] = tile_data['x_offset']
                    metadata['y_min'] = tile_data['y_offset']
                    metadata['x_max'] = tile_data['x_offset'] + tile_data['tile_width']
                    metadata['y_max'] = tile_data['y_offset'] + tile_data['tile_height']
                
                tiles_metadata.append(metadata)
                
                # Progress
                if (i + 1) % max(1, total_tiles // 10) == 0:
                    self.log_message.emit(f"   Encoded {i + 1}/{total_tiles} tiles...")
            
            # ================================================================
            # 5. SEND PREDICTION REQUEST TO SERVER
            # ================================================================
            self.log_message.emit(f"\n🚀 Sending {len(tiles_encoded)} tiles to server...")
            
            output_shape = (info['height'], info['width'])
            
            # Prepare task-specific parameters
            task = model_info_server['task']
            predict_kwargs = {
                'output_format': self.config.get('output_format', 'class_raster'),
                'confidence_threshold': self.config.get('threshold', 0.5),
                'apply_colors': self.config.get('apply_colors', True),
            }
            
            if task == 'instance_segmentation':
                predict_kwargs['nms_threshold'] = self.config.get('nms_threshold', 0.5)
                predict_kwargs['overlap'] = self.config['overlap']
            elif task.startswith('yolo_'):
                predict_kwargs['confidence_threshold'] = self.config.get('confidence_threshold', 0.5)
                predict_kwargs['iou_threshold'] = self.config.get('iou_threshold', 0.5)
                predict_kwargs['overlap'] = self.config['overlap']
                predict_kwargs['simplify_tolerance'] = self.config.get('simplify_tolerance', 0.5)
            
            success, error, result = self.client.predict(
                model_id=self.model_id,
                tiles=tiles_encoded,
                tile_metadata=tiles_metadata,
                output_shape=output_shape,
                batch_size=self.config['batch_size'],
                **predict_kwargs
            )
            
            if not success:
                self._unload_model_from_server()
                raise RuntimeError(error)
            
            self.log_message.emit("   ✅ Inference complete!")
            
            # ================================================================
            # 6. PROCESS RESULTS AND CREATE GEO FILES
            # ================================================================
            self.log_message.emit("\n💾 Creating output files...")
            
            output_dir = Path(self.config['output_dir'])
            layer_name = self.config['layer_name']
            
            final_output_path = self._process_prediction_results(
                result, 
                task, 
                output_dir, 
                layer_name
            )
            
            # ================================================================
            # 7. ADD TO QGIS
            # ================================================================
            self.log_message.emit("\n📍 Adding layer to QGIS...")
            self._add_layer_to_qgis(str(final_output_path), layer_name)
            
            # ================================================================
            # 8. UNLOAD MODEL FROM SERVER
            # ================================================================
            self._unload_model_from_server()
            
            self.log_message.emit("\n✅ Pipeline complete!")
            self.prediction_finished.emit(str(final_output_path))
            
        except Exception as e:
            error_msg = f"Prediction failed: {str(e)}"
            self.log_message.emit(f"\n❌ {error_msg}")
            QgsMessageLog.logMessage(error_msg, "QPredict", Qgis.Critical)
            self.prediction_failed.emit(error_msg)
            
            try:
                if hasattr(self, 'model_id') and self.model_id is not None:
                    self._unload_model_from_server()
            except Exception as cleanup_error:
                self.log_message.emit(f"⚠️ Could not unload model: {cleanup_error}")
            
        finally:
            self._cleanup()
            self.cleanup_finished.emit()
    
    def _unload_model_from_server(self):
        """Unload model from server to free memory."""
        if self.model_id is None:
            return
        
        try:
            self.log_message.emit("\n🧹 Unloading model from server...")
            
            success, error = self.client.unload_model(self.model_id)
            
            if success:
                self.log_message.emit("   ✅ Model unloaded from server")
            else:
                self.log_message.emit(f"   ⚠️ Could not unload model: {error}")
                
        except Exception as e:
            self.log_message.emit(f"   ⚠️ Error unloading model: {e}")
    
    def _process_prediction_results(
        self, 
        result: dict, 
        task: str, 
        output_dir: Path, 
        layer_name: str
    ) -> Path:
        """
        Process prediction results from server and create geo files.
        
        Args:
            result: Server response with prediction data
            task: Task type (semantic_segmentation, instance_segmentation, yolo_*)
            output_dir: Output directory
            layer_name: Layer name
            
        Returns:
            Path to final output file
        """
        if task == 'semantic_segmentation':
            return self._process_semantic_results(result, output_dir, layer_name)
        elif task == 'instance_segmentation':
            return self._process_instance_results(result, output_dir, layer_name)
        elif task.startswith('yolo_'):
            return self._process_yolo_results(result, task, output_dir, layer_name)
        else:
            raise ValueError(f"Unknown task: {task}")
    
    def _process_semantic_results(
        self, 
        result: dict, 
        output_dir: Path, 
        layer_name: str
    ) -> Path:
        """Process semantic segmentation results."""
        # Decode prediction array
        prediction = decode_prediction(result['prediction'])
        
        # Create GeoTIFF
        output_path = output_dir / f"{layer_name}.tif"
        self._create_geotiff(prediction, output_path, result)

        self.log_message.emit(f"   ✅ GeoTIFF created: {output_path}")

        # Create QML style file for reliability maps
        if result['output_format'] == 'reliability_map' and self.config.get('apply_colors', True):
            reliability_type = result.get('reliability_type', 'Uncertainty')
            num_classes = len(result.get('classes', []))
            if num_classes > 0:
                self._create_qml_style_file(output_path, reliability_type, num_classes)
        
        # Vectorize if needed
        if self.config['output_format'] == 'vector':
            vector_path = output_dir / f"{layer_name}.gpkg"
            self.log_message.emit("\n🔷 Vectorizing...")
            
            vectorizer = Vectorizer(
                raster_path=str(output_path),
                output_path=str(vector_path),
                model_info=self.model_info,
                simplify_tolerance=self.config.get('simplify_tolerance', 1.0),
                exclude_background=self.config.get('exclude_background', True)
            )
            
            vectorizer.vectorize(progress_callback=self._on_vectorization_progress)
            self.log_message.emit(f"   ✅ Vector created: {vector_path}")
            
            return vector_path
        
        return output_path
    
    def _process_instance_results(
        self, 
        result: dict, 
        output_dir: Path, 
        layer_name: str
    ) -> Path:
        """Process instance segmentation results."""
        
        # Parse instances JSON
        instances = json.loads(result['prediction'])
        
        self.log_message.emit(f"   Received {len(instances)} instances from server")
        
        # Convert JSON lists back to numpy arrays
        for inst in instances:
            # Convert bbox
            if 'bbox' in inst and isinstance(inst['bbox'], list):
                inst['bbox'] = np.array(inst['bbox'])
            
            # Convert mask_local (if present)
            if 'mask_local' in inst and isinstance(inst['mask_local'], list):
                inst['mask_local'] = np.array(inst['mask_local'], dtype=np.uint8)
            
            # Convert mask (if present - for full masks)
            if 'mask' in inst and isinstance(inst['mask'], list):
                inst['mask'] = np.array(inst['mask'], dtype=np.uint8)
            
            # Convert centroid (if present)
            if 'centroid' in inst and isinstance(inst['centroid'], list):
                inst['centroid'] = np.array(inst['centroid'])
        
        # Adjust geotransform for extent offset
        adjusted_geo_transform = list(self.raster_processor.geo_transform)
        adjusted_geo_transform[0] += self.raster_processor.extent_x_min * adjusted_geo_transform[1]
        adjusted_geo_transform[3] += self.raster_processor.extent_y_min * adjusted_geo_transform[5]
        
        # Check output format from config
        if self.config.get('output_format') == 'class_raster':
            # Create class raster from instances
            self.log_message.emit("\n🔷 Rasterizing instances...")
            class_raster = self._rasterize_instances(instances)
            
            # Create GeoTIFF
            output_path = output_dir / f"{layer_name}.tif"
            
            # Prepare result dict for _create_geotiff
            raster_result = {
                'output_format': 'class_raster',
                'classes': result.get('classes', [])
            }

            self._create_geotiff(class_raster, output_path, raster_result)
            self.log_message.emit(f"   ✅ Class raster created: {output_path}")
            
            return output_path
        else:
            # Create GeoPackage with InstanceVectorizer
            output_path = output_dir / f"{layer_name}.gpkg"
            
            try:
                vectorizer = InstanceVectorizer(
                instances=instances,
                output_path=str(output_path),
                model_info=self.model_info,
                geo_transform=tuple(adjusted_geo_transform),
                projection=self.raster_processor.projection,
                simplify_tolerance=self.config.get('simplify_tolerance', 1.0),
                morph_kernel_size=self.config.get('morph_kernel_size', 3),
                min_region_pixels=self.config.get('min_region_pixels', 9),
                chaikin_iterations=self.config.get('chaikin_iterations', 1)
            )
                
                vectorizer.vectorize(progress_callback=self._on_vectorization_progress)
                self.log_message.emit(f"   ✅ GeoPackage created: {output_path}")
                
            except Exception as e:
                raise RuntimeError(f"Instance vectorization failed: {str(e)}")
            
            return output_path
        
    def _process_yolo_results(self, result, task, output_dir, layer_name):
        """Process YOLO results - route to appropriate vectorizer."""
        detections = json.loads(result['prediction'])
        
        self.log_message.emit(f"   Received {len(detections)} detections from server")
        
        # Route by YOLO task type
        if task == 'yolo_segment':
            return self._process_yolo_segment(detections, output_dir, layer_name)
        else:
            return self._process_yolo_boxes(detections, output_dir, layer_name)

    def _process_yolo_boxes(self, detections, output_dir, layer_name):
        """Process YOLO Detection/OBB with BoxVectorizer."""
        # Transform to BoxVectorizer format
        transformed_detections = []
        for det in detections:
            det_type = det['type']
            
            if det_type == 'obb':
                geometry = {'type': 'obb', 'corners': det['corners']}
            elif det_type == 'bbox':
                geometry = {'type': 'bbox', 'coords': det['bbox']}
            else:
                self.log_message.emit(f"   ⚠️ Unknown detection type: {det_type}")
                continue
            
            transformed_detections.append({
                'geometry': geometry,
                'class_id': det['class_id'],
                'confidence': det['score']
            })
        
        # Create GeoPackage with BoxVectorizer
        output_path = output_dir / f"{layer_name}.gpkg"
        
        vectorizer = BoxVectorizer(
            detections=transformed_detections,
            output_path=str(output_path),
            model_info=self.model_info,
            geo_transform=self.raster_processor.geo_transform,
            projection=self.raster_processor.projection
        )
        
        vectorizer.vectorize()
        self.log_message.emit(f"   ✅ GeoPackage created: {output_path}")
        
        return output_path

    def _process_yolo_segment(self, detections, output_dir, layer_name):
        """Process YOLO Segment with PolygonVectorizer (optimized)."""
        from ..core.polygon_vectorizer import PolygonVectorizer
        
        # YOLO Segment detections: [{'type': 'polygon', 'polygon': [[x,y],...], 'class_id': ..., 'score': ...}]
        # polygon coords are in global raster pixel coordinates
        
        # Adjust coordinates relative to extent (like other YOLO tasks)
        for det in detections:
            polygon_coords = det['polygon']
            for coord in polygon_coords:
                coord[0] -= self.raster_processor.extent_x_min
                coord[1] -= self.raster_processor.extent_y_min
        
        # Adjust geotransform for extent offset
        adjusted_geo_transform = list(self.raster_processor.geo_transform)
        adjusted_geo_transform[0] += self.raster_processor.extent_x_min * adjusted_geo_transform[1]
        adjusted_geo_transform[3] += self.raster_processor.extent_y_min * adjusted_geo_transform[5]
        
        # Create GeoPackage with PolygonVectorizer
        output_path = output_dir / f"{layer_name}.gpkg"
        
        try:
            vectorizer = PolygonVectorizer(
                detections=detections,
                output_path=str(output_path),
                model_info=self.model_info,
                geo_transform=tuple(adjusted_geo_transform),
                projection=self.raster_processor.projection,
                simplify_tolerance=self.config.get('simplify_tolerance', 0.2),  # In map units
                chaikin_iterations=self.config.get('chaikin_iterations', 1)
            )
            
            vectorizer.vectorize(progress_callback=self._on_vectorization_progress)
            self.log_message.emit(f"   ✅ GeoPackage created: {output_path}")
            
        except Exception as e:
            raise RuntimeError(f"YOLO Segment vectorization failed: {str(e)}")
        
        return output_path

    def _rasterize_instances(self, instances: list) -> np.ndarray:
        """
        Rasterize instances to class raster (same logic as initial version).
        
        Args:
            instances: List of instances from server
            
        Returns:
            Numpy array (H, W) with class IDs
        """
        
        # Get output shape from raster processor extent
        height = self.raster_processor.extent_height
        width = self.raster_processor.extent_width
        
        # Initialize class raster (background = 0)
        class_raster = np.zeros((height, width), dtype=np.uint8)
        confidence_map = np.zeros((height, width), dtype=np.float32)
        
        # Sort instances by confidence (ascending, so highest conf overwrites)
        instances_sorted = sorted(instances, key=lambda x: x['score'])
        
        # Rasterize each instance
        for instance in instances_sorted:
            mask_local = instance['mask_local']  
            mask_offset = instance['mask_offset']  # (y_offset, x_offset)
            class_id = instance['class_id']
            score = instance['score']
            
            # Get mask dimensions and position
            mask_h, mask_w = mask_local.shape
            offset_y, offset_x = mask_offset
            
            # Calculate destination in class raster (with clipping)
            dst_y1 = max(0, offset_y)
            dst_y2 = min(height, offset_y + mask_h)
            dst_x1 = max(0, offset_x)
            dst_x2 = min(width, offset_x + mask_w)
            
            # Calculate source in local mask (if clipped)
            src_y1 = dst_y1 - offset_y
            src_y2 = src_y1 + (dst_y2 - dst_y1)
            src_x1 = dst_x1 - offset_x
            src_x2 = src_x1 + (dst_x2 - dst_x1)
            
            # Get the mask region
            mask_region = mask_local[src_y1:src_y2, src_x1:src_x2]
            
            # Get corresponding regions in output rasters
            class_region = class_raster[dst_y1:dst_y2, dst_x1:dst_x2]
            conf_region = confidence_map[dst_y1:dst_y2, dst_x1:dst_x2]
            
            # Update where mask is True AND score is higher
            update_mask = (mask_region > 0) & (score > conf_region)
            
            class_region[update_mask] = class_id
            conf_region[update_mask] = score
            
            # Write back (necessary for numpy arrays)
            class_raster[dst_y1:dst_y2, dst_x1:dst_x2] = class_region
            confidence_map[dst_y1:dst_y2, dst_x1:dst_x2] = conf_region
        
        return class_raster

    def _create_geotiff(self, prediction: np.ndarray, output_path: Path, result: dict):
        """Create GeoTIFF from prediction array."""
        
        # Get raster info
        height, width = prediction.shape
        
        # Create GeoTIFF
        driver = gdal.GetDriverByName('GTiff')
        
        # Determine data type and band count
        if result['output_format'] == 'reliability_map':
            dtype = gdal.GDT_Float32
            num_bands = 1
        else:
            dtype = gdal.GDT_Byte
            num_bands = 1
        
        dataset = driver.Create(
            str(output_path),
            width,
            height,
            num_bands,
            dtype,
            options=['COMPRESS=LZW', 'TILED=YES']
        )
        
        if dataset is None:
            raise RuntimeError(f"Could not create GeoTIFF: {output_path}")
        
        # Set geotransform
        geo_transform = list(self.raster_processor.geo_transform)
        geo_transform[0] += self.raster_processor.extent_x_min * geo_transform[1]
        geo_transform[3] += self.raster_processor.extent_y_min * geo_transform[5]
        dataset.SetGeoTransform(geo_transform)
        
        # Set projection
        srs = osr.SpatialReference()
        srs.ImportFromWkt(self.raster_processor.projection)
        dataset.SetProjection(srs.ExportToWkt())
        
        # Get band
        band = dataset.GetRasterBand(1)
        
        # Add band description for reliability maps
        if result['output_format'] == 'reliability_map':
            reliability_type = result.get('reliability_type', 'Uncertainty')
            if 'Confidence' in reliability_type:
                band.SetDescription('Confidence (Max Softmax Probability)')
                band.SetUnitType('Probability')
            else:
                band.SetDescription('Uncertainty (Shannon Entropy)')
                band.SetUnitType('Entropy (nats)')
        
        # Add color table if class raster
        if result['output_format'] == 'class_raster' and 'classes' in result:
            # Set color interpretation
            band.SetColorInterpretation(gdal.GCI_PaletteIndex)
            
            # Create and set color table
            color_table = gdal.ColorTable()
            
            for class_info in result['classes']:
                class_id = class_info['id']
                color = class_info['color']
                
                if class_id == 0 and color == [0, 0, 0]:
                    color_entry = (1, 1, 1, 255)
                else:
                    color_entry = (color[0], color[1], color[2], 255)
                
                color_table.SetColorEntry(class_id, color_entry)
            
            band.SetColorTable(color_table)
            
            # Add category names
            category_names = [c['name'] for c in result['classes']]
            band.SetCategoryNames(category_names)
            
            # Create Raster Attribute Table
            self._create_rat(band, prediction, result['classes'])
        
        # Write data
        band.WriteArray(prediction)
        
        # Cleanup
        band.FlushCache()
        dataset.FlushCache()
        dataset = None

    def _create_rat(self, band, prediction: np.ndarray, classes: list):
        """
        Create Raster Attribute Table for proper QGIS rendering.
        
        Args:
            band: GDAL raster band
            prediction: Prediction array (to get unique values)
            classes: List of class definitions from model_info
        """
        
        
        # Get unique values actually present in the raster
        unique_values = np.unique(prediction)
        
        # Create RAT
        rat = gdal.RasterAttributeTable()
        
        # Add columns
        rat.CreateColumn('Value', gdal.GFT_Integer, gdal.GFU_MinMax)
        rat.CreateColumn('Class', gdal.GFT_String, gdal.GFU_Name)
        rat.CreateColumn('Red', gdal.GFT_Integer, gdal.GFU_Red)
        rat.CreateColumn('Green', gdal.GFT_Integer, gdal.GFU_Green)
        rat.CreateColumn('Blue', gdal.GFT_Integer, gdal.GFU_Blue)
        rat.CreateColumn('Alpha', gdal.GFT_Integer, gdal.GFU_Alpha)
        
        # Set row count
        rat.SetRowCount(len(unique_values))
        
        # Fill RAT
        for row_idx, class_id in enumerate(unique_values):
            # Find class info
            class_info = next((c for c in classes if c['id'] == class_id), None)
            
            if class_info:
                rat.SetValueAsInt(row_idx, 0, int(class_id))
                rat.SetValueAsString(row_idx, 1, class_info['name'])
                
                # Workaround: Use (1,1,1) instead of (0,0,0) for background
                if class_id == 0 and class_info['color'] == [0, 0, 0]:
                    rat.SetValueAsInt(row_idx, 2, 1)
                    rat.SetValueAsInt(row_idx, 3, 1)
                    rat.SetValueAsInt(row_idx, 4, 1)
                else:
                    rat.SetValueAsInt(row_idx, 2, class_info['color'][0])
                    rat.SetValueAsInt(row_idx, 3, class_info['color'][1])
                    rat.SetValueAsInt(row_idx, 4, class_info['color'][2])
                
                rat.SetValueAsInt(row_idx, 5, 255)
        
        # Set RAT to band
        band.SetDefaultRAT(rat)

    def _create_qml_style_file(self, output_path: Path, reliability_type: str, num_classes: int):
        """
        Create QGIS .qml style file for automatic styling of reliability maps.
        
        Args:
            output_path: Path to the GeoTIFF file
            reliability_type: 'Confidence' or 'Uncertainty'
            num_classes: Number of classes (for uncertainty max value)
        """
        
        qml_path = output_path.with_suffix('.qml')
        
        # Determine min/max and color ramp
        if 'Confidence' in reliability_type:
            min_value = 0.0
            max_value = 1.0
            color_ramp = [
                (0.0, '#d73027'),
                (0.25, '#fc8d59'),
                (0.5, '#fee090'),
                (0.75, '#91cf60'),
                (1.0, '#1a9850')
            ]
            label = 'Confidence'
        else:
            min_value = 0.0
            max_value = np.log(num_classes)
            color_ramp = [
                (0.0, '#1a9850'),
                (0.25, '#91cf60'),
                (0.5, '#fee090'),
                (0.75, '#fc8d59'),
                (1.0, '#d73027')
            ]
            label = 'Uncertainty'
        
        # Generate QML XML
        qml_content = f'''<!DOCTYPE qgis PUBLIC 'http://mrcc.com/qgis.dtd' 'SYSTEM'>
<qgis version="3.28.0" styleCategories="AllStyleCategories">
<pipe-data-defined-properties>
    <Option type="Map">
    <Option value="" name="name" type="QString"/>
    <Option name="properties"/>
    <Option value="collection" name="type" type="QString"/>
    </Option>
</pipe-data-defined-properties>
<pipe>
    <provider>
    <resampling enabled="false" maxOversampling="2" zoomedInResamplingMethod="nearestNeighbour" zoomedOutResamplingMethod="nearestNeighbour"/>
    </provider>
    <rasterrenderer alphaBand="-1" band="1" classificationMax="{max_value}" classificationMin="{min_value}" nodataColor="" opacity="1" type="singlebandpseudocolor">
    <rasterTransparency/>
    <minMaxOrigin>
        <limits>None</limits>
        <extent>WholeRaster</extent>
        <statAccuracy>Estimated</statAccuracy>
        <cumulativeCutLower>0.02</cumulativeCutLower>
        <cumulativeCutUpper>0.98</cumulativeCutUpper>
        <stdDevFactor>2</stdDevFactor>
    </minMaxOrigin>
    <rastershader>
        <colorrampshader classificationMode="1" clip="0" colorRampType="INTERPOLATED" labelPrecision="4" maximumValue="{max_value}" minimumValue="{min_value}">
'''
        
        # Add color ramp entries
        for ratio, color in color_ramp:
            value = min_value + ratio * (max_value - min_value)
            qml_content += f'      <item alpha="255" color="{color}" label="{value:.4f}" value="{value}"/>\n'
        
        qml_content += '''      <ramp>
            <Option type="Map">
            <Option value="0" name="alpha" type="QString"/>
            <Option value="gradient" name="color_type" type="QString"/>
            <Option value="0" name="discrete" type="QString"/>
'''
        
        for i, (ratio, color) in enumerate(color_ramp):
            qml_content += f'          <Option value="{ratio};{color}" name="rampType" type="QString"/>\n'
        
        qml_content += '''          <Option value="rgb" name="spec" type="QString"/>
            </Option>
            <prop k="alpha" v="255"/>
            <prop k="color_type" v="gradient"/>
            <prop k="discrete" v="0"/>
            <prop k="rampType" v="gradient"/>
            <prop k="spec" v="rgb"/>
        </ramp>
        </colorrampshader>
    </rastershader>
    </rasterrenderer>
    <brightnesscontrast brightness="0" contrast="0" gamma="1"/>
    <huesaturation colorizeBlue="128" colorizeGreen="128" colorizeOn="0" colorizeRed="255" colorizeStrength="100" grayscaleMode="0" invertColors="0" saturation="0"/>
    <rasterresampler maxOversampling="2"/>
    <resamplingStage>resamplingFilter</resamplingStage>
</pipe>
<blendMode>0</blendMode>
</qgis>
'''
        
        # Write QML file
        try:
            with open(qml_path, 'w', encoding='utf-8') as f:
                f.write(qml_content)
            self.log_message.emit(f"   🎨 QML style file created: {qml_path}")
        except Exception as e:
            self.log_message.emit(f"   ⚠️ Could not create QML style file: {e}")
    
    def _on_progress(self, current: int, total: int):
        """
        Progress callback for predictor.
        
        Args:
            current: Number of tiles processed
            total: Total number of tiles
        """
        # Check stop flag
        if self._should_stop:
            return
        
        # Emit progress signal
        self.progress_updated.emit(current, total)
        
        # Log progress every 10% or at milestones
        percentage = (current / total) * 100
        if current % max(1, total // 10) == 0 or current == total:
            self.log_message.emit(f"   Progress: {current}/{total} tiles ({percentage:.1f}%)")

    def _on_vectorization_progress(self, current: int, total: int):
        """
        Progress callback for vectorizer.
        
        Args:
            current: Number of classes processed
            total: Total number of classes
        """
        # Check stop flag
        if self._should_stop:
            return
        
        # Emit progress signal (continue from where prediction left off)
        self.progress_updated.emit(current, total)
        
        # Log progress
        self.log_message.emit(f"   Class {current}/{total} processed")

    def _apply_categorized_symbology(self, vector_layer, skip_background=True):
        """Apply categorized symbology to vector layer based on class_id."""
        
        try:
            # Get class information from model
            classes = self.model_info['classes']
            
            # Create categories
            categories = []
            
            for class_info in classes:
                class_id = class_info['id']
                class_name = class_info['name']
                color = class_info['color']
                
                # Skip background conditionally
                if skip_background and class_id == 0:
                    continue
                
                # Create symbol based on geometry type
                geom_type = vector_layer.geometryType()
                
                if geom_type == 0:  # Point
                    symbol = QgsMarkerSymbol.createSimple({})
                elif geom_type == 1:  # Line
                    symbol = QgsLineSymbol.createSimple({})
                else:  # Polygon (most common for segmentation)
                    symbol = QgsFillSymbol.createSimple({})
                
                # Set symbol color (RGB from model)
                qcolor = QColor(color[0], color[1], color[2])
                symbol.setColor(qcolor)
                
                # Create category (value, symbol, label)
                category = QgsRendererCategory(
                    class_id, 
                    symbol,
                    class_name
                )
                
                categories.append(category)
            
            # Create categorized renderer
            renderer = QgsCategorizedSymbolRenderer(
                'class_id',
                categories
            )
            
            # Apply renderer to layer
            vector_layer.setRenderer(renderer)
            
            # Trigger repaint
            vector_layer.triggerRepaint()
            
            self.log_message.emit(f"   🎨 Categorized symbology applied ({len(categories)} classes)")
            
        except Exception as e:
            self.log_message.emit(f"   ⚠️ Could not apply symbology: {str(e)}")
    
    def _add_layer_to_qgis(self, output_path: str, layer_name: str):
        """
        Add output layer (raster or vector) to QGIS project.
        
        Routes to task-specific methods for proper symbology.
        
        Args:
            output_path: Path to output file (.tif or .gpkg)
            layer_name: Name for the layer in QGIS
        """
       
        try:
            # Determine if vector output (by extension)
            is_vector = Path(output_path).suffix.lower() == '.gpkg'
            
            if is_vector:
                # Vector output - route to task-specific handler
                task_type = self.model_info['configuration']['model'].get('task', 'semantic_segmentation')

                if task_type.startswith('yolo_'):
                    # Check if it's YOLO Segment (instance segmentation)
                    if task_type == 'yolo_segment':
                        # YOLO Segment → use instance symbology (by class_id)
                        layer = QgsVectorLayer(output_path, layer_name, "ogr")
                        
                        if layer.isValid():
                            self._apply_categorized_symbology(layer, skip_background=False)
                            QgsProject.instance().addMapLayer(layer)
                            
                            feature_count = layer.featureCount()
                            self.log_message.emit(f"   ✅ Vector layer '{layer_name}' added to project")
                            self.log_message.emit(f"   📊 {feature_count} instances loaded")
                        else:
                            self.log_message.emit(f"   ⚠️ Vector layer created but could not be added to project")
                    
                    else:
                        # YOLO Detection/OBB → use box symbology (by class_name)
                        layer = QgsVectorLayer(output_path, layer_name, "ogr")
                        
                        if layer.isValid():
                            # Apply categorized symbology by class_name
                            
                            classes = self.model_info['classes']
                            categories = []
                            
                            for class_info in classes:
                                class_name = class_info['name']
                                color = class_info['color']
                                
                                symbol = QgsFillSymbol.createSimple({
                                    'color': '255,255,255,0',
                                    'outline_width': '0.5',
                                    'outline_style': 'solid'
                                })
                                
                                qcolor = QColor(color[0], color[1], color[2])
                                symbol.setColor(QColor(255, 255, 255, 0))
                                symbol.symbolLayer(0).setStrokeColor(qcolor)
                                
                                category = QgsRendererCategory(
                                    class_name,
                                    symbol,
                                    class_name
                                )
                                
                                categories.append(category)
                            
                            renderer = QgsCategorizedSymbolRenderer('class_name', categories)
                            layer.setRenderer(renderer)
                            layer.triggerRepaint()
                            
                            QgsProject.instance().addMapLayer(layer)
                            
                            feature_count = layer.featureCount()
                            self.log_message.emit(f"   🎨 YOLO categorized symbology applied ({len(categories)} classes)")
                            self.log_message.emit(f"   ✅ Vector layer '{layer_name}' added to project")
                            self.log_message.emit(f"   📊 {feature_count} detections loaded")
                        else:
                            self.log_message.emit(f"   ⚠️ YOLO vector layer created but could not be added to project")

                else:
                    # PyTorch semantic/instance segmentation
                    layer = QgsVectorLayer(output_path, layer_name, "ogr")
                    
                    if layer.isValid():
                        # Apply categorized symbology before adding to project
                        self._apply_categorized_symbology(layer, skip_background=self.config.get('exclude_background', True))
                        
                        # Add to project
                        QgsProject.instance().addMapLayer(layer)
                        
                        # Get feature count for log
                        feature_count = layer.featureCount()
                        self.log_message.emit(f"   ✅ Vector layer '{layer_name}' added to project")
                        self.log_message.emit(f"   📊 {feature_count} polygons loaded")
                    else:
                        self.log_message.emit(f"   ⚠️ Vector layer created but could not be added to project")
            else:
                # Add raster layer
                layer = QgsRasterLayer(output_path, layer_name)
                
                if layer.isValid():
                    QgsProject.instance().addMapLayer(layer)
                    self.log_message.emit(f"   ✅ Raster layer '{layer_name}' added to project")
                else:
                    self.log_message.emit(f"   ⚠️ Raster layer created but could not be added to project")
                    
        except Exception as e:
            self.log_message.emit(f"   ⚠️ Could not add layer to QGIS: {str(e)}")
    
    def request_stop(self):
        """
        Request graceful stop of prediction.
        
        Sets stop flag, which is checked between tiles.
        """
        self._should_stop = True
        self.log_message.emit("\n⏸️ Stopping prediction (may take a few seconds)...")
    
    def _cleanup(self):
        """
        Clean up resources.
        """
        try:
            if self.raster_processor:
                self.raster_processor.cleanup()
            
            # Clear references
            self.raster_processor = None
            self.model_id = None
            
        except Exception as e:
            QgsMessageLog.logMessage(
                f"Cleanup error: {str(e)}",
                "QPredict",
                Qgis.Warning
            )