"""
/***************************************************************************
Name       :  geowrsi_climatological_analysis_controller.py
Description:  Geowrsi climatological analysis controller for FEWSTools plugin
copyright  :  (C) 2022-2023 by FEWS
email      :  jhowton@contractor.usgs.gov
Created    :  10/17/2022 - Jacob Howton
Modified   :
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   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 glob
import os

from PyQt5 import QtCore, QtGui
from PyQt5.QtWidgets import QMessageBox, QDialog, QFileDialog
from qgis.core import QgsMessageLog, Qgis
from qgis.utils import iface

from fews_tools import fews_tools_config as config
from fews_tools.workers.geowrsi_climatological_analysis_worker \
    import GeoWRSIClimatologicalAnalysisWorker as wrkr
from fews_tools.utilities import geoclim_utilities as util
from fews_tools.utilities import geoclim_qgs_utilities as qgs_util
from fews_tools.utilities import logging_utilities as log_util
from fews_tools.models.region_model import RegionModel
from fews_tools.models.datasets_model import DatasetsModel
from fews_tools.models.workspace_setup_model import WorkspaceSetupModel

from fews_tools.forms.Ui_GeoWRSIClimAnalysis import Ui_GeoWRSIClimAnalysis


class GeoWRSIClimatologicalAnalysisController(QDialog):
    '''
    class for GeoWRSI Climatological analysis controller
    '''

    def __init__(self):
        QDialog.__init__(self)

        self.is_processing = False  # Flag to indicate if running a process

        # set up UI and models
        self.wrksp_setup = WorkspaceSetupModel()
        self.region_info = RegionModel()
        self.ds_info = DatasetsModel()
        self.ui = Ui_GeoWRSIClimAnalysis()
        self.ui.setupUi(self)

        self.thread = None
        self.worker = None
        self.closed = False  # A flag to indicate if close_event is triggered

        self.files_available = []
        self.years_available = []

        self.years_to_analyze = []
        self.files_to_analyze = []
        self.all_years_selected = False

        self.parameter = ""
        self.smart_trend = False
        self.min_r2 = 0
        self.percentile = 0

        self.output_path = ""
        self.analysis_method = ""

        self.parameter_fname = ""

        self.set_dic = {}
        self.reg_dic = {}

        self.close_event = None

        # handle UI
        self.ui.progressBar.setValue(0)

        self.ui.browseWRSIButton.clicked.connect(self.browse_wrsi)
        self.ui.browseOutputButton.clicked.connect(self.browse_output)
        self.ui.regionComboBox.currentIndexChanged.connect(self.region_changed)
        self.ui.parameterComboBox.currentIndexChanged.connect(
            self.check_input_files)
        self.ui.extensionComboBox.currentIndexChanged.connect(
            self.check_input_files)
        self.ui.analysisMethodComboBox.currentIndexChanged.connect(
            self.method_changed)
        self.ui.smartTrendsCheckBox.stateChanged.connect(
            self.smart_trend_check)
        self.ui.selectAllYearsButton.clicked.connect(
            self.toggle_select_all_years)
        self.ui.outputPrefixLineEdit.textChanged.connect(
            self.check_input_files)
        self.fill_settings_dict()
        util.fill_misc_widget_no_sort(
            self.ui.regionComboBox, self.region_info.get_all_region_names())
        util.fill_misc_widget_no_sort(
            self.ui.analysisMethodComboBox, wrkr.CLIM_ANAL_METHODS)
        util.fill_misc_widget_no_sort(self.ui.parameterComboBox, [
                                      "WRSI", "Start of Season"])
        util.fill_misc_widget_no_sort(self.ui.extensionComboBox, [
                                      config.BIL_SUFFIX, config.TIF_SUFFIX])

        self.ui.WRSIFolderLineEdit.setText(self.wrksp_setup.output_path)
        self.ui.OutputFolderLineEdit.setText(self.wrksp_setup.output_path)
        self.ui.regionComboBox.setCurrentText(self.set_dic['analysis_region'])

        self.ui.analyzeButton.clicked.connect(
            self.start_climatological_analysis)

    def browse_wrsi(self):
        '''
        Function to browse to WRSI output directory and set line edit control.
        '''
        dir_name = \
            QFileDialog.getExistingDirectory(
                None,
                "Select WRSI Output Data Directory",
                self.wrksp_setup.get_output_path())
        if dir_name:
            self.ui.WRSIFolderLineEdit.setText(
                self.wrksp_setup.fix_os_sep_in_path(dir_name))

        self.region_changed()

    def browse_output(self):
        '''
        Function to browse to output directory and set line edit control.
        '''
        dir_name = \
            QFileDialog.getExistingDirectory(
                None,
                "Select Output Data Directory",
                self.wrksp_setup.get_output_path())
        if dir_name:
            self.ui.OutputFolderLineEdit.setText(
                self.wrksp_setup.fix_os_sep_in_path(dir_name))

    def fill_settings_dict(self):
        '''
        This function will open the geowrsi_settings.txt file in order to fill
        the set_dic with needed data from the geowrsi_settings_controller
        returns(bool) - err - if there was an error or not loading set_dic
        '''
        err = True
        if os.path.exists(os.path.join(self.wrksp_setup.get_workspace(),
                                       config.PROGRAM_SETTINGS,
                                       config.GEOWRSI_SETTINGS_FILE)):
            lines = []
            with open(os.path.join(self.wrksp_setup.get_workspace(),
                                   config.PROGRAM_SETTINGS,
                                   config.GEOWRSI_SETTINGS_FILE), "r") as f_obj:
                lines = f_obj.readlines()
            for line in lines:
                x_obj = line.split(' ', 1)
                self.set_dic[x_obj[0]] = x_obj[1].strip()
            err = False
        else:
            QMessageBox.information(
                self, 'geowrsi_settings.txt File: Path Not Found',
                'The path to the settings file cannot be found!',
                QMessageBox.Ok)
        return err

    def check_input_files(self):
        '''
        Check what files are available
        '''

        self.ui.filesAvailableListWidget.clear()
        self.ui.yearsListWidget.clear()

        self.parameter = self.ui.parameterComboBox.currentText()

        if self.parameter == "WRSI":
            self.parameter_fname = "wrsi_eos"

        else:
            self.parameter_fname = "sos"

        self.ds_info.query_named_dataset(
            self.ds_info.query_default_dataset_name())
        ds_dic = self.ds_info.get_ds_dictionary()

        prefix = self.ui.outputPrefixLineEdit.text()
        suffix = self.ui.extensionComboBox.currentText()

        if os.path.exists(self.ui.WRSIFolderLineEdit.text()):
            region = self.ui.regionComboBox.currentText().replace(
                ' ', '_').replace('-', '_')

            self.files_available = glob.glob(os.path.join(
                self.ui.WRSIFolderLineEdit.text(),
                prefix + region + '_w*' + 
                self.parameter_fname + suffix))

            if self.files_available != []:
                cross_year = (
                    self.region_info.reg_dic['InitialPeriod'] > 
                    self.region_info.reg_dic['FinalPeriod'])
                
                year_characters = 4
                if cross_year:
                    year_characters = 9

                suffix_characters = len(self.parameter_fname)
                    
                self.years_available = [os.path.splitext(
                    f)[0][-(year_characters+suffix_characters):
                    -suffix_characters] for f in self.files_available]
                
                for file in self.files_available:
                    self.ui.filesAvailableListWidget.addItem(
                        os.path.basename(file))

                for year in self.years_available:
                    self.ui.yearsListWidget.addItem(year)

    def toggle_select_all_years(self):
        '''
        Select all years and swap the label to deselect
        '''
        if not self.all_years_selected:
            self.ui.selectAllYearsButton.setText("Clear All Years")
            self.ui.yearsListWidget.selectAll()
        else:
            self.ui.selectAllYearsButton.setText("Select All Years")
            self.ui.yearsListWidget.clearSelection()

        self.all_years_selected = not self.all_years_selected

    def region_changed(self):
        '''
        Function for when the region selection is changed
        '''
        self.all_years_selected = False
        self.ui.selectAllYearsButton.setText("Select All Years")
        self.ui.yearsListWidget.clearSelection()

        self.region_info.query_named_region(
            self.ui.regionComboBox.currentText())
        self.reg_dic = self.region_info.get_region_dictionary()

        self.check_input_files()

    def method_changed(self):
        '''
        If the method is changed then check for secondary options
        '''
        self.analysis_method = self.ui.analysisMethodComboBox.currentText()

        if self.analysis_method == "Trend":
            self.ui.smartTrendsCheckBox.setVisible(True)
        else:
            self.ui.smartTrendsCheckBox.setVisible(False)
            self.ui.lblMinR2.setVisible(False)
            self.ui.minR2DoubleSpinBox.setVisible(False)

        if self.analysis_method == "Percentiles":
            self.ui.percentileSpinBox.setVisible(True)
            self.ui.lblPercentile.setVisible(True)

        else:
            self.ui.lblPercentile.setVisible(False)
            self.ui.percentileSpinBox.setVisible(False)

    def smart_trend_check(self):
        '''
        Function to check toggle minR2 spin box
        '''
        self.smart_trend = self.ui.smartTrendsCheckBox.isChecked()
        self.ui.lblMinR2.setVisible(self.smart_trend)
        self.ui.minR2DoubleSpinBox.setVisible(self.smart_trend)

    def get_files_to_analize(self):
        '''
        Function that sets the years to analize from the widget and sets the
        files to analize by selecting the files with the specified years.
        '''
        self.years_to_analyze = [f.text()
                                 for f in self.ui.yearsListWidget.selectedItems()]

        self.files_to_analyze = [file for file in self.files_available if any(
            year in file for year in self.years_to_analyze)]

    def start_climatological_analysis(self):
        '''
        Function to run the analysis.
        '''

        self.get_files_to_analize()

        if len(self.years_to_analyze) < 2:
            QMessageBox.critical(
                self, 'Error!', "Multiple years must be selected.",
                QMessageBox.Ok)
            return

        if self.ui.OutputFilenameLineEdit.text() == "":
            QMessageBox.critical(
                self, 'Error!', "An output filename must be provided.",
                QMessageBox.Ok)
            return

        self.min_r2 = self.ui.minR2DoubleSpinBox.value()
        self.percentile = self.ui.percentileSpinBox.value()
        self.output_path = os.path.join(
            self.ui.OutputFolderLineEdit.text(),
            self.ui.OutputFilenameLineEdit.text())
        self.worker = wrkr(
            self.analysis_method,
            self.parameter,
            self.smart_trend,
            self.min_r2,
            self.percentile,
            self.years_to_analyze,
            self.files_to_analyze,
            self.output_path,
            self.reg_dic)

        self.check_open_files()

        # lock UI
        self.ui.CancelButton.clicked.connect(
            self.worker.kill)
        self.ui.CancelButton.setEnabled(True)
        self.ui.closeButton.setEnabled(False)
        self.ui.browseWRSIButton.setEnabled(False)
        self.ui.filesAvailableListWidget.setEnabled(False)
        self.ui.yearsListWidget.setEnabled(False)
        self.ui.selectAllYearsButton.setEnabled(False)
        self.ui.browseOutputButton.setEnabled(False)
        self.ui.regionComboBox.setEnabled(False)
        self.ui.outputPrefixLineEdit.setEnabled(False)
        self.ui.parameterComboBox.setEnabled(False)
        self.ui.analysisMethodComboBox.setEnabled(False)
        self.ui.OutputFilenameLineEdit.setEnabled(False)
        self.ui.analyzeButton.setEnabled(False)

        # start the process in a new thread
        iface.messageBar().clearWidgets()
        self.thread = QtCore.QThread(self)
        self.worker.moveToThread(self.thread)
        self.thread.started.connect(self.worker.run)
        self.thread.start()
        self.worker.finished.connect(self.finish_clim_analysis)
        self.worker.error.connect(self.clim_analysis_error)
        self.worker.progress.connect(self.ui.progressBar.setValue)
        self.is_processing = True

    def finish_clim_analysis(self, return_value):
        '''
        Clean up the climatological analysis worker object and thread.
        return_value - Returned code and string from
                                climatological analysis worker
        '''

        self.is_processing = False
        self.thread.quit()
        self.thread.wait()
        if self.closed:
            self.thread_worker_cleanup()
            self.close_event.accept()
            QDialog.closeEvent(self, self.close_event)
            return
        if self.worker.killed:
            self.ui.progressBar.setValue(0)
            msg = ("Climatological analysis aborted by user")
            iface.messageBar().pushCritical("Climatological analysis", msg)
            QgsMessageLog.logMessage(msg, level=Qgis.Critical)
            iface.mainWindow().statusBar().showMessage(msg, 3000)
        elif return_value is not None:
            color_path = None

            # Average and Median
            if self.analysis_method in wrkr.CLIM_ANAL_METHODS[0:2]:
                if self.parameter == "WRSI":
                    color_path = os.path.join(
                        self.wrksp_setup.get_colors_path(),
                        config.WRSI_COLOR_FILE)

                else:
                    color_path = self.reg_dic['SOSColor']

            # Standard Deviation
            elif self.analysis_method == wrkr.CLIM_ANAL_METHODS[2]:
                color_path = os.path.join(
                    self.wrksp_setup.get_colors_path(),
                    config.RF_600_RASTER_COLOR_FILE)

            elif self.analysis_method == wrkr.CLIM_ANAL_METHODS[3]:  # Range
                color_path = os.path.join(
                    self.wrksp_setup.get_colors_path(),
                    config.COUNT_COLOR_FILE)

            elif self.analysis_method == wrkr.CLIM_ANAL_METHODS[4]:  # Count
                color_path = os.path.join(
                    self.wrksp_setup.get_colors_path(),
                    config.COUNT_COLOR_FILE)

            # Coefficient of Variation
            elif self.analysis_method == wrkr.CLIM_ANAL_METHODS[5]:
                color_path = os.path.join(
                    self.wrksp_setup.get_colors_path(), config.RF_CV_COLOR_FILE)

            elif self.analysis_method == wrkr.CLIM_ANAL_METHODS[6]:  # Trend
                color_path = os.path.join(
                    self.wrksp_setup.get_colors_path(),
                    config.RF_TREND_COLOR_FILE)

            # Percentiles
            elif self.analysis_method == wrkr.CLIM_ANAL_METHODS[7]:
                color_path = os.path.join(
                    self.wrksp_setup.get_colors_path(),
                    config.RF_600_RASTER_COLOR_FILE)

            qgs_util.display_raster_layer(
                self.output_path, color_path)

        # Unlock UI
        self.ui.closeButton.setEnabled(True)
        self.ui.browseWRSIButton.setEnabled(True)
        self.ui.filesAvailableListWidget.setEnabled(True)
        self.ui.yearsListWidget.setEnabled(True)
        self.ui.selectAllYearsButton.setEnabled(True)
        self.ui.browseOutputButton.setEnabled(True)
        self.ui.regionComboBox.setEnabled(True)
        self.ui.outputPrefixLineEdit.setEnabled(True)
        self.ui.parameterComboBox.setEnabled(True)
        self.ui.analysisMethodComboBox.setEnabled(True)
        self.ui.OutputFilenameLineEdit.setEnabled(True)
        self.ui.analyzeButton.setEnabled(True)
        self.thread_worker_cleanup()

        log_util.add_log("Climatological Analysis",
                         log_util.log_string("Region",
                                             [self.set_dic['analysis_region']]) +
                         log_util.log_string("Years to Analyze",
                                             self.years_to_analyze) +
                         log_util.log_string("Files to Analye",
                                             self.files_to_analyze) +
                         log_util.log_string("Analysis Parameter",
                                             [self.parameter]) +
                         log_util.log_string("Analysis Method",
                                             [self.analysis_method]) +
                         log_util.log_string("Smart Trend",
                                             [self.smart_trend]) +
                         log_util.log_string("Min r2",
                                             [self.min_r2]) +
                         log_util.log_string("Percentile",
                                             [self.percentile]),
                         [self.output_path])

        all_done = util.notify_process_completed('Climatological Analysis')
        if all_done:
            self.ui.closeButton.click()

        util.remove_temp_files(self.output_path)

    def check_open_files(self):
        '''
        Check if the output file is currently open in QGIS and close if so.
        '''
        self.open_file_info = qgs_util.get_open_files_info()

        for entry in self.open_file_info:
            if os.path.splitext(os.path.basename(self.output_path))[0] == entry[1]:
                qgs_util.notify_loaded_file(self, entry[1], entry[0])

    def clim_analysis_error(self, err, ex_str):
        '''
        Function to log any errors from the object and thread.
        Args:  err - not used
               ex_str string: Exception info
        '''
        self.is_processing = False
        QgsMessageLog.logMessage(ex_str, level=Qgis.Critical)
        QMessageBox.critical(self, 'Error!', str(err) + ex_str, QMessageBox.Ok)
        iface.mainWindow().statusBar().showMessage(ex_str)

    def closeEvent(self, a_0: QtGui.QCloseEvent) -> None:
        '''
        Close event used to handle a case where user closes the form
        while the process is running.
        '''
        self.closed = True
        if self.worker:
            self.worker.kill()
        self.close_event = a_0
        if self.is_processing:
            self.close_event.ignore()
        else:
            self.close_event.accept()

    def thread_worker_cleanup(self):
        """
        Cleanup thread if it exists.
        """
        if self.thread is not None:
            self.thread.deleteLater()
