"""
Dock widget UI for QPredict

This module provides the user interface for configuring and monitoring
deep learning inference within QGIS. It handles user inputs for model
selection, input raster, tiling parameters, export parameters and displays
prediction progress and logs.
"""

import os
from qgis.PyQt.QtWidgets import (
    QDockWidget, QWidget, QVBoxLayout, QHBoxLayout, QPushButton,
    QLabel, QLineEdit, QComboBox, QCheckBox, QSpinBox, QDoubleSpinBox,
    QFileDialog, QTextEdit, QProgressBar, QGroupBox, QFormLayout,
    QScrollArea
)
from .advanced_vectorization_dialog import AdvancedVectorizationDialog
from .rectangle_map_tool import RectangleMapTool
from qgis.gui import QgsMapLayerComboBox, QgsRubberBand
from qgis.PyQt.QtCore import Qt, pyqtSignal
from qgis.PyQt.QtGui import QFont, QColor
from qgis.core import QgsMapLayerProxyModel, QgsWkbTypes, QgsPointXY


class PredictDockWidget(QDockWidget):
    """
    Main prediction dock widget for QPredict.
    
    This widget provides a comprehensive interface for:
    - Model selection and validation (.qmtp files)
    - Input raster selection with extent configuration
    - Tiling parameters (tile size, overlap, padding, batch size)
    - Device selection (CPU/CUDA)
    - Output configuration (path, format, colors)
    - Prediction monitoring (progress bar and logs)
    
    Signals:
        prediction_requested: Emitted when user clicks Start button
        prediction_cancelled: Emitted when user clicks Stop button
    """
    
    # Signals for communication with PredictionManager
    prediction_requested = pyqtSignal()
    prediction_cancelled = pyqtSignal()
    
    # Extent modes
    EXTENT_MODES = ['Full Raster', 'Current Viewport', 'Selection (Draw)']
    
    # Output formats
    OUTPUT_FORMATS = {
        'Class Raster': 'class_raster',              
        'Prediction Reliability': 'reliability_map',  
        'Vector Layer': 'vector'                      
    }
    
    def __init__(self, parent=None):
        """
        Initialize the dock widget.
        
        Args:
            parent: Parent widget (typically QGIS main window)
        """
        super().__init__("QPredict", parent)
        self.setObjectName("QPredictDock")  
        
        # State variables
        self._model_loaded = False
        self._model_info = None
        self._qmtp_path = None  
        self._model = None  
        self._model_loader = None  
        self._qmtp_loader = None   
        self._rectangle_tool = None 
        self._selection_extent = None  
        self._selection_rubber_band = None 

        # Advanced vectorization parameters
        self._advanced_vectorization_params = {
            'morph_kernel_size': 3,
            'min_region_pixels': 9,
            'chaikin_iterations': 1
        }
        
        self.init_ui()
    
    def init_ui(self):
        """
        Initialize the user interface.
        
        Creates all UI components organized in groups:
        1. Model selection and info
        2. Input raster and extent
        3. Tiling parameters
        4. Device selection
        5. Output configuration
        6. Progress monitoring
        """
        # Main container
        main_widget = QWidget()
        layout = QVBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        
        # ===== Header =====
        header_widget = self._create_header()
        layout.addWidget(header_widget)
        
        # ===== Content Area =====
        # Create scroll area for controls
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setFrameShape(QScrollArea.NoFrame)
        scroll_area.setStyleSheet("QScrollArea { border: none; }")
        
        # Scroll content widget
        scroll_content = QWidget()
        scroll_layout = QVBoxLayout()
        scroll_layout.setContentsMargins(10, 10, 10, 10)
        scroll_layout.setSpacing(10)
        
        # ===================================================================
        # MODEL CONFIGURATION
        # ===================================================================
        model_group = QGroupBox("Model")
        model_layout = QFormLayout()
        
        # Model path with browse button
        self.model_path_edit = QLineEdit()
        self.model_path_edit.setPlaceholderText("Select .qmtp model file...")
        self.model_path_edit.textChanged.connect(self.auto_load_model)
        model_btn = QPushButton("📁")
        model_btn.setMaximumWidth(40)
        model_btn.setToolTip("Browse for .qmtp model file")
        model_btn.clicked.connect(self.browse_model)
        
        model_row = QHBoxLayout()
        model_row.addWidget(self.model_path_edit)
        model_row.addWidget(model_btn)
        model_layout.addRow("Path:", model_row)
        
        # Model info label (updated after loading)
        self.model_info_label = QLabel("No model loaded")
        self.model_info_label.setStyleSheet("color: gray; font-size: 10px;")
        self.model_info_label.setWordWrap(True)
        model_layout.addRow(self.model_info_label)
        
        model_group.setLayout(model_layout)
        scroll_layout.addWidget(model_group)
        
        # ===================================================================
        # INPUT CONFIGURATION
        # ===================================================================
        input_group = QGroupBox("Input")
        input_layout = QFormLayout()
        
        # Raster layer selection
        self.raster_combo = QgsMapLayerComboBox()
        self.raster_combo.setFilters(QgsMapLayerProxyModel.RasterLayer)
        self.raster_combo.setAllowEmptyLayer(True)
        self.raster_combo.setShowCrs(True)
        input_layout.addRow("Raster Layer:", self.raster_combo)
        
        # Extent mode
        extent_container = QWidget()
        extent_layout = QHBoxLayout()
        extent_layout.setContentsMargins(0, 0, 0, 0)

        self.extent_combo = QComboBox()
        self.extent_combo.addItems(self.EXTENT_MODES)
        self.extent_combo.setToolTip(
            "Full Raster: Process entire raster\n"
            "Current Viewport: Process only visible area\n"
            "Selection: Draw rectangle to define area"
        )
        self.extent_combo.currentTextChanged.connect(self.on_extent_mode_changed)
        extent_layout.addWidget(self.extent_combo)

        # Draw button (only visible for Selection mode)
        self.draw_extent_btn = QPushButton("Draw Rectangle")
        self.draw_extent_btn.setMaximumWidth(120)
        self.draw_extent_btn.setToolTip("Click to draw rectangle on map")
        self.draw_extent_btn.clicked.connect(self.on_draw_extent_clicked)
        self.draw_extent_btn.setVisible(False)
        extent_layout.addWidget(self.draw_extent_btn)

        extent_layout.addStretch()
        extent_container.setLayout(extent_layout)
        input_layout.addRow("Extent:", extent_container)

        # Store selected extent for Selection mode
        self._selection_extent = None
        self._rectangle_tool = None
        
        # Input validation label
        self.input_info_label = QLabel("Select a raster layer")
        self.input_info_label.setStyleSheet("color: gray; font-size: 10px;")
        self.input_info_label.setWordWrap(True)
        input_layout.addRow(self.input_info_label)
        
        input_group.setLayout(input_layout)
        scroll_layout.addWidget(input_group)
        
        # ===================================================================
        # TILING CONFIGURATION
        # ===================================================================
        tiling_group = QGroupBox("Tiling")
        tiling_layout = QFormLayout()
        
        # Tile size
        self.tile_size_spin = QSpinBox()
        self.tile_size_spin.setRange(128, 2048)
        self.tile_size_spin.setSingleStep(32)
        self.tile_size_spin.setValue(512)
        self.tile_size_spin.setToolTip(
            "Size of each tile for inference.\n"
            "Should match training tile size for best results."
        )
        tiling_layout.addRow("Tile Size:", self.tile_size_spin)
        
        # Overlap
        self.overlap_spin = QSpinBox()
        self.overlap_spin.setRange(0, 1024)
        self.overlap_spin.setSingleStep(32)
        self.overlap_spin.setValue(256)
        self.overlap_spin.setToolTip(
            "Overlap between adjacent tiles (in pixels).\n"
            "Higher overlap = smoother results but slower processing."
        )
        # Show overlap percentage dynamically
        overlap_container = QWidget()
        overlap_layout = QHBoxLayout()
        overlap_layout.setContentsMargins(0, 0, 0, 0)
        overlap_layout.addWidget(self.overlap_spin)
        self.overlap_pct_label = QLabel("(50%)")
        self.overlap_pct_label.setStyleSheet("color: gray; font-size: 9px;")
        overlap_layout.addWidget(self.overlap_pct_label)
        overlap_layout.addStretch()
        overlap_container.setLayout(overlap_layout)
        tiling_layout.addRow("Overlap:", overlap_container)
        
        # Connect tile size and overlap to update percentage
        self.tile_size_spin.valueChanged.connect(self.update_overlap_percentage)
        self.overlap_spin.valueChanged.connect(self.update_overlap_percentage)
        
        # Padding
        self.padding_spin = QSpinBox()
        self.padding_spin.setRange(0, 256)
        self.padding_spin.setSingleStep(16)
        self.padding_spin.setValue(64)
        self.padding_spin.setToolTip(
            "Padding added around each tile to reduce edge artifacts.\n"
            "Padded regions are cropped after prediction.\n"
            "Recommended: 32-64 pixels."
        )
        tiling_layout.addRow("Padding:", self.padding_spin)
        
        # Batch size
        self.batch_size_spin = QSpinBox()
        self.batch_size_spin.setRange(1, 32)
        self.batch_size_spin.setValue(4)
        self.batch_size_spin.setToolTip(
            "Number of tiles processed simultaneously.\n"
            "Higher = faster but more GPU memory.\n"
            "Reduce if out-of-memory errors occur."
        )
        tiling_layout.addRow("Batch Size:", self.batch_size_spin)
        
        # Padding note for YOLO
        self.padding_note_label = QLabel("Padding locked to 0 for YOLO models")
        self.padding_note_label.setStyleSheet("color: gray; font-size: 9px; font-style: italic;")
        self.padding_note_label.setVisible(False)
        tiling_layout.addRow("", self.padding_note_label)
        
        tiling_group.setLayout(tiling_layout)
        scroll_layout.addWidget(tiling_group)
        
        # ===================================================================
        # DEVICE CONFIGURATION
        # ===================================================================
        device_group = QGroupBox("Device")
        device_layout = QFormLayout()
        
        # Device selection (CPU or CUDA)
        self.device_combo = QComboBox()
        self.device_combo.addItems(['CPU', 'CUDA'])
        self.device_combo.currentTextChanged.connect(self.on_device_changed)
        self.device_combo.currentTextChanged.connect(self._on_device_changed_for_yolo)
        device_layout.addRow("Device:", self.device_combo)
        
        # CUDA device ID (only visible when CUDA is selected)
        self.cuda_container = QWidget()
        cuda_layout = QHBoxLayout()
        cuda_layout.setContentsMargins(0, 0, 0, 0)
        self.cuda_device_label = QLabel("GPU ID:")
        self.cuda_device_spin = QSpinBox()
        self.cuda_device_spin.setRange(0, 7)
        self.cuda_device_spin.setValue(0)
        self.cuda_device_spin.setFixedWidth(60)
        self.cuda_device_spin.setToolTip("CUDA device ID (0-7)")
        cuda_layout.addWidget(self.cuda_device_label)
        cuda_layout.addWidget(self.cuda_device_spin)
        cuda_layout.addStretch()
        self.cuda_container.setLayout(cuda_layout)
        device_layout.addRow("", self.cuda_container)
        
        device_group.setLayout(device_layout)
        scroll_layout.addWidget(device_group)
        
        # Initialize CUDA visibility
        self.on_device_changed(self.device_combo.currentText())
        
        # ===================================================================
        # OUTPUT CONFIGURATION
        # ===================================================================
        output_group = QGroupBox("Output")
        output_layout = QFormLayout()
        
        # Output directory
        self.output_dir_edit = QLineEdit()
        self.output_dir_edit.setPlaceholderText("Select output folder...")
        output_btn = QPushButton("📁")
        output_btn.setMaximumWidth(40)
        output_btn.setToolTip("Browse for output folder")
        output_btn.clicked.connect(self.browse_output)
        
        output_row = QHBoxLayout()
        output_row.addWidget(self.output_dir_edit)
        output_row.addWidget(output_btn)
        output_layout.addRow("Path:", output_row)
        
        # Layer name
        self.layer_name_edit = QLineEdit()
        self.layer_name_edit.setPlaceholderText("prediction")
        self.layer_name_edit.setToolTip("Name for the output layer in QGIS")
        output_layout.addRow("Layer Name:", self.layer_name_edit)
        
        # Output format
        self.format_combo = QComboBox()
        self.format_combo.addItems(list(self.OUTPUT_FORMATS.keys()))
        self.format_combo.setToolTip(
            "Class Raster: 8-bit integer with class IDs\n"
            "Prediction Reliability: 32-bit float confidence/uncertainty maps\n"
            "Vector Layer: Polygons with class attributes"
        )
        self.format_combo.currentTextChanged.connect(self.on_format_changed)
        output_layout.addRow("Format:", self.format_combo)

        # Reliability sub-options (only visible when Prediction Reliability is selected)
        self.reliability_container = QWidget()
        reliability_layout = QFormLayout()
        reliability_layout.setContentsMargins(20, 0, 0, 0)  

        self.reliability_type_combo = QComboBox()
        self.reliability_type_combo.addItems(['Confidence Map', 'Uncertainty Map (Entropy)'])
        self.reliability_type_combo.setToolTip(
            "Confidence Map: Max softmax probability (0=uncertain, 1=confident)\n"
            "Uncertainty Map: Shannon entropy (0=certain, higher=uncertain)"
        )
        reliability_layout.addRow("Type:", self.reliability_type_combo)

        self.reliability_container.setLayout(reliability_layout)
        output_layout.addRow("", self.reliability_container)

        # Vector sub-options (only visible when Vector Layer is selected)
        self.vector_container = QWidget()
        vector_layout = QFormLayout()
        vector_layout.setContentsMargins(20, 0, 0, 0)  
        
        # Simplify tolerance
        simplify_row = QHBoxLayout()
        self.simplify_spin = QDoubleSpinBox()
        self.simplify_spin.setRange(0.0, 10.0)
        self.simplify_spin.setSingleStep(0.5)
        self.simplify_spin.setValue(1.0)
        self.simplify_spin.setDecimals(1)
        self.simplify_spin.setSuffix(" px")
        self.simplify_spin.setToolTip(
        "Topology-aware simplification tolerance in pixels.\n"
        "Higher value = simpler geometries (fewer vertices).\n"
        "0 = no simplification.\n\n"
        "Uses Shapely simplify with preserve_topology=True\n"
        "to avoid self-intersections and maintain validity."
        )
        simplify_row.addWidget(self.simplify_spin)
        simplify_row.addStretch()
        vector_layout.addRow("Simplify Tolerance:", simplify_row)
        
        # Exclude background checkbox
        self.exclude_background_cb = QCheckBox("Exclude background (class 0)")
        self.exclude_background_cb.setChecked(True)
        self.exclude_background_cb.setToolTip(
            "Skip background class (class 0) during vectorization.\n"
            "Recommended for most use cases to avoid large\n"
            "background polygons."
        )
        vector_layout.addRow("", self.exclude_background_cb)
        
        # Advanced vectorization parameters button
        advanced_btn_layout = QHBoxLayout()
        self.advanced_vectorization_btn = QPushButton("Advanced Parameters...")
        self.advanced_vectorization_btn.setMaximumWidth(180)
        self.advanced_vectorization_btn.setToolTip(
            "Configure advanced vectorization parameters:\n"
            "• Morphology kernel size\n"
            "• Min region filtering\n"
            "• Chaikin smoothing iterations"
        )
        self.advanced_vectorization_btn.clicked.connect(self.on_advanced_vectorization_clicked)
        advanced_btn_layout.addWidget(self.advanced_vectorization_btn)
        advanced_btn_layout.addStretch()
        vector_layout.addRow("", advanced_btn_layout)
        
        self.vector_container.setLayout(vector_layout)
        output_layout.addRow("", self.vector_container)
        
        # Apply class colors (only for raster outputs)
        self.apply_colors_cb = QCheckBox("Apply class colors")
        self.apply_colors_cb.setChecked(True)
        self.apply_colors_cb.setToolTip(
            "Apply color palette from model to output raster.\n"
            "Only available for raster outputs."
        )
        output_layout.addRow("", self.apply_colors_cb)

        # Confidence threshold (only for instance segmentation)
        self.threshold_container = QWidget()
        threshold_layout = QHBoxLayout()
        threshold_layout.setContentsMargins(0, 0, 0, 0)
        
        self.threshold_spin = QDoubleSpinBox()
        self.threshold_spin.setRange(0.0, 1.0)
        self.threshold_spin.setSingleStep(0.05)
        self.threshold_spin.setValue(0.5)
        self.threshold_spin.setDecimals(2)
        self.threshold_spin.setToolTip(
            "Confidence threshold for filtering detections.\n"
            "Only instances with confidence above this threshold\n"
            "will be kept in the output."
        )
        threshold_layout.addWidget(self.threshold_spin)
        threshold_layout.addStretch()
        
        self.threshold_container.setLayout(threshold_layout)
        output_layout.addRow("Threshold:", self.threshold_container)
        
        # YOLO-specific thresholds (only visible for YOLO models)
        self.yolo_thresholds_container = QWidget()
        yolo_thresholds_layout = QFormLayout()
        yolo_thresholds_layout.setContentsMargins(0, 0, 0, 0)
        
        # Confidence threshold
        self.conf_threshold_spin = QDoubleSpinBox()
        self.conf_threshold_spin.setRange(0.0, 1.0)
        self.conf_threshold_spin.setSingleStep(0.05)
        self.conf_threshold_spin.setValue(0.5)
        self.conf_threshold_spin.setDecimals(2)
        self.conf_threshold_spin.setToolTip(
            "Confidence threshold for YOLO detections.\n"
            "Detections with confidence below this value are filtered out."
        )
        yolo_thresholds_layout.addRow("Confidence Threshold:", self.conf_threshold_spin)
        
        # IoU threshold
        self.iou_threshold_spin = QDoubleSpinBox()
        self.iou_threshold_spin.setRange(0.0, 1.0)
        self.iou_threshold_spin.setSingleStep(0.05)
        self.iou_threshold_spin.setValue(0.5)
        self.iou_threshold_spin.setDecimals(2)
        self.iou_threshold_spin.setToolTip(
            "IoU threshold for Non-Maximum Suppression (NMS).\n"
            "Lower values = more aggressive suppression of overlapping detections.\n"
            "Higher values = keep more overlapping detections (useful for dense scenes)."
        )
        yolo_thresholds_layout.addRow("IoU Threshold (NMS):", self.iou_threshold_spin)
        
        self.yolo_thresholds_container.setLayout(yolo_thresholds_layout)
        self.yolo_thresholds_container.setVisible(False) 
        output_layout.addRow("", self.yolo_thresholds_container)
        
        output_group.setLayout(output_layout)
        scroll_layout.addWidget(output_group)
        
        # Finalize scroll area
        scroll_content.setLayout(scroll_layout)
        scroll_area.setWidget(scroll_content)
        
        # Add scroll area to main layout
        layout.addWidget(scroll_area, 1)  
        layout.addSpacing(5)
        
        # ===================================================================
        # PROGRESS & LOG (FIXED AT BOTTOM)
        # ===================================================================
        # Container for progress/log section with padding
        progress_log_container = QWidget()
        progress_log_layout = QVBoxLayout()
        progress_log_layout.setContentsMargins(10, 5, 10, 5)
        progress_log_layout.setSpacing(5)

        # Progress bar
        progress_label = QLabel("Progress:")
        progress_label_font = QFont()
        progress_label_font.setBold(True)
        progress_label.setFont(progress_label_font)
        progress_log_layout.addWidget(progress_label)

        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        self.progress_bar.setTextVisible(True)
        progress_log_layout.addWidget(self.progress_bar)

        # Add spacing between progress and log
        progress_log_layout.addSpacing(10)

        # Prediction log
        log_header = QHBoxLayout()
        log_label = QLabel("Log:")
        log_label_font = QFont()
        log_label_font.setBold(True)
        log_label.setFont(log_label_font)
        log_header.addWidget(log_label)

        # Clear log button
        self.clear_log_btn = QPushButton("Clear")
        self.clear_log_btn.setMaximumWidth(60)
        self.clear_log_btn.clicked.connect(self.clear_log)
        log_header.addWidget(self.clear_log_btn)
        log_header.addStretch()

        progress_log_layout.addLayout(log_header)

        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setMaximumHeight(180)
        self.log_text.setStyleSheet(
            "QTextEdit { "
            "font-family: 'Courier New', monospace; "
            "font-size: 9pt; "
            "}"
        )
        progress_log_layout.addWidget(self.log_text)

        progress_log_container.setLayout(progress_log_layout)
        layout.addWidget(progress_log_container)

        # ===================================================================
        # CONTROL BUTTONS
        # ===================================================================
        btn_container = QWidget()
        btn_layout = QHBoxLayout()
        btn_layout.setContentsMargins(10, 0, 10, 10)
        btn_layout.setSpacing(10)
        
        self.start_btn = QPushButton("Start Prediction")
        self.start_btn.setStyleSheet(
            "QPushButton {"
            "    background-color: #4CAF50; "
            "    color: white; "
            "    padding: 8px; "
            "    font-weight: bold; "
            "    border: none; "
            "    border-radius: 4px;"
            "}"
            "QPushButton:hover {"
            "    background-color: #45a049;"
            "}"
            "QPushButton:disabled {"
            "    background-color: #cccccc;"
            "}"
        )
        self.start_btn.clicked.connect(self.on_start_clicked)
        
        self.stop_btn = QPushButton("Stop")
        self.stop_btn.setStyleSheet(
            "QPushButton {"
            "    background-color: #f44336; "
            "    color: white; "
            "    padding: 8px; "
            "    border: none; "
            "    border-radius: 4px;"
            "}"
            "QPushButton:hover {"
            "    background-color: #da190b;"
            "}"
            "QPushButton:disabled {"
            "    background-color: #cccccc;"
            "}"
        )
        self.stop_btn.setEnabled(False)
        self.stop_btn.clicked.connect(self.on_stop_clicked)
        
        btn_layout.addWidget(self.start_btn)
        btn_layout.addWidget(self.stop_btn)
        btn_container.setLayout(btn_layout)
        layout.addWidget(btn_container)

        # ===== Footer =====
        footer_widget = self._create_footer()
        layout.addWidget(footer_widget)

        # Set main layout
        main_widget.setLayout(layout)
        self.setWidget(main_widget)

        # Connect raster combo signal after input_info_label is created
        self.raster_combo.layerChanged.connect(self.on_raster_changed)

        # Initialize format-dependent UI visibility
        self.on_format_changed(self.format_combo.currentText())
        
        # Dock widget configuration
        self.setAllowedAreas(Qt.RightDockWidgetArea | Qt.LeftDockWidgetArea)
        self.setFeatures(
            QDockWidget.DockWidgetMovable |
            QDockWidget.DockWidgetFloatable |
            QDockWidget.DockWidgetClosable
        )
    
    # =======================================================================
    # DYNAMIC UI UPDATES
    # =======================================================================
    def _create_header(self) -> QWidget:
        """Create header widget"""
        header = QWidget()
        header.setStyleSheet(
            "background-color: #2e5872; "
            "color: white; "
            "padding: 15px;"
        )
        
        layout = QVBoxLayout()
        layout.setContentsMargins(15, 15, 15, 15)
        layout.setSpacing(5)
        
        # Title
        title = QLabel("QPredict")
        title_font = QFont()
        title_font.setPointSize(14)
        title_font.setBold(True)
        title.setFont(title_font)
        title.setStyleSheet("color: white;")
        layout.addWidget(title)
        
        # Subtitle
        subtitle = QLabel("Inference tool")
        subtitle.setStyleSheet("color: #BDC3C7; font-size: 10px;")
        layout.addWidget(subtitle)
        
        header.setLayout(layout)
        return header

    def _create_footer(self) -> QWidget:
        """Create footer widget"""
        footer = QWidget()
        footer.setStyleSheet(
            "background-color: #ECF0F1; "
            "border-top: 1px solid #BDC3C7; "
            "padding: 10px;"
        )
        
        layout = QVBoxLayout()
        layout.setContentsMargins(15, 10, 15, 10)
        
        footer_label = QLabel(
            "<b>QGeoAI</b> | edited by "
            "<a href='https://qgeoai.nextelia.fr/' style='color: #3498DB;'>Nextelia</a>"
        )
        footer_label.setOpenExternalLinks(True)
        footer_label.setAlignment(Qt.AlignCenter)
        footer_label.setStyleSheet("color: #7F8C8D; font-size: 9px;")
        layout.addWidget(footer_label)
        
        footer.setLayout(layout)
        return footer    
    
    def on_raster_changed(self, layer):
        """
        Handle raster layer selection change.
        Validates compatibility with loaded model.
        
        Args:
            layer: QgsRasterLayer or None
        """
        if layer is None:
            self.input_info_label.setText("Select a raster layer")
            self.input_info_label.setStyleSheet("color: gray; font-size: 10px;")
            return
        
        if not layer.isValid():
            self.input_info_label.setText("⚠️ Invalid layer")
            self.input_info_label.setStyleSheet("color: orange; font-size: 10px;")
            return
        
        # Basic validation
        num_bands = layer.bandCount()
        crs = layer.crs().authid()
        
        # Check if model is loaded
        if not self._model_loaded or not self._model_info:
            self.input_info_label.setText(f"✅ {num_bands} bands, {crs}")
            self.input_info_label.setStyleSheet("color: green; font-size: 10px;")
            return
        
        # Validate against model requirements
        model_bands = self._model_info['configuration']['model']['in_channels']
        model_crs = self._model_info['configuration']['geospatial']['crs']
        
        messages = []
        
        # Check bands
        if num_bands != model_bands:
            messages.append(f"⚠️ Band mismatch: model expects {model_bands}, layer has {num_bands}")
        else:
            messages.append(f"✅ {num_bands} bands")
        
        # Check CRS (warning only, not blocking)
        if crs != model_crs:
            messages.append(f"⚠️ CRS differs: model={model_crs}, layer={crs}")
        
        # Check resolution (warning only)
        pixel_x = layer.rasterUnitsPerPixelX()
        pixel_y = layer.rasterUnitsPerPixelY()
        model_pixel_x = self._model_info['configuration']['geospatial']['pixel_size_x']
        model_pixel_y = self._model_info['configuration']['geospatial']['pixel_size_y']
        
        if abs(pixel_x - model_pixel_x) > 0.01 or abs(pixel_y - model_pixel_y) > 0.01:
            messages.append(
                f"⚠️ Resolution differs: trained=({model_pixel_x:.2f}, {model_pixel_y:.2f}), "
                f"input=({pixel_x:.2f}, {pixel_y:.2f})"
            )
        
        # Update label
        self.input_info_label.setText("\n".join(messages))
        
        # Color based on severity
        if any("⚠️" in msg for msg in messages):
            self.input_info_label.setStyleSheet("color: orange; font-size: 10px;")
        else:
            self.input_info_label.setStyleSheet("color: green; font-size: 10px;")
    
    def update_overlap_percentage(self):
        """
        Update overlap percentage label dynamically.
        """
        tile_size = self.tile_size_spin.value()
        overlap = self.overlap_spin.value()
        
        if tile_size > 0:
            percentage = (overlap / tile_size) * 100
            self.overlap_pct_label.setText(f"({percentage:.0f}%)")
    
    def on_device_changed(self, device_text):
        """
        Show/hide CUDA device selector based on device selection.
        
        Args:
            device_text: Selected device ('CPU' or 'CUDA')
        """
        self.cuda_container.setVisible(device_text == "CUDA")

    def update_available_formats(self):
        """
        Update available output formats based on loaded model task.
        
        Different tasks support different output formats:
        - Semantic: Class Raster, Prediction Reliability, Vector Layer
        - Instance: Vector Layer, Class Raster (no Reliability)
        - YOLO: Vector Layer ONLY
        """
        if not self._model_info:
            # No model loaded - show all formats
            return
        
        task_type = self._model_info.get('task_type', 'semantic_segmentation')
        
        # Clear current items
        self.format_combo.clear()
        
        # Determine available formats and whether combo should be locked
        is_locked = False
        
        if task_type == 'semantic_segmentation':
            # All formats available
            self.format_combo.addItems(['Class Raster', 'Prediction Reliability', 'Vector Layer'])
        
        elif task_type == 'instance_segmentation':
            # No Prediction Reliability (not applicable to instance segmentation)
            self.format_combo.addItems(['Vector Layer', 'Class Raster'])
            # Set default to Vector Layer (primary for instance seg)
            self.format_combo.setCurrentText('Vector Layer')
        
        elif task_type.startswith('yolo_'):
            # YOLO: Only Vector Layer
            self.format_combo.addItems(['Vector Layer'])
            # Lock to Vector Layer
            is_locked = True
        
        elif task_type == 'object_detection':
            # Only Vector Layer (future YOLO)
            self.format_combo.addItems(['Vector Layer'])
        
        else:
            # Unknown task - show all formats
            self.format_combo.addItems(['Class Raster', 'Prediction Reliability', 'Vector Layer'])
        
        # Apply lock state
        self.format_combo.setEnabled(not is_locked)
    
    def on_format_changed(self, format_text):
        """
        Enable/disable options based on output format.
        
        Args:
            format_text: Selected output format
        """
        is_class_raster = format_text == 'Class Raster'
        is_reliability = format_text == 'Prediction Reliability'
        is_vector = format_text == 'Vector Layer'
        
        # Check if YOLO model is loaded
        is_yolo = False
        if self._model_info:
            task_type = self._model_info.get('task_type', 'semantic_segmentation')
            is_yolo = task_type.startswith('yolo_')
        
        # Apply colors: only for Class Raster
        self.apply_colors_cb.setVisible(is_class_raster)
        self.apply_colors_cb.setEnabled(is_class_raster)
        
        # Reliability sub-options only for Prediction Reliability
        self.reliability_container.setVisible(is_reliability)
        
        # Vector sub-options: conditional visibility for YOLO
        if is_yolo:
            # For YOLO Segment, show simplify tolerance
            # For YOLO Detection/OBB, hide
            task_type = self._model_info.get('task_type', 'semantic_segmentation')
            if task_type == 'yolo_segment':
                self.vector_container.setVisible(is_vector)
                # Hide exclude_background checkbox
                self.exclude_background_cb.setVisible(False)
            else:
                # YOLO Detection/OBB: hide all vector options
                self.vector_container.setVisible(False)
        else:
            # PyTorch: show all vector options
            self.vector_container.setVisible(is_vector)
            self.exclude_background_cb.setVisible(True)
        
        # Threshold: Only for instance segmentation
        if is_yolo:
            self.threshold_container.setVisible(False)
        else:
            # Check if instance segmentation
            task_type = self._model_info.get('task_type', 'semantic_segmentation') if self._model_info else 'semantic_segmentation'
            is_instance = task_type == 'instance_segmentation'
            
            # Threshold only visible for instance segmentation
            self.threshold_container.setVisible(is_instance)
    
    # =======================================================================
    # FILE BROWSER METHODS
    # =======================================================================
    
    def browse_model(self):
        """Open file dialog to select .qmtp model file."""
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "Select QModel Trainer Package (.qmtp)",
            "",
            "QMTP Files (*.qmtp);;All Files (*)"
        )
        if file_path:
            self.model_path_edit.setText(file_path)
    
    def browse_output(self):
        """Open file dialog to select output directory."""
        directory = QFileDialog.getExistingDirectory(
            self,
            "Select Output Directory"
        )
        if directory:
            self.output_dir_edit.setText(directory)
    
    # =======================================================================
    # MODEL LOADING
    # =======================================================================
    
    def auto_load_model(self, text):
        """
        Auto-load and validate model when path changes.
        
        Args:
            text: New model path
        """
        from ..core.qmtp_loader import QMTPLoader
        
        # Reset state
        self._model_loaded = False
        self._model_info = None
        self._qmtp_path = None  
        self._model = None
        self._model_loader = None
        self._qmtp_loader = None
        
        model_path = text.strip()
        
        if not model_path:
            self.model_info_label.setText("No model loaded")
            self.model_info_label.setStyleSheet("color: gray; font-size: 10px;")
            return
        
        # Check if file exists (basic check before full validation)
        if not os.path.exists(model_path):
            self.model_info_label.setText("⚠️ File does not exist")
            self.model_info_label.setStyleSheet("color: orange; font-size: 10px;")
            return
        
        # Show loading status
        self.model_info_label.setText("🔄 Loading model...")
        self.model_info_label.setStyleSheet("color: blue; font-size: 10px;")
        
        # Load and validate .qmtp file
        try:
            # Create loader
            loader = QMTPLoader(model_path)
            
            # Validate .qmtp structure
            is_valid, error = loader.validate()
            
            if not is_valid:
                # Validation failed - show error
                self.model_info_label.setText(f"❌ {error}")
                self.model_info_label.setStyleSheet("color: red; font-size: 10px;")
                self._model_loaded = False
                self._model_info = None
                self.append_log(f"\n❌ Model validation failed: {error}\n")
                return
            
            # Load model information
            model_info = loader.load_info()
            
            # Format model info for display
            info_text = loader.format_info_for_display()
            
            # Log model info
            self.append_log("\n" + "="*50)
            self.append_log("Model Information")
            self.append_log("="*50)
            self.append_log(f"File: {os.path.basename(model_path)}")
            self.append_log(info_text)
            self.append_log(f"Input: {model_info['in_channels']} bands, {model_info['tile_size']}x{model_info['tile_size']} pixels")
            self.append_log(f"CRS: {model_info['crs']}")
            self.append_log(f"Pixel size: ({model_info['pixel_size_x']:.4f}, {model_info['pixel_size_y']:.4f})")
            self.append_log("="*50)
            self.append_log("✅ Model validated and ready for inference")
            self.append_log("")
            
            # Store loader and info
            self._qmtp_loader = loader
            self._qmtp_path = model_path  
            self._model_info = model_info
            self._model_loaded = True
            self._model = None 
            self._model_loader = None
            
            # Format and display model info in label
            self.model_info_label.setText(f"✅ Model validated\n{info_text}")
            self.model_info_label.setStyleSheet("color: green; font-size: 10px;")
            
            # Update tile size to match model (if different)
            if self.tile_size_spin.value() != model_info['tile_size']:
                self.tile_size_spin.setValue(model_info['tile_size'])
                self.append_log(f"ℹ️ Tile size adjusted to {model_info['tile_size']} (model requirement)\n")
            
            # Update available output formats based on task
            self.update_available_formats()
            # Update UI for YOLO models
            self._update_ui_for_yolo()
            
            # Re-validate input raster if one is selected
            current_layer = self.raster_combo.currentLayer()
            if current_layer is not None:
                self.on_raster_changed(current_layer)
            
        except Exception as e:
            self.model_info_label.setText(f"❌ Failed to load: {str(e)}")
            self.model_info_label.setStyleSheet("color: red; font-size: 10px;")
            self._model_loaded = False
            self._model_info = None
            self._model = None 
            self.append_log(f"\n❌ Error loading model: {str(e)}\n")
    
    # =======================================================================
    # LOG MANAGEMENT
    # =======================================================================
    
    def append_log(self, message):
        """
        Append a message to the prediction log.
        
        Args:
            message: Text message to append
        """
        self.log_text.append(message)
        # Auto-scroll to bottom
        self.log_text.verticalScrollBar().setValue(
            self.log_text.verticalScrollBar().maximum()
        )
    
    def clear_log(self):
        """Clear the prediction log."""
        self.log_text.clear()

    # =======================================================================
    # EXTENT SELECTION
    # =======================================================================
    
    def on_extent_mode_changed(self, mode):
        """
        Handle extent mode change.
        
        Args:
            mode: Selected extent mode
        """
        # Show/hide draw button
        is_selection = mode == 'Selection (Draw)'
        self.draw_extent_btn.setVisible(is_selection)
        
        # Reset selection extent if mode changed
        if not is_selection:
            self._selection_extent = None
            
            # Deactivate and clear rectangle tool if active
            if self._rectangle_tool:
                from qgis.utils import iface
                iface.mapCanvas().unsetMapTool(self._rectangle_tool)
                self._rectangle_tool.clear_rectangle()
            
            # Clear persistent rubber band
            if self._selection_rubber_band:
                self._selection_rubber_band.reset()

    def on_draw_extent_clicked(self):
        """
        Activate rectangle drawing tool.
        """
        from qgis.utils import iface
        
        # Create tool if not exists
        if not self._rectangle_tool:
            self._rectangle_tool = RectangleMapTool(iface.mapCanvas())
            self._rectangle_tool.rectangle_created.connect(self.on_rectangle_drawn)
        else:
            # Clear previous rectangles before drawing new one
            self._rectangle_tool.clear_rectangle()
            if self._selection_rubber_band:
                self._selection_rubber_band.reset()
        
        # Activate tool
        iface.mapCanvas().setMapTool(self._rectangle_tool)
        
        # Update button text
        self.draw_extent_btn.setText("Drawing...")
        self.draw_extent_btn.setEnabled(False)
        
        self.append_log("\n📐 Draw rectangle on map (click two corners, ESC to cancel)")

    def on_rectangle_drawn(self, rectangle):
        """
        Handle rectangle drawn by user.
        
        Args:
            rectangle: QgsRectangle in map coordinates
        """
        from qgis.utils import iface

        # Store extent
        self._selection_extent = rectangle
        
        # Create persistent rubber band if not exists
        if not self._selection_rubber_band:
            self._selection_rubber_band = QgsRubberBand(
                iface.mapCanvas(),
                QgsWkbTypes.PolygonGeometry
            )
            # Style: transparent fill with blue border
            self._selection_rubber_band.setColor(QColor(0, 0, 255, 50))
            self._selection_rubber_band.setWidth(2)
            self._selection_rubber_band.setLineStyle(Qt.SolidLine)
        
        # Update rubber band with rectangle
        self._selection_rubber_band.reset(QgsWkbTypes.PolygonGeometry)
        self._selection_rubber_band.addPoint(QgsPointXY(rectangle.xMinimum(), rectangle.yMinimum()))
        self._selection_rubber_band.addPoint(QgsPointXY(rectangle.xMaximum(), rectangle.yMinimum()))
        self._selection_rubber_band.addPoint(QgsPointXY(rectangle.xMaximum(), rectangle.yMaximum()))
        self._selection_rubber_band.addPoint(QgsPointXY(rectangle.xMinimum(), rectangle.yMaximum()))
        self._selection_rubber_band.addPoint(QgsPointXY(rectangle.xMinimum(), rectangle.yMinimum()))
        self._selection_rubber_band.show()
        
        # Deactivate tool
        iface.mapCanvas().unsetMapTool(self._rectangle_tool)
        
        # Update button
        self.draw_extent_btn.setText("Redraw Rectangle")
        self.draw_extent_btn.setEnabled(True)
        
        # Log extent
        self.append_log(
            f"✅ Rectangle drawn: "
            f"X[{rectangle.xMinimum():.2f}, {rectangle.xMaximum():.2f}], "
            f"Y[{rectangle.yMinimum():.2f}, {rectangle.yMaximum():.2f}]"
        )
    
    # =======================================================================
    # PREDICTION CONTROL
    # =======================================================================
    
    def on_start_clicked(self):
        """
        Handle Start button click.
        Validates configuration before starting.
        """
        # Check if model is loaded
        if not self._model_loaded:
            self.append_log("\n❌ Please load a valid .qmtp model first.\n")
            return
        
        # Check if raster is selected
        if self.raster_combo.currentLayer() is None:
            self.append_log("\n❌ Please select an input raster layer.\n")
            return
        
        # Check if selection extent is drawn (for Selection mode)
        if self.extent_combo.currentText() == 'Selection (Draw)' and not self._selection_extent:
            self.append_log("\n❌ Please draw a rectangle on the map first.\n")
            return
        
        # Check if output directory is set
        if not self.output_dir_edit.text():
            self.append_log("\n❌ Please select an output directory.\n")
            return
        
        # Log start
        self.append_log("\n" + "="*50)
        self.append_log("Starting prediction...")
        self.append_log("="*50)
        
        # Disable start button, enable stop button
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        
        # Emit signal
        self.prediction_requested.emit()
    
    def on_stop_clicked(self):
        """
        Handle Stop button click.
        Emits prediction_cancelled signal.
        """
        self.append_log("\n⚠️ Stop requested by user...")
        self.prediction_cancelled.emit()
    
    def on_prediction_cleanup(self):
        """
        Called when prediction completes or is cancelled.
        Resets UI to ready state for next prediction.
        
        This method:
        - Resets button states
        - Resets progress bar to 0
        - Clears selection rectangle
        - Displays ready message
        - Unloads model from memory
        """
        # Reset button states
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        
        # Reset progress bar to 0
        self.progress_bar.setValue(0)
        
        # Clear selection rectangle if it was used
        if self.extent_combo.currentText() == 'Selection (Draw)':
            if self._rectangle_tool:
                self._rectangle_tool.clear_rectangle()
            if self._selection_rubber_band:
                self._selection_rubber_band.reset()
        
        # Display ready message
        self.append_log("\n" + "="*50)
        self.append_log("✅ Ready for next prediction")
        self.append_log("="*50 + "\n")

    # =======================================================================
    # YOLO-SPECIFIC UI UPDATES
    # =======================================================================
    
    def _update_ui_for_yolo(self):
        """
        Update UI when YOLO model is loaded.
        
        - Lock padding to 0
        - Adjust batch size based on device
        - Show YOLO thresholds (conf, iou)
        - Hide PyTorch threshold
        - Update format visibility
        """
        if not self._model_info:
            return
        
        task_type = self._model_info.get('task_type', 'semantic_segmentation')
        is_yolo = task_type.startswith('yolo_')
        
        if is_yolo:
            # Lock padding to 0 for YOLO
            self.padding_spin.setValue(0)
            self.padding_spin.setEnabled(False)
            self.padding_note_label.setVisible(True)
            
            # Adjust batch size based on current device
            current_device = self.device_combo.currentText()
            if current_device == 'CPU':
                self.batch_size_spin.setValue(1)
                self.batch_size_spin.setEnabled(False)
                self.batch_size_spin.setToolTip(
                    "Batch size locked to 1 for YOLO on CPU.\n"
                    "CPU does not benefit from batching."
                )
            else:
                # CUDA: set to 1 by default but allow user to change
                self.batch_size_spin.setValue(1)
                self.batch_size_spin.setEnabled(True)
                self.batch_size_spin.setToolTip(
                    "Number of tiles processed simultaneously on GPU.\n"
                    "Higher = faster but more GPU memory.\n"
                    "Start with 1, increase if you have GPU memory available.\n"
                    "Recommended: 1-4 for YOLO models."
                )
            
            # Show YOLO thresholds
            self.yolo_thresholds_container.setVisible(True)
            
            # Hide PyTorch threshold
            self.threshold_container.setVisible(False)
            
            self.append_log("\n YOLO model detected:")
            self.append_log("   - Padding locked to 0 (YOLO doesn't use padding)")
            self.append_log("   - Batch size adjusted for YOLO inference")
            self.append_log("   - Using confidence + IoU thresholds instead of softmax threshold")
            self.append_log("   - Output format: Vector Layer only")
        else:
            # Reset for PyTorch models
            self.padding_spin.setEnabled(True)
            self.padding_note_label.setVisible(False)
            
            # Reset batch size to default for PyTorch
            self.batch_size_spin.setValue(4)
            self.batch_size_spin.setEnabled(True)
            self.batch_size_spin.setToolTip(
                "Number of tiles processed simultaneously.\n"
                "Higher = faster but more GPU memory.\n"
                "Reduce if out-of-memory errors occur."
            )
            
            self.yolo_thresholds_container.setVisible(False)
            

    def _on_device_changed_for_yolo(self, device_text):
        """
        Adjust batch size when device changes if YOLO model is loaded.
        
        Args:
            device_text: Selected device ('CPU' or 'CUDA')
        """
        # Only apply if YOLO model is loaded
        if not self._model_info:
            return
        
        task_type = self._model_info.get('task_type', 'semantic_segmentation')
        is_yolo = task_type.startswith('yolo_')
        
        if not is_yolo:
            return
        
        # Adjust batch size for YOLO based on device
        if device_text == 'CPU':
            self.batch_size_spin.setValue(1)
            self.batch_size_spin.setEnabled(False)
            self.batch_size_spin.setToolTip(
                "Batch size locked to 1 for YOLO on CPU.\n"
                "CPU does not benefit from batching."
            )
        else:
            # CUDA: allow user choice
            self.batch_size_spin.setValue(1)
            self.batch_size_spin.setEnabled(True)
            self.batch_size_spin.setToolTip(
                "Number of tiles processed simultaneously on GPU.\n"
                "Higher = faster but more GPU memory.\n"
                "Start with 1, increase if you have GPU memory available.\n"
                "Recommended: 1-4 for YOLO models."
            )
    # =======================================================================
    # ADVANCED PARAMETERS
    # =======================================================================
    
    def on_advanced_vectorization_clicked(self):
        """
        Open advanced vectorization parameters dialog.
        """
        dialog = AdvancedVectorizationDialog(
            parent=self,
            current_params=self._advanced_vectorization_params
        )
        
        if dialog.exec_():
            # User clicked OK - save parameters
            self._advanced_vectorization_params = dialog.get_parameters()
            
            self.append_log("\n[Advanced Vectorization] Parameters updated:")
            self.append_log(f"   Morphology kernel: {self._advanced_vectorization_params['morph_kernel_size']}")
            self.append_log(f"   Min region: {self._advanced_vectorization_params['min_region_pixels']} pixels")
            self.append_log(f"   Chaikin iterations: {self._advanced_vectorization_params['chaikin_iterations']}")
    
    # =======================================================================
    # CONFIGURATION GETTERS
    # =======================================================================
    
    def get_prediction_config(self):
        """
        Collect all prediction configuration from UI.
        
        Returns:
            dict: Complete prediction configuration dictionary
        """
        # Get selected raster layer
        layer = self.raster_combo.currentLayer()
        
        return {
            # Model
            'model_path': self.model_path_edit.text(),
            
            # Input
            'input_layer': layer,
            'extent_mode': self.extent_combo.currentText(),
            'selection_extent': self._selection_extent,
            
            # Tiling
            'tile_size': self.tile_size_spin.value(),
            'overlap': self.overlap_spin.value(),
            'padding': self.padding_spin.value(),
            'batch_size': self.batch_size_spin.value(),
            
            # Device
            'device': self.device_combo.currentText(),
            'cuda_device_id': self.cuda_device_spin.value() if self.device_combo.currentText() == 'CUDA' else None,
            
            # Output
            'output_dir': self.output_dir_edit.text(),
            'layer_name': self.layer_name_edit.text() or 'prediction',
            'output_format': self.OUTPUT_FORMATS[self.format_combo.currentText()],
            'apply_colors': self.apply_colors_cb.isChecked(),
            'threshold': self.threshold_spin.value(),
            'reliability_type': self.reliability_type_combo.currentText() if self.format_combo.currentText() == 'Prediction Reliability' else None,
            'simplify_tolerance': self.simplify_spin.value() if self.format_combo.currentText() == 'Vector Layer' else 1.0,
            'exclude_background': self.exclude_background_cb.isChecked() if self.format_combo.currentText() == 'Vector Layer' else True,
            # Advanced vectorization parameters
            'morph_kernel_size': self._advanced_vectorization_params['morph_kernel_size'],
            'min_region_pixels': self._advanced_vectorization_params['min_region_pixels'],
            'chaikin_iterations': self._advanced_vectorization_params['chaikin_iterations'],
            # YOLO-specific thresholds
            'confidence_threshold': self.conf_threshold_spin.value(),
            'iou_threshold': self.iou_threshold_spin.value()
        }
    
    def get_qmtp_loader(self):
        """
        Get the QMTP loader for the validated model.
        
        Returns:
            QMTPLoader instance if model is validated, None otherwise
        """
        return self._qmtp_loader if self._model_loaded else None
    
    def update_progress(self, current: int, total: int):
        """
        Update progress bar.
        
        Args:
            current: Current number of tiles processed
            total: Total number of tiles
        """
        if total > 0:
            percentage = int((current / total) * 100)
            self.progress_bar.setValue(percentage)