'''
/***************************************************************************
Name          :  extract_stats_controller.py
Description   :  Extract statistics from Raster Data using Shapefile
copyright     :  (C) 2019-2023 by FEWS
email         :  minxuansun@contractor.usgs.gov
Created       :  02/20/2020 - ACHRISTIANSON
Modified      :  06/15/2020 - cholen - Adjust check for getSaveFileName result
                 07/13/2020 - MSun - add lock and unlock functions,
                                     fix NULL stats and Add button issues.
                 07/17/2020 - cholen - Adjust progress bar, disable path line
                                     edits
                 07/23/2020 - cholen - Add check and update to use nodata
                 08/28/2020 - cholen - Notify on completion, cleanup
                 12/03/2020 - cholen - Handle OSError
                 01/20/2022 - cholen - New gdal utilities
                 02/10/2022 - cholen - Handle point shapefiles, add worker
                 03/25/2022 - cholen - Add tiff support, changes for code consistency,
                                       and testability
                 07/28/2023 - jhowton - Added update CSV
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   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
import shutil
import csv

from PyQt5 import QtCore, QtGui
from PyQt5.QtWidgets import QDialog, QFileDialog, QTableWidgetItem, QMessageBox

from qgis.core import QgsVectorLayer, QgsMessageLog, Qgis, QgsFeatureRequest
from qgis.utils import iface

from fews_tools.models.workspace_setup_model import WorkspaceSetupModel
from fews_tools.controllers.time_series_list_controller import TimeSeriesListController
from fews_tools.forms.Ui_ExtractStats import Ui_ExtractStats

from fews_tools import fews_tools_config as config
from fews_tools.utilities import geoclim_qgs_utilities as qgs_util
from fews_tools.utilities import geoclim_utilities as util
from fews_tools.utilities import logging_utilities as log_util
from fews_tools.utilities.help_utilities import view_manual

from fews_tools.workers.extract_stats_worker import ExtractStatsWorker as wrkr


class ExtractStatsController(QDialog):
    '''
    Class to extract raster statistics
    '''

    RASTER_LABELS = (['ROW', 'FILENAME'])

    def __init__(self):

        QDialog.__init__(self)
        self.ui = Ui_ExtractStats()
        self.ui.setupUi(self)
        self.wrksp_setup = WorkspaceSetupModel()
        self.temp_path = self.wrksp_setup.get_temp_data_path()
        self.temp_shp = ''
        # adding list and header values to form
        util.fill_misc_widget(
            self.ui.summaryComboBox, wrkr.Z_STATS_TYPES_DIC.keys())
        self.ui.rstrTableWidget.setRowCount(0)
        self.ui.rstrTableWidget.setColumnCount(2)
        self.ui.rstrTableWidget.setHorizontalHeaderLabels(self.RASTER_LABELS)
        self.ui.rstrTableWidget.resizeRowsToContents()
        # making connections
        self.ui.vectorButton.clicked.connect(self.browse_vector_file)
        self.ui.addButton.clicked.connect(self.add_raster_file_to_table)
        self.ui.removeButton.clicked.connect(
            self.remove_raster_file_from_table)
        self.ui.saveButton.clicked.connect(self.save_new_list_file)
        self.ui.loadButton.clicked.connect(self.browse_list_file)
        self.ui.existingCSVButton.clicked.connect(self.browse_csv_file)
        self.ui.outputButton.clicked.connect(self.browse_output_file)
        self.ui.processButton.clicked.connect(self.start_extract_stats)
        self.ui.rasterListButton.clicked.connect(self.time_series)
        self.ui.summaryComboBox.currentIndexChanged.connect(
            self.reset_progress)
        self.ui.uniqueIdComboBox.currentIndexChanged.connect(
            self.reset_progress)
        self.ui.outputLineEdit.textChanged.connect(self.reset_progress)
        self.ui.rasterLineEdit.textChanged.connect(self.reset_progress)
        self.ui.vectorLineEdit.textChanged.connect(self.vector_changed)
        self.ui.missingSpinBox.valueChanged.connect(self.reset_progress)
        self.ui.updateCheckBox.stateChanged.connect(self.existing_csv_changed)
        self.ui.helpButton.clicked.connect(self.view_manual)
        # set progress bar value
        self.ui.progressBar.setValue(0)
        # disabled form boxes
        self.ui.vectorLineEdit.setEnabled(False)
        self.ui.rasterLineEdit.setEnabled(False)
        self.ui.outputLineEdit.setEnabled(False)
        self.ui.existingCSVLineEdit.setEnabled(False)

        # Hide existing CSV until checked
        self.ui.existingCSVLineEdit.setVisible(False)
        self.ui.lblExistingCSV.setVisible(False)
        self.ui.existingCSVButton.setVisible(False)

        self.ui.missingSpinBox.setValue(config.DEFAULT_NO_DATA)
        self.orig_vector_path = ""
        self.output_file = ""
        self.update_csv_path = ""
        self.raster_dic = {}
        self.is_processing = False
        self.thread = None
        self.worker = None
        self.closed = False  # A flag to indicate if close_event is triggered
        self.close_event = None
        self.ui.cancelButton.setEnabled(False)

    def __lock_controls__(self):
        '''
        Lock controls for processing
        '''
        self.ui.vectorButton.setEnabled(False)
        self.ui.addButton.setEnabled(False)
        self.ui.removeButton.setEnabled(False)
        self.ui.saveButton.setEnabled(False)
        self.ui.loadButton.setEnabled(False)
        self.ui.outputButton.setEnabled(False)
        self.ui.processButton.setEnabled(False)
        self.ui.rasterListButton.setEnabled(False)
        self.ui.uniqueIdComboBox.setEnabled(False)
        self.ui.summaryComboBox.setEnabled(False)
        self.ui.missingSpinBox.setEnabled(False)
        self.ui.cancelButton.setEnabled(True)

    def __unlock_controls__(self):
        '''
        Unlock controls when not processing
        '''
        self.ui.vectorButton.setEnabled(True)
        self.ui.addButton.setEnabled(True)
        self.ui.removeButton.setEnabled(True)
        self.ui.saveButton.setEnabled(True)
        self.ui.loadButton.setEnabled(True)
        self.ui.outputButton.setEnabled(True)
        self.ui.processButton.setEnabled(True)
        self.ui.rasterListButton.setEnabled(True)
        self.ui.uniqueIdComboBox.setEnabled(True)
        self.ui.summaryComboBox.setEnabled(True)
        self.ui.missingSpinBox.setEnabled(True)
        self.ui.cancelButton.setEnabled(False)

    def add_raster_file_to_table(self):
        '''
        Add raster entry to GUI table.
        '''
        self.ui.progressBar.setValue(0)
        f_name = QFileDialog.getOpenFileNames(self,
                                              'Select Raster file ',
                                              self.wrksp_setup.get_climate_data_path(),
                                              config.RST_FILE_TYPES)
        for entry in f_name[0]:
            current_row_count = self.ui.rstrTableWidget.rowCount()
            self.ui.rstrTableWidget.insertRow(current_row_count)
            fixed = self.wrksp_setup.fix_os_sep_in_path(entry)
            base_name = os.path.splitext(os.path.basename(entry))[0]
            self.ui.rstrTableWidget.setItem(current_row_count, 0,
                                            QTableWidgetItem(base_name))
            self.ui.rstrTableWidget.setItem(current_row_count, 1,
                                            QTableWidgetItem(fixed))

    def browse_vector_file(self):
        '''
        Browse to vector data folder
        '''
        f_name = QFileDialog.getOpenFileName(self,
                                             'Please Select a Vector Map File',
                                             self.wrksp_setup.get_map_data_path(),
                                             config.VEC_FILE_TYPES)
        if f_name[0]:
            file_path = self.wrksp_setup.fix_os_sep_in_path(f_name[0])
            self.ui.vectorLineEdit.setText(file_path)

    def browse_csv_file(self):
        '''
        Browse to find existing csv
        '''
        f_name = QFileDialog.getOpenFileName(
            self, 'Please Select a CSV File',
            self.wrksp_setup.get_output_path(), config.CSV_FILE_TYPES)
        if f_name[0]:
            file_path = self.wrksp_setup.fix_os_sep_in_path(f_name[0])
            self.ui.existingCSVLineEdit.setText(file_path)

    def browse_list_file(self):
        '''
        Browse to timeseries file and fill GUI table.
        '''
        list_path = self.wrksp_setup.get_workspace()
        f_name = QFileDialog.getOpenFileName(
            self, 'Please Select an Image List File',
            list_path,
            'List Files (*.lst);;'
            'Comma Seperate Values (*.csv);;'
            'All Files (*.*)')
        if f_name[0]:
            file_path = self.wrksp_setup.fix_os_sep_in_path(f_name[0])
            self.ui.rasterLineEdit.setText(file_path)
            self.ui.rstrTableWidget.setRowCount(0)
            with open(f_name[0], 'r', newline=None) as file:
                reader = file.readlines()
                for row in reader:
                    current_row_count = self.ui.rstrTableWidget.rowCount()
                    file_basename = os.path.splitext(os.path.basename(row))[0]
                    self.ui.rstrTableWidget.insertRow(current_row_count)
                    self.ui.rstrTableWidget.setItem(
                        current_row_count, 0,
                        QTableWidgetItem(str(file_basename)))
                    self.ui.rstrTableWidget.setItem(
                        current_row_count, 1,
                        QTableWidgetItem(str(row).rstrip()))

    def browse_output_file(self):
        '''
        Browse to output file
        '''
        output_path = self.wrksp_setup.get_output_path()
        output_file = QFileDialog.getSaveFileName(
            self, 'Please Select an Output File',
            output_path,
            'Comma Seperated Values (*.csv);;'
            'List File (*.lst);;'
            'Text File (*.txt);;'
            'All Files (*.*)')
        if output_file[0]:
            file_path = self.wrksp_setup.fix_os_sep_in_path(output_file[0])
            self.ui.outputLineEdit.setText(file_path)

    def start_extract_stats(self):
        '''
        Start the extract stats process.
        '''
        # whatever setup needed here
        if os.path.exists(self.temp_path):
            shutil.rmtree(self.temp_path, ignore_errors=True)
        os.makedirs(self.temp_path, exist_ok=True)
        self.ui.progressBar.setValue(0)
        self.orig_vector_path = self.ui.vectorLineEdit.text()
        self.output_file = self.ui.outputLineEdit.text()
        self.update_csv_path = self.ui.existingCSVLineEdit.text()
        # checks for text in file fields
        if not self.orig_vector_path or not\
                self.output_file or not\
                self.ui.rstrTableWidget.rowCount():
            QMessageBox.information(self, "Error - Missing information",
                                    "Please specify all parameters before "
                                    "processing",
                                    QMessageBox.Ok)
            return

        self.build_raster_dic()

        # Check if the update CSV box is checked
        if self.ui.updateCheckBox.isChecked():
            # Check to make sure the CSV path is valid and has the same
            # number of features as the current selected vector file
            if self.update_csv_path:
                existing_csv_same_region = self.check_existing_csv()
                if not existing_csv_same_region:
                    QMessageBox.information(
                        self, "Error - CSV missmatch",
                        "Please select an existing CSV with the "
                        "same features as the selected vector file",
                        QMessageBox.Ok)
                    return
            else:
                QMessageBox.information(
                    self, "Error - Missing information",
                    "Please specify an existing CSV file.",
                    QMessageBox.Ok)
                return

        self.worker = wrkr(
            self.orig_vector_path,
            self.ui.uniqueIdComboBox.currentText(),
            self.raster_dic,
            self.temp_path,
            self.output_file,
            self.ui.summaryComboBox.currentText(),
            self.ui.missingSpinBox.value(),
            self.ui.updateCheckBox.isChecked(),
            self.update_csv_path)

        # lock UI
        self.ui.cancelButton.clicked.connect(
            self.worker.kill)
        self.__lock_controls__()

        # 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_extract_stats)
        self.worker.error.connect(self.extract_stats_error)
        self.worker.progress.connect(self.ui.progressBar.setValue)
        self.is_processing = True

    def finish_extract_stats(self, return_value):
        '''
        Clean up the extract stats worker object and thread.
         return_value - Returned code and string from
                                extract stats 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 = (u"Extract Stats aborted by user")
            iface.messageBar().pushCritical("Extract Stats", msg)
            QgsMessageLog.logMessage(msg, level=Qgis.Critical)
            iface.mainWindow().statusBar().showMessage(msg, 3000)
        elif return_value is not None:
            os.startfile(self.output_file)
            if os.path.exists(self.temp_path):
                shutil.rmtree(self.temp_path, ignore_errors=True)
            log_util.add_log("Extract Stats",
                             log_util.log_string("Vector File",
                                                 [self.orig_vector_path]) +
                             log_util.log_string("Unique ID Field for Extraction",
                                                 [self.ui.uniqueIdComboBox.currentText()]) +
                             log_util.log_string("Summary",
                                                 [self.ui.summaryComboBox.currentText()]),
                             [self.output_file])

            self.ui.progressBar.setValue(100)
            msg = "Extract Stats Completed"
            QgsMessageLog.logMessage(msg, level=Qgis.Info)
            iface.mainWindow().statusBar().showMessage(msg, 3000)
        else:
            # notify of an error
            msg = (u"Error: Extract Stats did not finish")
            iface.messageBar().pushCritical("Extract Stats", msg)
            QgsMessageLog.logMessage(msg, level=Qgis.Critical)
            iface.mainWindow().statusBar().showMessage(msg, 3000)

        # Unlock UI
        self.__unlock_controls__()
        self.thread_worker_cleanup()
        all_done = util.notify_process_completed("Extract Stats")
        if all_done:
            QDialog.closeEvent(self, self.close_event)

    def extract_stats_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!', 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()

    def build_raster_dic(self):
        '''
        Build raster dic, structure looks like:

         # {v2p0chirps199801: {"path": "c:/....../v2p0chirps199801.bil",
         #                     "attr": "199801",
         # 					 "prefix":  "v2p0chirps"},
         # 					....
         # 					}
        '''
        self.raster_dic = {}
        for i in range(self.ui.rstrTableWidget.rowCount()):
            f_name = self.ui.rstrTableWidget.item(i, 0).text()
            self.raster_dic[f_name] = {}
            self.raster_dic[f_name]["path"] =\
                self.ui.rstrTableWidget.item(i, 1).text()
            if len(f_name) > qgs_util.MAX_FIELDNAME_LEN:
                self.raster_dic[f_name]["attr"] =\
                    f_name[(-1 * qgs_util.MAX_FIELDNAME_LEN):]
                self.raster_dic[f_name]["prefix"] =\
                    f_name[:(-1 * qgs_util.MAX_FIELDNAME_LEN)]
            else:
                self.raster_dic[f_name]["attr"] = f_name
                self.raster_dic[f_name]["prefix"] = ""

    def remove_raster_file_from_table(self):
        '''
        Remove raster file from GUI table
        '''
        self.ui.progressBar.setValue(0)
        rows = self.ui.rstrTableWidget.selectionModel().selectedRows()
        indexes = []
        for row in rows:
            indexes.append(row.row())
        indexes = sorted(indexes, reverse=True)
        for rm_index in indexes:
            self.ui.rstrTableWidget.removeRow(rm_index)
        self.ui.progressBar.setValue(0)

    def reset_progress(self):
        '''
        Reset progress bar on form change
        '''
        self.ui.progressBar.setValue(0)

    def save_new_list_file(self):
        '''
        Save the timeseries table entries to csv
        '''
        list_table = self.wrksp_setup.get_workspace()
        f_name = QFileDialog.getSaveFileName(
            self, 'Please Select an Image List File',
            list_table,
            'List File (*.lst);;'
            'Comma Seperated Values (*.csv);;'
            'Text File (*.txt);;'
            'All files (*.*)')
        if f_name[0]:
            with open(f_name[0], 'w', newline=None) as file_save:
                row_count = self.ui.rstrTableWidget.rowCount()
                for i in range(0, row_count):
                    item = self.ui.rstrTableWidget.item(i, 1)
                    file_name = item.text()
                    file_save.writelines(file_name + '\n')

    def time_series(self):
        '''
        Open new form to select time series files, create a time series file
        and then add entries to GUI table
        '''
        self.ui.rstrTableWidget.setRowCount(0)
        t_seriesdlg = TimeSeriesListController()
        t_seriesdlg.exec_()
        selected_files_list = t_seriesdlg.available_files_list
        list_file_name = self.wrksp_setup.fix_os_sep_in_path(
            t_seriesdlg.get_file_information())
        self.ui.rasterLineEdit.setText(list_file_name)
        for file in selected_files_list:
            current_row_count = self.ui.rstrTableWidget.rowCount()
            self.ui.rstrTableWidget.insertRow(current_row_count)
            row = os.path.splitext(os.path.basename(file))[0]
            self.ui.rstrTableWidget.setItem(
                current_row_count, 0, QTableWidgetItem(str(row)))
            self.ui.rstrTableWidget.setItem(
                current_row_count, 1, QTableWidgetItem(str(file)))

    def vector_changed(self):
        '''
        Update the unique ID when the vector file changes
        '''
        self.reset_progress()
        f_name = self.ui.vectorLineEdit.text()
        layer = QgsVectorLayer(f_name, 'Shp', config.OGR)
        self.ui.uniqueIdComboBox.clear()
        for field in layer.fields():
            self.ui.uniqueIdComboBox.addItem(field.name())

    def existing_csv_changed(self):
        '''
        Update the form when the existing CSV box is checked
        '''

        if self.ui.updateCheckBox.isChecked():
            self.ui.existingCSVLineEdit.setVisible(True)
            self.ui.lblExistingCSV.setVisible(True)
            self.ui.existingCSVButton.setVisible(True)

        else:
            self.ui.existingCSVLineEdit.setVisible(False)
            self.ui.lblExistingCSV.setVisible(False)
            self.ui.existingCSVButton.setVisible(False)

    def view_manual(self):
        '''
        Function to browse to output directory and set line edit control.
        '''
        clim_manual_webpage_section = "chapter-10-extracting-raster-statistics-and-time-s"
        view_manual(clim_manual_webpage_section)

    def check_existing_csv(self):
        '''
        Check to see if the existing CSV has the same first
        column as the selected vector's fields
        '''

        csv_data = []

        # Read in the existing CSV
        with open(self.update_csv_path, 'r', newline='') as csvfile:
            reader = csv.reader(csvfile)
            for row in reader:
                csv_data.append(row)

        # Get the list of features from the vector file
        vector_layer = QgsVectorLayer(
            self.ui.vectorLineEdit.text(), "Shp", config.OGR)
        col_0 = ["Feature"]
        for feature in vector_layer.getFeatures(QgsFeatureRequest()):
            col_0.append(str(feature[self.ui.uniqueIdComboBox.currentText()]))

        if [row[0] for row in csv_data] == col_0:
            return True

        return False
