# -*- coding: utf-8 -*-
"""
/***************************************************************************
 OMRAT
                                 A QGIS plugin
 This is an open source implementation of Pedersens equations from 1995
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2022-12-22
        git sha              : $Format:%H$
        copyright            : (C) 2022 by Axel Hörteborn
        email                : axel.horteborn@ri.se
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""
from typing import Any, Callable
from functools import partial
from qgis.PyQt.QtCore import QSettings, QTranslator, QCoreApplication, QMetaType
from qgis.PyQt.QtGui import QIcon, QAction, QCloseEvent
from qgis.PyQt.QtWidgets import QMenuBar, QWidget, QFileDialog, QToolBar, QMessageBox
from qgis._core import QgsVectorDataProvider
from qgis.core import (QgsVectorLayer, QgsFeature, QgsGeometry, QgsLineString, QgsPoint, QgsProject,
                       QgsField, QgsFields)
from qgis.core import QgsPointXY, QgsMessageLog, Qgis, QgsApplication

from qgis.gui import QgsMapToolPan, QgisInterface
import gc
import sys
import os
import os.path
import time
from pathlib import Path


def _fmt(v: float | None) -> str:
    """Format a probability total for the previous-runs table."""
    if v is None:
        return ''
    if v == 0:
        return '0'
    return f'{v:.3e}'


def _format_duration(seconds: float | None) -> str:
    """Human-readable elapsed-time string for the Duration column."""
    if seconds is None:
        return ''
    s = float(seconds)
    if s < 1:
        return f"{s*1000:.0f} ms"
    if s < 60:
        return f"{s:.1f} s"
    m, sec = divmod(s, 60)
    if m < 60:
        return f"{int(m)}:{int(sec):02d}"
    h, m = divmod(m, 60)
    return f"{int(h)}:{int(m):02d}:{int(sec):02d}"


def _qt_enum(klass, name: str, *scoped_paths: str):
    """Resolve a Qt enum value across Qt5 (flat) and Qt6 (scoped).

    PyQt5 used flat enum members (``Qt.UserRole``).  PyQt6 (= QGIS 4)
    moved them under their typed inner class
    (``Qt.ItemDataRole.UserRole``).  Some PyQt6 builds expose the flat
    form for backward compatibility, others don't.

    This helper tries flat access first and falls back to each scoped
    path until one resolves.

    Example::

        UserRole = _qt_enum(Qt, 'UserRole', 'ItemDataRole.UserRole')
    """
    val = getattr(klass, name, None)
    if val is not None and not isinstance(val, type):
        return val
    for path in scoped_paths:
        obj = klass
        ok = True
        for part in path.split('.'):
            obj = getattr(obj, part, None)
            if obj is None:
                ok = False
                break
        if ok:
            return obj
    raise AttributeError(
        f"{klass.__name__} has no enum member {name!r} "
        f"(tried {[name, *scoped_paths]})"
    )
sys.path.append('.')
# Initialize Qt resources from file resources.py
from resources import *

# Import the code for the DockWidget
from compute.run_calculations import Calculation
from compute.calculation_task import CalculationTask
from compute.iwrap_convertion import write_iwrap_xml, parse_iwrap_xml
from geometries.handle_qgis_iface import HandleQGISIface
from omrat_utils.causation_factors import CausationFactors
from omrat_utils.handle_ais import AIS
from omrat_utils.handle_settings import DriftSettings
from omrat_utils.handle_traffic import Traffic
from omrat_utils.handle_distributions import Distributions
from omrat_utils.storage import Storage
from omrat_utils.handle_object import OObject
from omrat_utils.handle_ship_cat import ShipCategories
from omrat_utils.gather_data import GatherData
from omrat_widget import OMRATMainWidget
from geometries.drift import DriftCorridorGenerator
from geometries.drift_corridor_task_v2 import DriftCorridorTask


class OMRAT:
    """QGIS Plugin Implementation."""
    def __init__(self, iface:QgisInterface, testing:bool=False):
        """Constructor.

        :param iface: An interface instance that will be passed to this class
            which provides the hook by which you can manipulate the QGIS
            application at run time.
        :type iface: QgsInterface
        """
        # Save reference to the QGIS interface
        self.iface = iface
        self.testing = testing
        # initialize plugin directory
        self.plugin_dir = os.path.dirname(__file__)
        if not self.testing:
            # initialize locale
            locale: str = QSettings().value('locale/userLocale')[0:2]
            locale_path: str = os.path.join(
                self.plugin_dir,
                'i18n',
                f'OMRAT_{locale}.qm')

            if os.path.exists(locale_path):
                self.translator: QTranslator = QTranslator()
                self.translator.load(locale_path)
                QCoreApplication.installTranslator(self.translator)

        # Declare instance attributes
        self.actions: list[Any] = []
        self.menu: str = self.tr(u'&Open Maritime Risk Analysis Tool')
        self.toolbar: QToolBar = self.iface.addToolBar(u'OMRAT')
        self.toolbar.setObjectName(u'OMRAT')

            #print "** INITIALIZING OMRAT"
        self.pluginIsActive = False
        self.been_closed = False
        self.segment_id = 0
        self.traffic_data: dict[str, dict[str, dict[str, Any]]] = {}
        self.segment_data: dict[str, Any] = {}
        self.cur_route_id = 1
        self.main_widget = OMRATMainWidget(plugin=self)
        self.qgis_geoms: HandleQGISIface = HandleQGISIface(self)
        self.drift_settings: DriftSettings = DriftSettings(self)
        self.causation_f: CausationFactors = CausationFactors(self)
        self.calc: Calculation | None = Calculation(self)
        self.ship_cat = ShipCategories(self)
        self.ais = AIS(self)
        self.traffic = Traffic(self, self.main_widget)
        self.distributions = Distributions(self)
        self.object = OObject(self)
        self.drift_values = self.drift_settings.drift_values
        self.causation_values = self.causation_f.data
        self.drift_corridor_layers: list[QgsVectorLayer] = []  # Track corridor layers for cleanup
        # Layers loaded from a saved per-run GeoPackage via the
        # "Add results to map" button.  Tracked here so they get
        # removed when the plugin is unloaded or the project is
        # cleared -- otherwise they linger in the layer panel after
        # a Plugin Reloader cycle.
        self._history_layers: list[QgsVectorLayer] = []
        # Strong reference to the active CalculationTask while it runs.
        # Cleared from ``_on_calculation_finished`` /
        # ``_on_calculation_failed``.  See ``run_calculation`` for why.
        self._current_task = None

    def tr(self, message:str):
        """Get the translation for a string using Qt translation API.

        We implement this ourselves since we do not inherit QObject.

        :param message: String for translation.
        :type message: str, QString

        :returns: Translated version of message.
        :rtype: QString
        """
        return QCoreApplication.translate('OMRAT', message)

    def add_action(
        self,
        icon_path:str,
        text:str,
        callback: Callable[[], None],
        enabled_flag:bool=True,
        add_to_menu:bool=True,
        add_to_toolbar:bool=True,
        status_tip:str|None=None,
        whats_this:str|None=None,
        parent:QWidget|None=None) -> QAction:
        """Add a toolbar icon to the toolbar.

        :param icon_path: Path to the icon for this action. Can be a resource
            path (e.g. ':/plugins/foo/bar.png') or a normal file system path.
        :type icon_path: str

        :param text: Text that should be shown in menu items for this action.
        :type text: str

        :param callback: Function to be called when the action is triggered.
        :type callback: function

        :param enabled_flag: A flag indicating if the action should be enabled
            by default. Defaults to True.
        :type enabled_flag: bool

        :param add_to_menu: Flag indicating whether the action should also
            be added to the menu. Defaults to True.
        :type add_to_menu: bool

        :param add_to_toolbar: Flag indicating whether the action should also
            be added to the toolbar. Defaults to True.

        :param status_tip: Optional text to show in a popup when mouse pointer
            hovers over the action.
        :type status_tip: str

        :param parent: Parent widget for the new action. Defaults None.
        :type parent: QWidget

        :param whats_this: Optional text to show in the status bar when the
            mouse pointer hovers over the action.

        :returns: The action that was created. Note that the action is also
            added to self.actions list.
        :rtype: QAction
        """

        icon = QIcon(icon_path)
        action = QAction(icon, text, parent)
        action.triggered.connect(callback)
        action.setEnabled(enabled_flag)

        if status_tip is not None:
            action.setStatusTip(status_tip)

        if whats_this is not None:
            action.setWhatsThis(whats_this)

        if add_to_toolbar:
            self.toolbar.addAction(action)

        if add_to_menu:
            self.iface.addPluginToMenu(
                self.menu,
                action)

        self.actions.append(action)

        return action

    def initGui(self):
        """Create the menu entries and toolbar icons inside the QGIS GUI."""

        icon_path = self.plugin_dir + '/img/icon.png'
        self.add_action(
            icon_path,
            text=self.tr(u'Omrat'),
            callback=self.run,
            parent=self.iface.mainWindow())

        # Connect to project cleared signal to clean up drift corridors
        QgsProject.instance().cleared.connect(self._on_project_cleared)

    #--------------------------------------------------------------------------

    def _on_project_cleared(self):
        """Called when the QGIS project is cleared/reloaded.

        Clears drift corridor layers and any cached data.
        """
        QgsMessageLog.logMessage("Project cleared - cleaning up drift corridors", "OMRAT", Qgis.Info)
        self.drift_corridor_layers.clear()
        # The QGIS project drop also wipes any history-loaded layers,
        # so just drop our own references; no removeMapLayer needed.
        self._history_layers.clear()
        # Reset any task reference
        if hasattr(self, '_drift_task'):
            self._drift_task = None

    def _remove_history_layers_from_project(self) -> None:
        """Remove every history-loaded layer from the QGIS project.

        Called from ``clear_model`` and ``unload``.  Layers are
        appended to ``self._history_layers`` by
        :meth:`_add_selected_run_to_map`.
        """
        for layer in self._history_layers:
            try:
                project = QgsProject.instance()
                if project is not None and layer is not None:
                    project.removeMapLayer(layer.id())
            except Exception:
                pass
        self._history_layers.clear()

    def show_error_popup(self, message: str, function_name:str) -> None:
        msg_box = QMessageBox()
        msg_box.setIcon(QMessageBox.Icon.Critical)
        msg_box.setWindowTitle(self.tr(f"Failure in {function_name}"))
        msg_box.setText(self.tr("A failure occurred:"))
        msg_box.setInformativeText(message)
        msg_box.exec()

    def _confirm_before_load(self, action_label: str) -> str:
        """Show a confirmation dialog before loading/importing data.

        Returns 'clear', 'merge', or 'cancel'.
        If the current model is empty the dialog is skipped and 'clear' is returned.
        """
        has_data = (
            len(self.segment_data) > 0
            or len(self.traffic_data) > 0
            or self.main_widget.twRouteList.rowCount() > 0
            or self.main_widget.twDepthList.rowCount() > 0
            or self.main_widget.twObjectList.rowCount() > 0
        )
        if not has_data:
            return 'clear'

        msg_box = QMessageBox(self.main_widget)
        msg_box.setIcon(QMessageBox.Icon.Question)
        msg_box.setWindowTitle(self.tr(f'{action_label} Project'))
        msg_box.setText(self.tr(
            'The current model contains data.\n\n'
            'Do you want to clear all current data before loading,\n'
            'or merge the new data into the existing model?'
        ))
        btn_clear = msg_box.addButton(self.tr('Clear && Load'), QMessageBox.ButtonRole.AcceptRole)
        btn_merge = msg_box.addButton(self.tr('Merge'), QMessageBox.ButtonRole.ActionRole)
        msg_box.addButton(QMessageBox.StandardButton.Cancel)
        msg_box.setDefaultButton(btn_clear)
        msg_box.exec()

        clicked = msg_box.clickedButton()
        if clicked == btn_clear:
            return 'clear'
        elif clicked == btn_merge:
            return 'merge'
        return 'cancel'

    def clear_model(self) -> None:
        """Reset all model state and UI to defaults."""
        QgsMessageLog.logMessage("Clearing current model", "OMRAT", Qgis.Info)

        # --- 1. Clear QGIS layers ---
        self.qgis_geoms.clear()
        self.object.clear()

        # Remove drift corridor layers
        for layer in self.drift_corridor_layers:
            try:
                project = QgsProject.instance()
                if project is not None and layer is not None:
                    project.removeMapLayer(layer.id())
            except Exception:
                pass
        self.drift_corridor_layers.clear()

        # Remove any layers loaded from previous-run GeoPackages.
        self._remove_history_layers_from_project()

        # Reset cached Calculation result-layer attributes.  The plugin
        # no longer auto-adds memory layers to the canvas during a run --
        # results are persisted to per-run GeoPackages instead -- so
        # there is nothing to remove from the project here.  Just clear
        # the attribute references so a re-run starts clean.
        if self.calc is not None:
            for attr in (
                'allision_result_layer', 'grounding_result_layer',
                'collision_line_layer', 'collision_point_layer',
                'powered_grounding_layer', 'powered_allision_layer',
            ):
                setattr(self.calc, attr, None)

        # --- 2. Clear internal data structures ---
        self.traffic_data = {}
        self.segment_data = {}
        self.segment_id = 0
        self.cur_route_id = 1

        # Reset drift settings to defaults
        rose = {'0': .125, '45': .125, '90': .125, '135': .125,
                '180': .125, '225': .125, '270': .125, '315': .125}
        repair: dict[str, str | float | bool] = {
            'func': '', 'std': .95, 'loc': .2, 'scale': .85, 'use_lognormal': True
        }
        default_drift: dict[str, Any] = {
            'drift_p': 1, 'anchor_p': .95, 'anchor_d': 7,
            'speed': 1 * 3600 / 1852, 'rose': rose, 'repair': repair
        }
        self.drift_values = default_drift
        self.drift_settings.drift_values = default_drift

        # Reset causation factors to defaults.  Use the full default
        # set defined in ``CausationFactors.__init__`` so collision
        # keys (``headon``, ``overtaking``, ...) survive Clear.
        from omrat_utils.causation_factors import CausationFactors
        self.causation_f.data = CausationFactors(self).data
        self.causation_values = self.causation_f.data

        # Reset distributions tracking
        self.distributions.last_id = '1'
        self.distributions.new_id = '1'

        # Prevent traffic table updates during clearing
        self.traffic.run_update = False
        self.traffic.traffic_data = self.traffic_data
        self.traffic.c_seg = ''

        # --- 3. Clear UI widgets ---
        self.reset_route_table()

        self.main_widget.twDepthList.setRowCount(0)
        self.main_widget.twObjectList.setRowCount(0)

        self.main_widget.cbTrafficSelectSeg.clear()
        self.main_widget.cbTrafficDirectionSelect.clear()
        self.traffic.set_table_headings()

        # Distribution fields
        dist_widgets = [
            self.main_widget.leNormMean1_1, self.main_widget.leNormMean1_2,
            self.main_widget.leNormMean1_3, self.main_widget.leNormMean2_1,
            self.main_widget.leNormMean2_2, self.main_widget.leNormMean2_3,
            self.main_widget.leNormStd1_1, self.main_widget.leNormStd1_2,
            self.main_widget.leNormStd1_3, self.main_widget.leNormStd2_1,
            self.main_widget.leNormStd2_2, self.main_widget.leNormStd2_3,
            self.main_widget.leNormWeight1_1, self.main_widget.leNormWeight1_2,
            self.main_widget.leNormWeight1_3, self.main_widget.leNormWeight2_1,
            self.main_widget.leNormWeight2_2, self.main_widget.leNormWeight2_3,
            self.main_widget.leUniformMin1, self.main_widget.leUniformMax1,
            self.main_widget.leUniformMin2, self.main_widget.leUniformMax2,
        ]
        for w in dist_widgets:
            w.setText('0')
        self.main_widget.sbUniformP1.setValue(0)
        self.main_widget.sbUniformP2.setValue(0)

        # Result fields
        result_fields = [
            self.main_widget.LEPDriftAllision,
            self.main_widget.LEPPoweredAllision,
            self.main_widget.LEPDriftingGrounding,
            self.main_widget.LEPPoweredGrounding,
            self.main_widget.LEPOvertakingCollision,
            self.main_widget.LEPHeadOnCollision,
            self.main_widget.LEPCrossingCollision,
            self.main_widget.LEPMergingCollision,
        ]
        for field in result_fields:
            field.setText('')

        self.main_widget.LEModelName.setText('')
        self.main_widget.TWPreviousRuns.setRowCount(0)
        self.main_widget.laDir1.setText('')
        self.main_widget.laDir2.setText('')

        if hasattr(self.main_widget, 'label_drift_status'):
            self.main_widget.label_drift_status.setText('')

        # Remove distribution plot canvas if present
        if self.distributions.canvas is not None:
            try:
                self.main_widget.DistributionWidget.removeWidget(self.distributions.canvas)
                self.distributions.canvas.deleteLater()
                self.distributions.canvas = None
            except Exception:
                pass

        # Re-enable traffic updates
        self.traffic.run_update = True

        self.iface.mapCanvas().refresh()
        QgsMessageLog.logMessage("Model cleared successfully", "OMRAT", Qgis.Info)

    def onClosePlugin(self):
        """Cleanup necessary items here when plugin main_widget is closed"""

        # Disconnect signals
        if self.main_widget.is_active:
            self.main_widget.pbAddRoute.clicked.disconnect()
            self.main_widget.pbStopRoute.clicked.disconnect()
            self.main_widget.pbLoadRoute.clicked.disconnect()
            self.main_widget.pbRemoveRoute.clicked.disconnect()
            self.main_widget.pbAddSimpleDepth.clicked.disconnect()
            self.main_widget.pbLoadDepth.clicked.disconnect()
            self.main_widget.pbRemoveDepth.clicked.disconnect()
            self.main_widget.pbAddSimpleObject.clicked.disconnect()
            self.main_widget.pbLoadObject.clicked.disconnect()
            self.main_widget.pbRemoveObject.clicked.disconnect()
            self.main_widget.pbRunModel.clicked.disconnect()

        # Set plugin state to inactive
        self.pluginIsActive = False
        self.main_widget.is_active = False

    def unload(self):
        """Removes the plugin menu item and icon from QGIS GUI."""
        # Disconnect project cleared signal
        try:
            QgsProject.instance().cleared.disconnect(self._on_project_cleared)
        except Exception:
            pass

        # Call the unload method of Traffic
        self.traffic.unload()

        # Call the unload method of AIS
        self.ais.unload()

        # Unload geometries
        self.object.unload()
        self.qgis_geoms.unload()

        # Remove drift corridor layers
        for layer in self.drift_corridor_layers:
            try:
                project = QgsProject.instance()
                if project is not None and layer is not None:
                    project.removeMapLayer(layer.id())
            except Exception:
                pass
        self.drift_corridor_layers.clear()

        # Remove any layers loaded from previous-run GeoPackages.
        self._remove_history_layers_from_project()

        # Remove result layers (allision/grounding/collision/powered visualization)
        if self.calc is not None:
            result_layer_attrs = [
                'allision_result_layer', 'grounding_result_layer',
                'collision_line_layer', 'collision_point_layer',
                'powered_grounding_layer', 'powered_allision_layer',
            ]
            for attr in result_layer_attrs:
                layer = getattr(self.calc, attr, None)
                try:
                    project = QgsProject.instance()
                    if project is not None and layer is not None:
                        project.removeMapLayer(layer.id())
                except Exception:
                    pass
                setattr(self.calc, attr, None)

        # Remove actions from menu and toolbar
        for action in self.actions:
            self.iface.removePluginMenu(
                self.tr(u'&Open Maritime Risk Analysis Tool'),
                action)
            self.iface.removeToolBarIcon(action)

        # Remove the toolbar
        if hasattr(self, 'toolbar'):
            del self.toolbar    

        # Disconnect signals and clean up dock widget
        try:
            self.onClosePlugin()
        except Exception as e:
            print(f"Error during onClosePlugin: {e}")
        try:
            self.main_widget.pbGetGebcoDephts.clicked.disconnect()
            self.main_widget.PBUpdateDepthIntervals.clicked.disconnect()
            self.main_widget.deleteLater()
        except TypeError:
            pass
        gc.collect()
        # Remove all plugin modules from sys.modules to ensure fresh code on reload
        for folder in ['omrat_utils', 'compute', 'geometries']:
            to_remove = [m for m in sys.modules if m.startswith(folder)]
            for m in to_remove:
                QgsMessageLog.logMessage(f"{m} unloaded", "OMRAT", Qgis.Info)
                del sys.modules[m]
        self.been_close = True
        QgsMessageLog.logMessage("Plugin unloaded", "OMRAT", Qgis.Info)
        
    def point4326_from_wkt(self, coord_str:str, crs:str) -> QgsPoint:
        if '(' in coord_str:
            coord_str = coord_str.split('(')[1].split(')')[0]
        coords_str: list[str] = coord_str.split(' ')
        coords: list[float] = [float(coord) for coord in coords_str]
        q_point = QgsPoint(coords[0], coords[1])
        return q_point
        
    def load_lines(self, data:dict[str, dict[str, Any]]) -> None:
        self.segment_id = 0
        crs = self.iface.mapCanvas().mapSettings().destinationCrs().authid()
        for key, seg_data in data["segment_data"].items():
            name = f"Segment {seg_data['Route_Id']} - {seg_data['Segment_Id']}"
            vl = QgsVectorLayer(f"LineString?crs=EPSG:4326", name, "memory")
            QgsProject.instance().addMapLayer(vl)
            
            self.segment_id += 1
            if not vl.isEditable():
                vl.startEditing()
            fields = self.qgis_geoms.create_fields()
            provider: QgsVectorDataProvider | None = vl.dataProvider()
            if provider is None:
                return
            provider.addAttributes(fields.toList())
            vl.updateFields()

            # Create feature
            fet = QgsFeature(fields)
            start = self.point4326_from_wkt(seg_data["Start_Point"], crs)
            end = self.point4326_from_wkt(seg_data["End_Point"], crs)
            fet.setGeometry(QgsLineString([start, end]))
            fet.setAttributes([seg_data["Segment_Id"], seg_data["Route_Id"], seg_data["Start_Point"], 
                               seg_data["End_Point"], f"LEG_{seg_data['Segment_Id']}_{seg_data['Route_Id']}"])
            self.qgis_geoms.style_layer(vl)
            # Ensure feature id matches Segment_Id so geometry edits update correct row
            try:
                fid = int(str(seg_data["Segment_Id"]))
            except Exception:
                fid = self.segment_id
            fet.setId(fid)
            provider.addFeature(fet)
            self.qgis_geoms.label_layer(vl)
            vl.updateExtents()
            
            # Validate geometry
            if not fet.geometry().isGeosValid():
                print(f"Invalid geometry for segment {seg_data['Segment_Id']}")
                continue
            # Keep layer editable and connect geometry change handler
            edit_buffer = vl.editBuffer()
            if edit_buffer is not None:
                edit_buffer.geometryChanged.connect(partial(self.qgis_geoms.on_geometry_changed_wrapper, fid))
            vl.triggerRepaint()
            self.qgis_geoms.vector_layers.append(vl)

            # Preserve direction labels mapping for AIS
            try:
                self.qgis_geoms.leg_dirs[str(seg_data["Segment_Id"])]=list(seg_data.get("Dirs", []))
            except Exception:
                pass

        # After loading, ensure subsequent new segments use a higher id
        try:
            max_id = max(int(str(v.get('Segment_Id'))) for v in data['segment_data'].values())
        except Exception:
            max_id = self.segment_id
        self.qgis_geoms.segment_id = max_id
        self.segment_id = max_id
        # Refresh map canvas
        self.iface.mapCanvas().refresh()    

    def reset_route_table(self) -> None:
        self.main_widget.twRouteList.setColumnCount(7)
        self.main_widget.twRouteList.setHorizontalHeaderLabels(['Segment_Id', 'Route_Id', 'Leg_name',
                                                            'Start_Point', 'End_Point', 'Width', 'Update AIS'])
        self.main_widget.twRouteList.setColumnWidth(0, 75)
        self.main_widget.twRouteList.setColumnWidth(1, 75)
        self.main_widget.twRouteList.setColumnWidth(2, 75)
        self.main_widget.twRouteList.setColumnWidth(3, 125)
        self.main_widget.twRouteList.setColumnWidth(4, 125)
        self.main_widget.twRouteList.setColumnWidth(5, 75)
        self.main_widget.twRouteList.setColumnWidth(6, 75)
        self.main_widget.twRouteList.setRowCount(0)
    
    def run_traffic_module(self) -> None:
        self.traffic.traffic_data = self.traffic_data
        self.traffic.fill_cbTrafficSelectSeg()
        self.traffic.update_direction_select()
        
    def save_work(self):
        store = Storage(self)
        store.store_all()

    def load_work(self):
        store = Storage(self)
        file_path = store.select_file()
        if not file_path:
            return
        choice = self._confirm_before_load('Load')
        if choice == 'cancel':
            return
        if choice == 'clear':
            self.clear_model()
        store.load_from_path(file_path)
               
    def run_calculation(self):
        """Run calculation in a background QgsTask for better UI responsiveness."""
        gd = GatherData(self)
        data = gd.get_all_for_save()
        if self.calc is None:
            return

        # Refuse to run unless an output folder is configured.  The
        # button should already be disabled, but defend against direct
        # calls / keyboard shortcuts.
        out_dir = self._get_output_dir()
        if out_dir is None:
            QMessageBox.warning(
                self.main_widget, 'Output folder required',
                "Pick an output folder on the Run Analysis tab before "
                "running the model.  Result layers from each run are "
                "written to that folder as a single .gpkg file.",
            )
            return

        # Create and configure the background task
        task = CalculationTask(
            description='OMRAT: Starting calculations...',
            calc_object=self.calc,
            data=data
        )

        # Connect signals for progress updates and completion
        task.progress_updated.connect(self._on_calculation_progress)
        task.calculation_finished.connect(self._on_calculation_finished)
        task.calculation_failed.connect(self._on_calculation_failed)

        # CRITICAL: keep a strong Python reference to the task object.
        # QGIS's task manager only holds a C++ pointer; without a Python
        # reference, the QObject wrapper can be garbage-collected before
        # ``finished()`` fires, which silently breaks our signals.
        self._current_task = task

        # Track elapsed time for the run-history "Duration" column.
        self._run_started_at = time.monotonic()

        # Add task to QGIS task manager (runs in background)
        QgsApplication.taskManager().addTask(task)

        QgsMessageLog.logMessage(
            'Calculation started in background. Check task manager for progress.',
            'OMRAT',
            Qgis.Info
        )

    def _on_calculation_progress(self, completed: int, total: int, message: str) -> None:
        """Handle progress updates from the calculation task."""
        # Log progress to QGIS message log
        QgsMessageLog.logMessage(
            f"Progress: {completed}/{total} - {message}",
            'OMRAT',
            Qgis.Info
        )

    def _on_calculation_finished(self, calc_object: Any) -> None:
        """Handle successful completion of calculation."""
        QgsMessageLog.logMessage(
            'Calculation completed successfully!',
            'OMRAT',
            Qgis.Success
        )
        # Auto-save the run to the history GeoPackage so the user can
        # reload / compare it later.  Failures here are non-fatal --
        # the on-canvas results and line-edits are already populated.
        try:
            self._auto_save_run(calc_object)
        except Exception as e:
            import traceback
            QgsMessageLog.logMessage(
                f'Failed to save run history: {e}\n{traceback.format_exc()}',
                'OMRAT', Qgis.Warning,
            )
        try:
            self.refresh_previous_runs_table()
        except Exception as e:
            import traceback
            QgsMessageLog.logMessage(
                f'Failed to refresh previous-runs table: {e}\n'
                f'{traceback.format_exc()}',
                'OMRAT', Qgis.Warning,
            )
        # Drop the strong reference now that the task has finished.
        self._current_task = None

    def _on_calculation_failed(self, error_msg: str) -> None:
        """Handle calculation failure (overrides earlier definition).

        Keeps the message-box behaviour but also drops the strong task
        reference so the next run starts clean.
        """
        QgsMessageLog.logMessage(
            f'Calculation failed: {error_msg}',
            'OMRAT', Qgis.Critical,
        )
        QMessageBox.critical(
            None, 'Calculation Failed',
            f'The calculation failed with the following error:\n\n{error_msg}',
        )
        self._current_task = None

    def _auto_save_run(self, calc_object: Any) -> None:
        """Write per-run GeoPackage + record metadata in the master DB."""
        from omrat_utils.run_history import (
            RunHistory, totals_from_calc, make_run_filename,
        )
        from omrat_utils.run_persistence import write_run_results

        out_dir = self._get_output_dir()
        if out_dir is None:
            raise RuntimeError(
                "No output folder configured on the Run Analysis tab. "
                "Set one and re-run."
            )

        name = ''
        try:
            name = self.main_widget.LEModelName.text().strip()
        except Exception:
            pass
        if not name:
            name = time.strftime('run_%Y%m%d_%H%M%S')

        # Per-run GeoPackage path: <out_dir>/<slug(name)>_<timestamp>.gpkg
        ts_struct = time.localtime()
        ts_text = time.strftime('%Y-%m-%d %H:%M:%S', ts_struct)
        filename = make_run_filename(name, ts_struct)
        gpkg_path = Path(out_dir) / filename

        QgsMessageLog.logMessage(
            f"Writing per-run results to {gpkg_path}", 'OMRAT', Qgis.Info,
        )
        try:
            written_layers = write_run_results(
                calc_object, gpkg_path,
                structures=getattr(calc_object, '_last_structures', None),
                depths=getattr(calc_object, '_last_depths', None),
                depths_original=getattr(
                    calc_object, '_last_depths_original', None,
                ),
                segment_data=self.segment_data,
            )
        except Exception as exc:
            import traceback
            QgsMessageLog.logMessage(
                f"GeoPackage write failed: {exc}\n{traceback.format_exc()}",
                'OMRAT', Qgis.Critical,
            )
            written_layers = []

        # Compute elapsed time, if we tracked the start.
        duration = None
        try:
            start = getattr(self, '_run_started_at', None)
            if start is not None:
                duration = float(time.monotonic() - start)
        except Exception:
            duration = None

        history = RunHistory()
        run_id = history.save_run(
            name=name,
            timestamp=ts_text,
            duration_seconds=duration,
            totals=totals_from_calc(calc_object),
            output_dir=str(out_dir),
            output_filename=filename,
        )
        dur_text = (
            f"duration={duration:.1f}s" if duration is not None else "duration=n/a"
        )
        QgsMessageLog.logMessage(
            f"Run '{name}' saved (run_id={run_id}, layers={len(written_layers)}, "
            f"{dur_text}).  GeoPackage: {gpkg_path}",
            'OMRAT', Qgis.Success,
        )

    # ------------------------------------------------------------------
    # Output folder + Run Model gating
    # ------------------------------------------------------------------
    def _get_output_dir(self) -> Path | None:
        """Resolved + validated output folder, or ``None`` if not set."""
        try:
            text = (self.main_widget.LEReportPath.text() or '').strip()
        except Exception:
            text = ''
        if not text:
            return None
        path = Path(text)
        if not path.is_dir():
            return None
        return path

    def choose_output_folder(self) -> None:
        """Pop a folder picker into ``LEReportPath``.

        Each Run Model writes its result layers into a single
        ``<name>_<timestamp>.gpkg`` file in this folder.
        """
        initial = ''
        try:
            initial = (self.main_widget.LEReportPath.text() or '').strip()
        except Exception:
            initial = ''
        if not initial:
            initial = str(Path.home())
        chosen = QFileDialog.getExistingDirectory(
            self.main_widget,
            'Select output folder for OMRAT result GeoPackages',
            initial,
        )
        if chosen:
            self.main_widget.LEReportPath.setText(chosen)
            try:
                QSettings().setValue('omrat/output_dir', chosen)
            except Exception:
                pass
            self._update_run_model_enabled()

    def _restore_output_dir(self) -> None:
        """At plugin start, restore the last-used output folder."""
        try:
            value = QSettings().value('omrat/output_dir', '', type=str)
            if value:
                self.main_widget.LEReportPath.setText(value)
        except Exception:
            pass

    def _update_run_model_enabled(self) -> None:
        """Run Model is gated on a writable output folder being selected.

        When disabled the button shows the QSS ``:disabled`` style
        (gray) defined in ``omrat_base.ui`` plus a tooltip telling the
        user what to do.  When enabled the tooltip is cleared so the
        regular blue look takes over without a stale hint.
        """
        try:
            btn = self.main_widget.pbRunModel
        except Exception:
            return
        enabled = self._get_output_dir() is not None
        try:
            btn.setEnabled(enabled)
            btn.setToolTip(
                '' if enabled else 'You must select a folder first'
            )
        except Exception:
            pass

    # ------------------------------------------------------------------
    # Previous-runs table on the Results tab
    # ------------------------------------------------------------------
    def _setup_previous_runs_table(self) -> None:
        """Configure ``TWPreviousRuns`` columns + add the load button.

        Columns are slim now -- Name / Date / Duration only.  Detail
        probabilities flow into the ``LEPDriftAllision`` etc. fields
        below the table when the user picks a row (or several rows).
        """
        from qgis.PyQt import QtCore, QtWidgets
        Qt = QtCore.Qt
        AIV = QtWidgets.QAbstractItemView
        tw = self.main_widget.TWPreviousRuns

        headers = ['Name', 'Date', 'Duration']
        tw.setColumnCount(len(headers))
        tw.setHorizontalHeaderLabels(headers)
        tw.horizontalHeader().setStretchLastSection(True)
        tw.setSelectionBehavior(_qt_enum(
            AIV, 'SelectRows', 'SelectionBehavior.SelectRows',
        ))
        tw.setSelectionMode(_qt_enum(
            AIV, 'ExtendedSelection', 'SelectionMode.ExtendedSelection',
        ))
        tw.setEditTriggers(_qt_enum(
            AIV, 'NoEditTriggers', 'EditTrigger.NoEditTriggers',
        ))
        tw.setContextMenuPolicy(_qt_enum(
            Qt, 'CustomContextMenu', 'ContextMenuPolicy.CustomContextMenu',
        ))
        tw.customContextMenuRequested.connect(
            self._previous_runs_context_menu,
        )

        # Drive the result-LineEdits from the row selection.
        sel_model = tw.selectionModel()
        if sel_model is not None:
            sel_model.selectionChanged.connect(
                lambda *_: self._on_previous_runs_selection_changed(),
            )

        # Add an "Add results to map" button programmatically next to
        # the table.  We append it to the same vertical layout the
        # table sits in so it appears just below.
        try:
            parent_widget = tw.parentWidget()
            layout = parent_widget.layout() if parent_widget is not None else None
            btn = QtWidgets.QPushButton('Add selected run results to map')
            btn.setEnabled(False)
            btn.clicked.connect(self._add_selected_run_to_map)
            self._pb_add_run_to_map = btn
            if layout is not None:
                idx = layout.indexOf(tw)
                if idx >= 0:
                    layout.insertWidget(idx + 1, btn)
                else:
                    layout.addWidget(btn)
        except Exception as exc:
            QgsMessageLog.logMessage(
                f'Could not add "Add to map" button: {exc}',
                'OMRAT', Qgis.Warning,
            )

    def refresh_previous_runs_table(self) -> None:
        """Reload ``TWPreviousRuns`` from the master history DB."""
        from omrat_utils.run_history import RunHistory
        from qgis.PyQt import QtWidgets, QtCore
        user_role = _qt_enum(QtCore.Qt, 'UserRole', 'ItemDataRole.UserRole')
        tw = self.main_widget.TWPreviousRuns
        try:
            runs = RunHistory().list_runs()
        except Exception as e:
            QgsMessageLog.logMessage(
                f'Could not read run history: {e}', 'OMRAT', Qgis.Warning,
            )
            tw.setRowCount(0)
            return
        tw.setRowCount(len(runs))
        for i, run in enumerate(runs):
            cells = [
                run.name,
                run.timestamp,
                _format_duration(run.duration_seconds),
            ]
            for j, text in enumerate(cells):
                item = QtWidgets.QTableWidgetItem(text)
                if j == 0:
                    item.setData(user_role, run.run_id)
                tw.setItem(i, j, item)
        tw.resizeColumnsToContents()

    def _selected_run_ids(self) -> list[int]:
        from qgis.PyQt.QtCore import Qt
        user_role = _qt_enum(Qt, 'UserRole', 'ItemDataRole.UserRole')
        tw = self.main_widget.TWPreviousRuns
        out: list[int] = []
        seen: set[int] = set()
        sel_model = tw.selectionModel()
        if sel_model is None:
            return out
        for idx in sel_model.selectedRows():
            item = tw.item(idx.row(), 0)
            if item is None:
                continue
            run_id = item.data(user_role)
            if run_id is not None and int(run_id) not in seen:
                out.append(int(run_id))
                seen.add(int(run_id))
        return out

    def _on_previous_runs_selection_changed(self) -> None:
        """Populate ``LEP*`` fields from the selected run(s)."""
        from omrat_utils.run_history import RunHistory
        run_ids = self._selected_run_ids()
        # Keep the "Add to map" button gated on having a single
        # selected run.
        try:
            self._pb_add_run_to_map.setEnabled(len(run_ids) == 1)
        except Exception:
            pass
        if not run_ids:
            return
        try:
            runs = RunHistory().compare_runs(run_ids)
        except Exception as e:
            QgsMessageLog.logMessage(
                f'Could not load previous-run details: {e}',
                'OMRAT', Qgis.Warning,
            )
            return
        if not runs:
            return
        self._fill_result_fields_from_runs(runs)

    def _fill_result_fields_from_runs(self, runs) -> None:
        """Write the per-run probabilities into the LEP*Collision /
        LEPDriftAllision etc. line-edits.

        Single run  -> LineEdit shows just the number.
        Multiple    -> LineEdit shows ``v1 | v2 (Δ+12.3%) | v3 (Δ-4.5%)``
                       where the first value is the baseline and each
                       subsequent value carries its relative difference.
        """
        # Map LineEdit objectName -> totals key from RunMeta.totals_dict.
        field_map = [
            ('LEPDriftAllision',     'drift_allision'),
            ('LEPDriftingGrounding', 'drift_grounding'),
            ('LEPPoweredAllision',   'powered_allision'),
            ('LEPPoweredGrounding',  'powered_grounding'),
            ('LEPHeadOnCollision',   'head_on'),
            ('LEPOvertakingCollision', 'overtaking'),
            ('LEPCrossingCollision', 'crossing'),
            ('LEPMergingCollision',  'bend'),
        ]
        for widget_name, key in field_map:
            le = getattr(self.main_widget, widget_name, None)
            if le is None:
                continue
            text = self._format_field_for_runs(runs, key)
            try:
                le.setText(text)
            except Exception:
                pass

    @staticmethod
    def _format_field_for_runs(runs, key: str) -> str:
        if not runs:
            return ''
        baseline = runs[0].totals_dict().get(key)
        parts = [_fmt(baseline)]
        for r in runs[1:]:
            v = r.totals_dict().get(key)
            base_text = _fmt(v)
            if v is None or baseline in (None, 0):
                parts.append(base_text)
            else:
                rel = (v - baseline) / baseline * 100.0
                parts.append(f"{base_text} (Δ{rel:+.1f}%)")
        return ' | '.join(parts)

    def _previous_runs_context_menu(self, pos) -> None:
        from qgis.PyQt.QtWidgets import QMenu
        tw = self.main_widget.TWPreviousRuns
        run_ids = self._selected_run_ids()
        if not run_ids:
            return
        menu = QMenu(tw)
        if len(run_ids) == 1:
            menu.addAction(
                'Add results to map',
                lambda: self._add_selected_run_to_map(),
            )
        menu.addSeparator()
        menu.addAction(
            'Delete from history',
            lambda: self._delete_runs(run_ids),
        )
        menu.addAction(
            'Delete from history + remove .gpkg file',
            lambda: self._delete_runs(run_ids, delete_gpkg=True),
        )
        menu.exec_(tw.viewport().mapToGlobal(pos))

    def _add_selected_run_to_map(self) -> None:
        """Load the selected run's per-run GeoPackage onto the canvas."""
        from omrat_utils.run_history import RunHistory
        from omrat_utils.run_persistence import load_run_results_to_map
        run_ids = self._selected_run_ids()
        if len(run_ids) != 1:
            return
        run = RunHistory().get_run(run_ids[0])
        if run is None:
            return
        gpkg_path = run.gpkg_path()
        if gpkg_path is None or not gpkg_path.is_file():
            QMessageBox.warning(
                self.main_widget, 'GeoPackage missing',
                f"The result file for run '{run.name}' could not be "
                f"found.\nExpected at: {gpkg_path}",
            )
            return
        try:
            new_layers = load_run_results_to_map(gpkg_path, run.name)
        except Exception as e:
            import traceback
            QgsMessageLog.logMessage(
                f'Failed to load run on map: {e}\n{traceback.format_exc()}',
                'OMRAT', Qgis.Warning,
            )
            return
        # Track the layers so plugin unload / clear_model can remove
        # them from the QGIS project.  Without this they survive a
        # Plugin Reloader cycle and pile up across runs.
        for layer in new_layers:
            if layer is not None:
                self._history_layers.append(layer)

    def _delete_runs(self, run_ids: list[int], *, delete_gpkg: bool = False) -> None:
        from omrat_utils.run_history import RunHistory
        yes = _qt_enum(QMessageBox, 'Yes', 'StandardButton.Yes')
        msg = (
            f'Delete {len(run_ids)} run(s) from history'
            + (' AND remove the .gpkg files from disk?' if delete_gpkg else '?')
        )
        confirm = QMessageBox.question(
            self.main_widget, 'Delete runs', msg,
        )
        if confirm != yes:
            return
        history = RunHistory()
        for run_id in run_ids:
            try:
                history.delete_run(run_id, delete_gpkg=delete_gpkg)
            except Exception as e:
                QgsMessageLog.logMessage(
                    f'Failed to delete run {run_id}: {e}', 'OMRAT', Qgis.Warning,
                )
        self.refresh_previous_runs_table()

    def open_previous_runs_dialog(self) -> None:
        """File menu shortcut -- focuses the Results tab + the runs table."""
        try:
            self.main_widget.tabWidget.setCurrentWidget(
                self.main_widget.tab_9,
            )
        except Exception:
            pass
        self.refresh_previous_runs_table()

    def choose_report_path(self) -> None:
        """Open a Save File dialog and store path in LEReportPath."""
        # Initial directory & filename from current field if present
        initial = ''
        try:
            if hasattr(self.main_widget, 'LEReportPath') and self.main_widget.LEReportPath is not None:
                initial = self.main_widget.LEReportPath.text() or ''
        except Exception:
            initial = ''
        filename, _ = QFileDialog.getSaveFileName(
            self.main_widget,
            self.tr('Select report file'),
            initial,
            'Markdown Files (*.md);;All Files (*.*)'
        )
        if filename:
            try:
                self.main_widget.LEReportPath.setText(filename)
            except Exception:
                pass

    def export_to_iwrap(self) -> None:
        """Export current project data to IWRAP XML format."""
        try:
            # Get file path from user
            filename, _ = QFileDialog.getSaveFileName(
                self.main_widget,
                self.tr('Export to IWRAP XML'),
                '',
                'IWRAP XML Files (*.xml);;All Files (*.*)'
            )

            if not filename:
                return  # User cancelled

            # Ensure .xml extension
            if not filename.lower().endswith('.xml'):
                filename += '.xml'

            # Gather current project data
            import tempfile
            import json
            gd = GatherData(self)
            data = gd.get_all_for_save()

            # Save to temporary JSON file
            temp_json = tempfile.NamedTemporaryFile(mode='w', suffix='.omrat', delete=False, encoding='utf-8')
            temp_json_path = temp_json.name
            try:
                json.dump(data, temp_json, indent=2)
                temp_json.close()

                # Convert to IWRAP XML
                write_iwrap_xml(temp_json_path, filename)

                # Show success message
                QMessageBox.information(
                    self.main_widget,
                    self.tr('Export Successful'),
                    self.tr(f'Project successfully exported to:\n{filename}')
                )
            finally:
                # Clean up temporary file
                try:
                    if os.path.exists(temp_json_path):
                        os.remove(temp_json_path)
                except Exception:
                    pass

        except Exception as e:
            self.show_error_popup(
                str(e),
                'export_to_iwrap'
            )

    def import_from_iwrap(self) -> None:
        """Import project data from IWRAP XML format."""
        try:
            # Get file path from user
            filename, _ = QFileDialog.getOpenFileName(
                self.main_widget,
                self.tr('Import from IWRAP XML'),
                '',
                'IWRAP XML Files (*.xml);;All Files (*.*)'
            )

            if not filename:
                return  # User cancelled

            # Check if file exists
            if not os.path.exists(filename):
                QMessageBox.warning(
                    self.main_widget,
                    self.tr('File Not Found'),
                    self.tr(f'The file does not exist:\n{filename}')
                )
                return

            # Ask whether to clear current data before importing
            choice = self._confirm_before_load('Import')
            if choice == 'cancel':
                return
            if choice == 'clear':
                self.clear_model()

            # Parse IWRAP XML to .omrat dictionary format
            import json
            # Enable debug output to help diagnose issues
            print("\n" + "="*70)
            print("IMPORTING IWRAP XML FILE")
            print("="*70)
            data = parse_iwrap_xml(filename, debug=True)
            print("\n" + "="*70)
            print("IMPORT PARSING COMPLETE")
            print("="*70 + "\n")

            # Validate and normalize the imported data
            from omrat_utils.storage import Storage
            storage = Storage(self)
            data = storage._normalize_legacy_to_schema(data)

            # Validate schema
            try:
                from omrat_utils.validate_data import RootModelSchema
                RootModelSchema.model_validate(data)
            except Exception as validation_error:
                QMessageBox.warning(
                    self.main_widget,
                    self.tr('Import Validation Error'),
                    self.tr(f'The imported data failed validation:\n{str(validation_error)}\n\nThe data will still be loaded, but there may be missing or incorrect values.')
                )

            # Load the imported data into the plugin
            # This will:
            # 1. Populate ship categories (rebuilds traffic table dimensions)
            # 2. Load traffic_data, segment_data, drift settings
            # 3. Populate depths, objects, routes
            gd = GatherData(self)
            gd.populate(data)

            # Update Traffic class reference to the new traffic_data
            # (populate() assigns a new dict, so we need to update the reference)
            self.traffic.traffic_data = self.traffic_data

            # Refresh the traffic table display with imported data
            self.traffic.update_direction_select()
            self.traffic.update_traffic_tbl('segment')

            # Show success message
            QMessageBox.information(
                self.main_widget,
                self.tr('Import Successful'),
                self.tr(f'Project successfully imported from:\n{filename}\n\nProject: {data.get("project_name", "Unknown")}\nSegments: {len(data.get("segment_data", {}))}')
            )

        except Exception as e:
            self.show_error_popup(
                str(e),
                'import_from_iwrap'
            )

    def show_drift_allision(self):
        gd = GatherData(self)
        data = gd.get_all_for_save()
        if self.calc is not None:
            self.calc.run_drift_visualization(data)

    def show_powered_allision(self):
        """Show interactive Cat II powered allision (missed turn -> hits object) visualisation."""
        gd = GatherData(self)
        data = gd.get_all_for_save()
        if self.calc is not None:
            self.calc.run_powered_allision_visualization(data)

    def show_powered_grounding(self):
        """Show interactive Cat II powered grounding (missed turn -> hits depth area) visualisation."""
        gd = GatherData(self)
        data = gd.get_all_for_save()
        if self.calc is not None:
            self.calc.run_powered_grounding_visualization(data)

    # ------------------------------------------------------------------
    # Placeholder slots for "View" buttons whose visualisation pipelines
    # have not been implemented yet.  The buttons exist in omrat_base.ui
    # and are wired up below so users see a clear "not implemented" message
    # instead of a silent dead button.  Replace the body of each method
    # with the real visualisation call when the pipeline is ready.
    # TODO(view-buttons):
    #   - show_drift_grounding:    drifting-grounding overlap visualisation
    #     (the drifting counterpart to show_powered_grounding)
    #   - show_overtaking_collision, show_head_on_collision,
    #     show_crossing_collision, show_merging_collision:
    #     ship-ship collision visualisations by encounter type
    # ------------------------------------------------------------------
    def _view_button_not_implemented(self, name: str) -> None:
        QgsMessageLog.logMessage(
            f"'{name}' visualisation is not implemented yet (TODO).",
            'OMRAT',
            Qgis.Warning,
        )

    def show_drift_grounding(self):
        self._view_button_not_implemented('Drift grounding')

    def show_overtaking_collision(self):
        self._view_button_not_implemented('Overtaking collision')

    def show_head_on_collision(self):
        self._view_button_not_implemented('Head-on collision')

    def show_crossing_collision(self):
        self._view_button_not_implemented('Crossing collision')

    def show_merging_collision(self):
        self._view_button_not_implemented('Merging collision')

    def open_drift(self):
        self.drift_settings.run()

    def open_ship_categories(self):
        self.ship_cat.run()

    def open_causation_factors(self):
        self.causation_f.run()

    def stop_route(self):
        """Stop the current route, commit changes, and set the active tool to 'Pan Map'."""
        # Disable the "Stop Route" button
        self.main_widget.pbStopRoute.setEnabled(False)

        for layer in self.qgis_geoms.vector_layers:
            if layer.isEditable():
                layer.commitChanges()  # Save changes
            layer.triggerRepaint()  # Ensure the layer is refreshed

        # Clear the current start point and increment the route ID
        self.cur_route_id += 1

        # Set the active tool to "Pan Map"
        canvas: QgsMapCanvas | None = self.iface.mapCanvas()
        map_tool_pan = QgsMapToolPan(canvas)
        if canvas is not None:
            canvas.setMapTool(map_tool_pan)
        actionPan: QAction | None = self.iface.actionPan()
        if actionPan is not None:
            actionPan.trigger()
        
    def ais_settings(self) -> None:
        self.ais.run()
        
    def update_ais(self) -> None:
        self.ais.update_legs()
    
    def remove_route(self)-> None:
        # To implement
        pass
    
    def load_route(self)-> None:
        # To implement
        pass

    def run_drift_analysis(self) -> None:
        """Run drift corridor analysis as a background task."""
        try:
            # Remove old corridor layers first to avoid confusion
            for layer in self.drift_corridor_layers:
                try:
                    project = QgsProject.instance()
                    if project is not None and layer is not None:
                        project.removeMapLayer(layer.id())
                except Exception:
                    pass
            self.drift_corridor_layers.clear()

            # Get thresholds from UI
            depth_threshold = float(self.main_widget.leDepthThreshold.text() or 10)
            height_threshold = float(self.main_widget.leHeightThreshold.text() or 10)

            self.main_widget.label_drift_status.setText("Collecting data...")
            self.main_widget.pbRunDriftAnalysis.setEnabled(False)

            # Create generator and pre-collect data in main thread
            # (Qt widgets can only be accessed from the main thread)
            generator = DriftCorridorGenerator(self)
            generator.precollect_data(depth_threshold, height_threshold)

            self.main_widget.label_drift_status.setText("Starting corridor generation...")

            # Create task
            task = DriftCorridorTask(
                "Generating Drift Corridors",
                generator,
                depth_threshold,
                height_threshold
            )

            # Connect signals for progress and completion
            task.progress_updated.connect(self._on_drift_progress)
            task.corridors_generated.connect(self._on_drift_complete)
            task.generation_failed.connect(self._on_drift_failed)

            # Store task reference to prevent garbage collection
            self._drift_task = task

            # Add task to QGIS task manager (runs in background)
            task_manager = QgsApplication.taskManager()
            if task_manager is not None:
                task_manager.addTask(task)
            else:
                # Fallback: run synchronously if task manager unavailable
                self.main_widget.label_drift_status.setText("Running (no task manager)...")
                corridors = generator.generate_corridors(depth_threshold, height_threshold)
                self._on_drift_complete(corridors)

        except Exception as e:
            self.main_widget.label_drift_status.setText(f"Error: {str(e)}")
            self.main_widget.pbRunDriftAnalysis.setEnabled(True)
            QgsMessageLog.logMessage(f"Drift analysis failed: {e}", "OMRAT", Qgis.Critical)

    def _on_drift_progress(self, completed: int, total: int, message: str) -> None:
        """Handle progress updates from drift corridor task."""
        if total > 0:
            pct = int((completed / total) * 100)
            self.main_widget.label_drift_status.setText(f"{message} ({pct}%)")
        else:
            self.main_widget.label_drift_status.setText(message)

    def _on_drift_complete(self, corridors: list) -> None:
        """Handle successful completion of drift corridor generation."""
        self.main_widget.pbRunDriftAnalysis.setEnabled(True)

        if not corridors:
            self.main_widget.label_drift_status.setText("No corridors generated. Add routes first.")
            return

        # Create QGIS layers for corridors (runs in main thread)
        self._create_corridor_layers(corridors)

        num_legs = len(set(c['leg_index'] for c in corridors))
        self.main_widget.label_drift_status.setText(
            f"Generated {len(corridors)} corridors for {num_legs} legs"
        )

    def _on_drift_failed(self, error_msg: str) -> None:
        """Handle drift corridor generation failure."""
        self.main_widget.pbRunDriftAnalysis.setEnabled(True)
        self.main_widget.label_drift_status.setText(f"Error: {error_msg}")
        QgsMessageLog.logMessage(f"Drift analysis failed: {error_msg}", "OMRAT", Qgis.Critical)

    def _create_corridor_layers(self, corridors: list) -> None:
        """Create QGIS vector layers from corridor polygons - one layer per leg."""
        # Group corridors by leg_index
        legs = {}
        for corridor in corridors:
            leg_idx = corridor['leg_index']
            if leg_idx not in legs:
                legs[leg_idx] = []
            legs[leg_idx].append(corridor)

        # Create a layer for each leg
        for leg_idx, leg_corridors in legs.items():
            layer_name = f"Drift Corridors - Leg {leg_idx + 1}"

            # Create memory layer
            vl = QgsVectorLayer("Polygon?crs=EPSG:4326", layer_name, "memory")

            # Add fields
            fields = QgsFields()
            fields.append(QgsField("direction", QMetaType.Type.QString))
            fields.append(QgsField("angle", QMetaType.Type.Int))
            fields.append(QgsField("leg_index", QMetaType.Type.Int))

            provider = vl.dataProvider()
            if provider is None:
                continue
            provider.addAttributes(fields.toList())
            vl.updateFields()

            # Add features for all 8 directions
            for corridor in leg_corridors:
                poly = corridor['polygon']
                dir_name = corridor['direction']
                if poly.is_empty:
                    continue

                # Handle MultiPolygon
                if poly.geom_type == 'MultiPolygon':
                    for geom in poly.geoms:
                        feat = QgsFeature(fields)
                        feat.setGeometry(QgsGeometry.fromWkt(geom.wkt))
                        feat.setAttributes([dir_name, corridor['angle'], corridor['leg_index']])
                        provider.addFeature(feat)
                else:
                    feat = QgsFeature(fields)
                    feat.setGeometry(QgsGeometry.fromWkt(poly.wkt))
                    feat.setAttributes([dir_name, corridor['angle'], corridor['leg_index']])
                    provider.addFeature(feat)

            vl.updateExtents()

            # Style the layer with categorized renderer by direction
            self._style_corridor_layer_categorized(vl)

            # Track layer for cleanup
            self.drift_corridor_layers.append(vl)

            # Add to project
            project = QgsProject.instance()
            if project is not None:
                project.addMapLayer(vl)

    def _style_corridor_layer_categorized(self, layer: QgsVectorLayer) -> None:
        """Apply categorized styling to corridor layer based on direction field."""
        from qgis.core import QgsFillSymbol, QgsRendererCategory, QgsCategorizedSymbolRenderer

        # Color mapping for directions
        colors = {
            'N': '#e41a1c',   # Red
            'NE': '#377eb8',  # Blue
            'E': '#4daf4a',   # Green
            'SE': '#984ea3',  # Purple
            'S': '#ff7f00',   # Orange
            'SW': '#ffff33',  # Yellow
            'W': '#a65628',   # Brown
            'NW': '#f781bf',  # Pink
        }

        categories = []
        for direction, color in colors.items():
            symbol = QgsFillSymbol.createSimple({
                'color': color,
                'outline_color': color,
                'outline_width': '0.5',
            })
            if symbol:
                symbol.setOpacity(0.3)
            category = QgsRendererCategory(direction, symbol, direction)
            categories.append(category)

        renderer = QgsCategorizedSymbolRenderer('direction', categories)
        layer.setRenderer(renderer)
        layer.triggerRepaint()

    #--------------------------------------------------------------------------

    def run(self):
        """Run method that loads and starts the plugin"""
        if not self.pluginIsActive:
            self.pluginIsActive = True
            self.main_widget.pbAddRoute.clicked.connect(self.qgis_geoms.add_new_route)
            self.main_widget.pbStopRoute.clicked.connect(self.stop_route)
            self.main_widget.pbRemoveRoute.clicked.connect(self.remove_route)
            self.main_widget.pbLoadRoute.clicked.connect(self.load_route)
            self.main_widget.pbUpdateAIS.clicked.connect(self.update_ais)
            self.main_widget.pbGetGebcoDephts.clicked.connect(self.object.obtain_gebco_data)
            self.main_widget.PBUpdateDepthIntervals.clicked.connect(self.object.update_depth_intervals)
            
            menubar = QMenuBar(self.main_widget)
            menubar.setMinimumSize(320,20)

            fileMenu = menubar.addMenu('File')
            SettingMenu = menubar.addMenu('Settings')
            analyse_sen_Menu = menubar.addMenu('Analyse sensitivity')
            HelpMenu = menubar.addMenu('Help')
            
            # mnuSub1 = viewMenu.addMenu('Sub-menu')
            fileMenu.addAction("Save", self.save_work)
            fileMenu.addAction("Load", self.load_work)
            fileMenu.addSeparator()
            fileMenu.addAction("Export to IWRAP XML", self.export_to_iwrap)
            fileMenu.addAction("Import from IWRAP XML", self.import_from_iwrap)
            SettingMenu.addAction("Drift settings", self.open_drift)
            SettingMenu.addAction("Ship Categories", self.open_ship_categories)
            SettingMenu.addAction("Causation Factors", self.open_causation_factors)
            SettingMenu.addAction("AIS connection settings", self.ais_settings)
            #self.main_widget.actionSave_project.clicked.connect(self.save_work)
            #self.main_widget.actionOpen_project.clicked.connect(self.load_work)
            self.main_widget.pbAddSimpleDepth.clicked.connect(self.object.add_simple_depth)
            self.main_widget.pbLoadDepth.clicked.connect(self.object.load_depths)
            self.main_widget.pbRemoveDepth.clicked.connect(self.object.remove_depth)
            self.main_widget.pbAddSimpleObject.clicked.connect(self.object.add_simple_object)
            self.main_widget.pbLoadObject.clicked.connect(self.object.load_objects)
            self.main_widget.pbRemoveObject.clicked.connect(self.object.remove_object)
            self.main_widget.pbViewDriftingAllision.clicked.connect(self.show_drift_allision)
            self.main_widget.pbViewPoweredAllision.clicked.connect(self.show_powered_allision)
            self.main_widget.pbViewPoweredGrounding.clicked.connect(self.show_powered_grounding)
            # TODO(view-buttons): replace the stubs below with real
            # visualisation pipelines.  See omrat.py::show_drift_grounding
            # et al. for the placeholders these connect to.
            self.main_widget.pbViewDriftingGrounding.clicked.connect(self.show_drift_grounding)
            self.main_widget.pbViewOvertakingCollision.clicked.connect(self.show_overtaking_collision)
            self.main_widget.pbViewHeadOnCollision.clicked.connect(self.show_head_on_collision)
            self.main_widget.pbViewCrossingCollision.clicked.connect(self.show_crossing_collision)
            self.main_widget.pbViewMergingCollision.clicked.connect(self.show_merging_collision)
            self.main_widget.pbRunModel.clicked.connect(self.run_calculation)
            # ``LEReportPath`` is the per-run output folder (one .gpkg
            # is written per Run Model into this directory).  The
            # picker button opens a folder-only dialog; live edits to
            # the line-edit re-evaluate the gate on the Run Model
            # button.
            try:
                if hasattr(self.main_widget, 'pbReportPath'):
                    # Replace the legacy "save report to file" wiring.
                    try:
                        self.main_widget.pbReportPath.clicked.disconnect()
                    except Exception:
                        pass
                    self.main_widget.pbReportPath.clicked.connect(
                        self.choose_output_folder,
                    )
                if hasattr(self.main_widget, 'LEReportPath'):
                    self.main_widget.LEReportPath.textChanged.connect(
                        lambda *_: self._update_run_model_enabled(),
                    )
                self._restore_output_dir()
                self._update_run_model_enabled()
            except Exception as exc:
                QgsMessageLog.logMessage(
                    f'Could not wire output-folder picker: {exc}',
                    'OMRAT', Qgis.Warning,
                )
            # Connect drift analysis button
            try:
                if hasattr(self.main_widget, 'pbRunDriftAnalysis'):
                    self.main_widget.pbRunDriftAnalysis.clicked.connect(self.run_drift_analysis)
            except Exception:
                pass
            # Wire up the "Previous runs" history table.
            try:
                self._setup_previous_runs_table()
                self.refresh_previous_runs_table()
            except Exception as e:
                QgsMessageLog.logMessage(
                    f'Could not initialise previous-runs table: {e}',
                    'OMRAT', Qgis.Warning,
                )
            # Add a "Manage previous runs" item to the File menu so
            # users can find load / compare / delete actions even if
            # they ignore the table on the Results tab.
            try:
                fileMenu.addSeparator()
                fileMenu.addAction(
                    'Manage previous runs...', self.open_previous_runs_dialog,
                )
            except Exception:
                pass

            self.reset_route_table()
            self.run_traffic_module()
            self.main_widget.show()
