qgis_pypoprf_dialog.py

# -*- coding: utf-8 -*-
"""
/***************************************************************************
 PyPopRFDialog
                                 A QGIS plugin
 A plugin for population prediction and dasymetric mapping using machine learning
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                             -------------------
        begin                : 2025-01-07
        git sha              : $Format:%H$
        copyright            : (C) 2025 by WorldPop SDI Team
        email                : b.nosatiuk@soton.ac.uk
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   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     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""

import os
from typing import Optional, Any

from qgis.PyQt import uic, QtWidgets
from qgis.PyQt.QtCore import QThread

from .q_models.config_manager import ConfigManager
from .q_models.console_handler import ConsoleHandler
from .q_models.covariate_table import CovariateTableHandler
from .q_models.file_handlers import FileHandler
from .q_models.process_executor import ProcessExecutor
from .q_models.settings_handler import SettingsHandler

FORM_CLASS, _ = uic.loadUiType(
    os.path.join(os.path.dirname(__file__), "qgis_pypoprf_dialog_base.ui")
)


class PyPopRFDialog(QtWidgets.QDialog, FORM_CLASS):
    """Main dialog for the PyPopRF plugin.

    This class handles the user interface and coordinates communication between
    different components of the plugin. It manages user input, file handling,
    settings configuration, and analysis execution.

    Attributes:
        iface: QGIS interface instance
        logger: Logger instance for output messages
        config_manager: Handler for configuration management
        covariate_handler: Handler for covariate management
        file_handler: Handler for file operations
        settings_handler: Handler for settings management
        process_executor: Handler for analysis execution
    """

    def __init__(
        self, parent: Optional[QtWidgets.QWidget] = None, iface: Optional[Any] = None
    ) -> None:
        """Initialize the dialog.

        Args:
            parent: Parent widget
            iface: QGIS interface instance
        """
        super(PyPopRFDialog, self).__init__(parent)
        self.iface = iface
        self.setupUi(self)

        # Initialize handlers
        self.console_handler = ConsoleHandler(self.scrollAreaWidgetContents)
        self.logger = self.console_handler.logger
        self.config_manager = ConfigManager(self.logger)
        self.covariate_handler = CovariateTableHandler(
            self.covariatesTable, self.config_manager, self.logger
        )
        self.file_handler = FileHandler("", self.logger)
        self.settings_handler = SettingsHandler(self.config_manager, self.logger)
        self.process_executor = ProcessExecutor(self, self.logger, iface)

        # Connect signals
        self._connect_signals()

        # Setup file widgets
        self._setup_file_widgets()

        # Disable widgets until project is initialized
        self._set_initial_state()

    def _connect_signals(self):
        """Connect all UI signals to their respective handlers."""
        # Button signals
        self.initProjectButton.clicked.connect(self.init_project)
        self.openProjectFolder.clicked.connect(self.open_project_folder)
        self.addCovariateButton.clicked.connect(self.add_covariate)

        # File widget signals
        self.workingDirEdit.fileChanged.connect(self.on_working_dir_changed)
        self.mastergridFileWidget.fileChanged.connect(
            lambda x: self._handle_file_change("mastergrid", x)
        )
        self.maskFileWidget.fileChanged.connect(lambda x: self._handle_file_change("mask", x))
        self.constrainFileWidget.fileChanged.connect(
            lambda x: self._handle_file_change("constrain", x)
        )
        self.populationCensusFileWidget.fileChanged.connect(
            lambda x: self._handle_file_change("census_data", x)
        )
        self.agesexCensusFileWidget.fileChanged.connect(
            lambda x: self._handle_file_change("agesex_data", x)
        )

        # Settings tab signals
        self.enableParallelCheckBox.stateChanged.connect(
            lambda: self.cpuCoresComboBox.setEnabled(self.enableParallelCheckBox.isChecked())
        )
        self.enableBlockProcessingCheckBox.stateChanged.connect(
            lambda: self.blockSizeComboBox.setEnabled(
                self.enableBlockProcessingCheckBox.isChecked()
            )
        )
        self.cpuCoresComboBox.currentTextChanged.connect(self._handle_cpu_cores_changed)
        self.blockSizeComboBox.currentTextChanged.connect(self._handle_block_size_changed)

        # Logging signals
        self.comboBox.currentTextChanged.connect(self._update_logging_settings)

        self.settings_handler.connect_census_fields_signals(self)
        self.settings_handler.connect_log_filename_signals(self)

        self.logScaleCheckBox.stateChanged.connect(lambda: self.settings_handler.save_settings(self))
        self.selectionThresholdLimit.valueChanged.connect(lambda: self.settings_handler.save_settings(self))


        # Analysis signals
        self.mainStartButton.setStyleSheet(
            "QPushButton { background-color: #878c87; color: black; font-size: 10pt; }"
        )
        self.mainStartButton.clicked.connect(self._handle_start_button)

    def _setup_file_widgets(self):
        """Configure file widgets with appropriate filters and titles."""

        # Mastergrid
        self.mastergridFileWidget.setDialogTitle("Select Mastergrid File (Required)")
        self.mastergridFileWidget.setFilter("GeoTIFF files (*.tif *.tiff)")

        # Mask
        self.maskFileWidget.setDialogTitle("Select Mask File (Optional)")
        self.maskFileWidget.setFilter("GeoTIFF files (*.tif *.tiff)")

        # Constrain
        self.constrainFileWidget.setDialogTitle("Select Constrain File (Optional)")
        self.constrainFileWidget.setFilter("GeoTIFF files (*.tif *.tiff)")

        # Population Census
        self.populationCensusFileWidget.setDialogTitle("Select Population Census CSV File")
        self.populationCensusFileWidget.setFilter("CSV files (*.csv)")

        # Age-Sex Census
        self.agesexCensusFileWidget.setDialogTitle("Select Age-Sex Census CSV File (Optional)")
        self.agesexCensusFileWidget.setFilter("CSV files (*.csv)")

    def _setup_cpu_cores_combo(self):
        """Configure CPU cores combo box based on available system resources.

        Sets up the combo box with values from 2 to max available cores in steps of 2.
        Default value is set to half of available cores + 2.
        """
        self.cpuCoresComboBox.clear()

        max_cores = QThread.idealThreadCount()
        self.logger.info(f"System has {max_cores} logical processors available")

        core_counts = list(range(2, max_cores + 1, 2))

        # Add items to combo box
        for count in core_counts:
            self.cpuCoresComboBox.addItem(str(count))

        default_cores = min(max_cores, (max_cores // 2) + 2)
        if default_cores % 2 != 0:
            default_cores -= 1

        default_index = core_counts.index(default_cores) if default_cores in core_counts else 0
        self.cpuCoresComboBox.setCurrentIndex(default_index)
        self.config_manager.update_config("max_workers", default_cores)

    def _set_initial_state(self):
        """Set initial state of UI widgets before project initialization."""

        self.initProjectButton.setEnabled(False)
        self.mainStartButton.setEnabled(False)
        self.openProjectFolder.setEnabled(False)

        # Disable input and settings tabs
        self.set_input_widgets_enabled(False)
        self.set_settings_widgets_enabled(False)

    # Set input main widgets enabled
    def _set_main_widgets_enabled(self, enabled: bool):
        """Enable/disable main widgets"""
        self.initProjectButton.setEnabled(enabled)
        self.workingDirEdit.setEnabled(enabled)
        self.openProjectFolder.setEnabled(enabled)

    def set_input_widgets_enabled(self, enabled: bool):
        """Enable/disable input widgets"""
        self.mastergridFileWidget.setEnabled(enabled)
        self.maskFileWidget.setEnabled(enabled)
        self.constrainFileWidget.setEnabled(enabled)
        self.populationCensusFileWidget.setEnabled(enabled)
        self.agesexCensusFileWidget.setEnabled(enabled)
        self.addCovariateButton.setEnabled(enabled)
        self.covariatesTable.setEnabled(enabled)

    def set_settings_widgets_enabled(self, enabled: bool):
        """Enable/disable settings widgets"""

        # Logging settings
        self.logsColumnEdit.setEnabled(enabled)
        self.comboBox.setEnabled(enabled)

        # Process settings
        self.enableParallelCheckBox.setEnabled(enabled)
        self.cpuCoresComboBox.setEnabled(self.enableParallelCheckBox.isChecked() and enabled)
        self.enableBlockProcessingCheckBox.setEnabled(enabled)
        self.blockSizeComboBox.setEnabled(
            self.enableBlockProcessingCheckBox.isChecked() and enabled
        )

        # Census settings
        self.popColumnEdit.setEnabled(enabled)
        self.idColumnEdit.setEnabled(enabled)

        # Additional settings
        self.addToQgisCheckBox.setEnabled(enabled)
        self.logScaleCheckBox.setEnabled(enabled)
        self.selectionThresholdLimit.setEnabled(enabled)

    def _handle_file_change(self, file_type: str, path: str):
        """Handle file selection changes for input files.

        Args:
            file_type: Type of file being changed ('mastergrid', 'mask', etc.)
            path: New file path
        """

        if path:
            filename = self.file_handler.copy_to_data_dir(path, file_type)
            if filename:
                self.config_manager.update_config(file_type, filename)
        else:
            self.config_manager.clear_config_value(file_type)

        has_mastergrid = bool(self.mastergridFileWidget.filePath())
        has_census = bool(self.populationCensusFileWidget.filePath())

        # Enable start button if all files are loaded
        self.mainStartButton.setEnabled(all([has_mastergrid, has_census]))
        if all([has_mastergrid, has_census]):
            self.mainStartButton.setStyleSheet(
                "QPushButton { background-color: #4CAF50; color: black; font-size: 10pt; }"
            )

    def on_working_dir_changed(self, path: str):
        """Handle working directory change"""
        self.initProjectButton.setEnabled(bool(path))

    def init_project(self):
        """Initialize new pypopRF project.

        Creates project directory structure, initializes configuration,
        and enables UI elements for further setup.

        Raises:
            Exception: If project initialization fails
        """

        self.console_handler.clear()
        self.mainProgressBar.setValue(0)

        working_dir = self.workingDirEdit.filePath()
        if not working_dir:
            self.logger.error("Please select working directory first")
            return

        try:
            # Create project structure and config
            if self.config_manager.create_initial_config(working_dir):
                self.file_handler.set_working_dir(working_dir)

                # Enable UI elements
                self.openProjectFolder.setEnabled(True)
                self.set_input_widgets_enabled(True)
                self.set_settings_widgets_enabled(True)
                self.addToQgisCheckBox.setChecked(True)

                self._setup_cpu_cores_combo()

                # Load initial settings
                self.settings_handler.load_settings(self)
                self._update_logging_settings()

                # Show next steps
                self.logger.info("Project initialized successfully!")
                self.logger.info(
                    '<span style="font-weight: bold; font-size: 11pt; color: '
                    '#050505;">Next steps: &#8628;</span>'
                )
                self.logger.info(
                    '<span style="font-weight: bold; color: #0066cc;">'
                    "1. Place input files in the data directory"
                    "</span>"
                )
                self.logger.info(
                    '<span style="font-weight: bold; color: #0066cc;">'
                    "2. Configure input data paths in the Input Data tab"
                    "</span>"
                )
                self.logger.info(
                    '<span style="font-weight: bold; color: #0066cc;">'
                    "3. Adjust processing settings in the Settings tab"
                    "</span>"
                )

        except Exception as e:
            self.logger.error(f"Error initializing project: {str(e)}")

    def add_covariate(self):
        """Add new covariate files to the project.

        Opens file dialog for selecting covariate files and adds them
        to the project data directory.
        """

        file_paths, _ = QtWidgets.QFileDialog.getOpenFileNames(
            self, "Select Covariate Files", "", "GeoTIFF files (*.tif *.tiff)"
        )
        filenames = []
        for path in file_paths:
            filename = self.file_handler.copy_to_data_dir(path, "covariate")
            if filename:
                filenames.append(filename)
        if filenames:
            self.covariate_handler.add_covariates(filenames)

    def open_project_folder(self):
        """Open project folder in system file explorer"""
        self.file_handler.open_folder(self.workingDirEdit.filePath())

    def _update_logging_settings(self):
        """Update logging configuration based on current UI state.

        Updates both config file and logger instance with new settings.
        """

        self.config_manager.update_config(
            "logging",
            {"level": self.comboBox.currentText(), "file": self.logsColumnEdit.text()},
        )

        # Then update logger
        self.console_handler.update_logging_settings(
            level=self.comboBox.currentText(),
            save_log=True,
            work_dir=self.workingDirEdit.filePath(),
            filename=self.logsColumnEdit.text(),
        )

    def _handle_cpu_cores_changed(self, i_cores):
        """Handle changes to CPU core count setting.

        Args:
            i_cores: New CPU core count value
        """

        if self.enableParallelCheckBox.isChecked():
            try:
                cores = int(i_cores) if i_cores else 0
                if cores > 0:
                    self.config_manager.update_config("max_workers", cores)
                    self.logger.debug(f"CPU cores set to {cores}")
            except ValueError:
                self.logger.warning(f"Invalid CPU cores value: {i_cores}")

    def _handle_block_size_changed(self, text):
        """Handle changes to processing block size.

        Args:
            text: New block size value in format "width, height"
        """

        if self.enableBlockProcessingCheckBox.isChecked():
            try:
                w, h = map(int, text.replace(" ", "").split(","))
                if w > 0 and h > 0:
                    self.config_manager.update_config("block_size", [w, h])
            except ValueError:
                self.logger.warning(f"Invalid block size format: {text}")

    def _handle_start_button(self):
        """Handle start/stop button click"""
        if self.mainStartButton.text() == "Start":
            self.process_executor.start_analysis()
        else:
            self.process_executor.stop_analysis()