# -*- coding: UTF-8 -*-.
"""
| ----------------------------------------------------------------------------------------------------------------------
| Date                : April 2020
| Copyright           : © 2020 by Ann Crabbé
| Email               : acrabbe.foss@gmail.com
| Acknowledgements    : Based on DINEOF 4.0 (ULG Liege) by Alida Alvera-Azcárte and Alexander Barth.
|
| This file is part of the EOF-based Time Series Reconstructor QGIS plugin package.
|
| 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 3 of the License, or any later version.
|
| This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied
| warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
|
| You should have received a copy of the GNU General Public License (COPYING.txt). If not see www.gnu.org/licenses.
| ----------------------------------------------------------------------------------------------------------------------
"""
import os
import random
import shutil

from dineof.interfaces.netcdf import netcdf_file
from osgeo import gdal
from qgis.gui import QgsFileWidget
from qgis.utils import iface
from qgis.PyQt.uic import loadUi
from qgis.PyQt.QtWidgets import QDialog, QDialogButtonBox

from dineof.core.dineof import Dineof

default_CDF_placeholder = "NetCDF format only"


class DineofWidget(QDialog):
    """ QDialog to interactively set up the Neural Network input and output. """

    def __init__(self):
        super(DineofWidget, self).__init__()
        loadUi(os.path.join(os.path.dirname(__file__), 'dineof.ui'), self)

        # Data
        self.browseData.lineEdit().setReadOnly(True)
        self.browseData.lineEdit().setText(default_CDF_placeholder)
        self.browseData.fileChanged.connect(self._browse)
        self.browseDataDropDown.currentTextChanged.connect(self._select_data_variable)
        self.maxEOFValue.valueChanged.connect(self._eof_changed)

        # Mask
        self.browseMask.lineEdit().setReadOnly(True)
        self.browseMask.lineEdit().setText(default_CDF_placeholder)
        self.browseMask.fileChanged.connect(self._browse)

        # Time options
        self.timeGroup.setCollapsed(True)
        self.browseTime.lineEdit().setReadOnly(True)
        self.browseTime.lineEdit().setText(default_CDF_placeholder)
        self.browseTime.fileChanged.connect(self._browse)

        # Advanced Options
        self.advancedOptionsGroup.setCollapsed(True)

        # Output
        self.browseOut.lineEdit().setReadOnly(True)
        self.browseOut.lineEdit().setPlaceholderText('Output directory ...')
        self.browseOut.setStorageMode(QgsFileWidget.GetDirectory)

        # Open in QGIS?
        try:
            iface.activeLayer
        except AttributeError:
            self.openInQGIS.setChecked(False)
            self.openInQGIS.setDisabled(True)

        # Run or cancel
        self.OKClose.button(QDialogButtonBox.Ok).setText("Run")
        self.OKClose.accepted.connect(self._run_dineof)
        self.OKClose.rejected.connect(self.close)

    def log(self, text):
        # append text to log window
        self.logBrowser.append(str(text) + '\n')
        # open the widget on the log screen
        self.tabWidget.setCurrentIndex(self.tabWidget.indexOf(self.tab_log))

    def _browse(self, path):
        """ When the users browses for a data set, fill the drop down with all variables.
        :param path: the absolute path to the data file
        """
        if path == '':
            return

        try:
            if not gdal.Info(path) or 'netCDF' not in gdal.Info(path):
                raise FileNotFoundError("The file '{}' does not exist or is not in netCDF format.".format(path))

            if self.sender() == self.browseData:
                drop_down_widget = self.browseDataDropDown
                change_output = True
            elif self.sender() == self.browseMask:
                drop_down_widget = self.browseMaskDropDown
                change_output = False
            else:
                drop_down_widget = self.browseTimeDropDown
                change_output = False

            # have the user select the correct variable
            drop_down_widget.clear()
            drop_down_widget.addItems(['Choose variable...'] + list(netcdf_file(path).variables.keys()))

            # set default output folder name
            if change_output:
                output = os.path.join(os.path.dirname(path), 'dineof_output')
                if not os.path.exists(output):
                    self.browseOut.lineEdit().setText(output)

        except Exception as e:
            self.sender().setFilePath(default_CDF_placeholder)
            self.log(e)

    def _select_data_variable(self, variable: str):
        """ When the users selects a variable, set the correct data path and make sure the EOF SpinBoxes have correct
        min and max values.
        :param variable: drop down selection ast text
        """
        if variable == 'Choose variable...' or variable == '':
            return
        try:
            path = 'NETCDF:"' + self.browseData.lineEdit().text() + '":' + variable
            n_time = gdal.Open(path).RasterCount
            if n_time < 6:
                raise ValueError("The data set needs at least 6 days of images, please select another variable.")
            self.maxEOFValue.setMaximum(n_time - 5)
            self.krylovValue.setMinimum(self.maxEOFValue.value() + 5)
            self.krylovValue.setMaximum(n_time)
            self.minEOFValue.setMaximum(self.maxEOFValue.value())
        except Exception as e:
            self.log(e)
            self.browseDataDropDown.setCurrentIndex(0)

    def _eof_changed(self, value: int):
        """ When the users changes the EOF value, make sure the other EOF SpinBoxes have correct min and max values.
        :param value: the EOF value chosen by the user
        """
        try:
            self.krylovValue.setMinimum(value + 5)
            self.minEOFValue.setMaximum(value)
        except Exception as e:
            self.log(e)

    def _run_dineof(self):
        """ Run DINEOF after checking the input """
        try:
            if self.browseData.lineEdit().text() == default_CDF_placeholder:
                raise Exception("Choose a data set!")

            if self.browseDataDropDown.currentIndex() == 0:
                raise Exception("Choose a data set variable from the drop down list!")

            if self.browseMask.lineEdit().text() != default_CDF_placeholder:
                if self.browseMaskDropDown.currentIndex() == 0:
                    raise Exception("Choose a mask variable from the drop down list!")
                mask = 'NETCDF:"' + self.browseMask.lineEdit().text() + '":' + self.browseMaskDropDown.currentText()
            else:
                mask = None

            if self.browseOut.lineEdit().text() == "":
                raise Exception("Choose an output folder!")
            else:
                output_folder = self.browseOut.lineEdit().text()

            if self.timeGroup.isChecked():
                if self.browseTime.lineEdit().text() == default_CDF_placeholder:
                    raise Exception("Choose a netCDF file with time!")
                if self.browseTimeDropDown.currentIndex() == 0:
                    raise Exception("Choose a time variable from the drop down list!")
                time = 'NETCDF:"' + self.browseTime.lineEdit().text() + '":' + self.browseTimeDropDown.currentText()
                alpha = self.timeAlphaValue.value()
            else:
                time = None
                alpha = 0

            data_path = self.browseData.lineEdit().text()
            data_variable = self.browseDataDropDown.currentText()

            Dineof(
                data_path='NETCDF:"' + data_path + '":' + data_variable,
                output_folder=output_folder,
                eof_max=self.maxEOFValue.value(),
                krylov_max=self.krylovValue.value(),
                mask_path=mask,
                eof_min=self.minEOFValue.value(),
                eof_max_it=self.iterationsMaxValue.value(),
                eof_it_stop=self.iterationsStopValue.value(),
                lanczos=self.lanczosValue.value(),
                time_path=time,
                time_alpha=alpha,
                time_it=self.timeIterationsValue.value(),
                reconstruct_all=int(self.reconstructCheckBox.isChecked()),
                save_eof_base=int(self.saveEOFCheckBox.isChecked()),
                normalize_input=int(self.normalizeInputCheckBox.isChecked()),
                cv_seed=random.randrange(0, 500000)
            ).execute()

            if len(os.listdir(output_folder)) <= 3:
                self.log("Algorithm failed.")
                if os.path.exists(os.path.join(output_folder, 'error_log.txt')):
                    with open(os.path.join(output_folder, 'error_log.txt')) as f:
                        self.log(f.read())
                return

            # copy the lon/lat and other information from the original netcdf file
            base_path_no_ext, ext = os.path.splitext(os.path.join(output_folder, os.path.basename(data_path)))

            # rename file created by DINEOF to _temporary
            temp_path = shutil.move(base_path_no_ext + ext, base_path_no_ext + "_temporary" + ext)

            # copy original file to output folder and rename _filled
            shutil.copy(data_path, output_folder)
            filled_path = shutil.move(base_path_no_ext + ext, base_path_no_ext + "_filled" + ext)

            filled_netcdf = netcdf_file(filled_path, 'a')
            temp_netcdf = netcdf_file(temp_path, 'r')

            var = filled_netcdf.createVariable(data_variable + '_fill', 'f', ('time', 'lat', 'lon'))
            var[:] = temp_netcdf.variables[data_variable][:]
            var.units = filled_netcdf.variables[data_variable].units
            var.missing_value = temp_netcdf.variables[data_variable].missing_value
            filled_netcdf.close()
            temp_netcdf.close()

            os.remove(temp_path)

            with open(os.path.join(output_folder, 'log.txt')) as log_file:
                content = log_file.read()
                self.log(content)

            self.progressBar.setValue(100)

            if self.openInQGIS.isChecked():
                out = 'NETCDF:"' + filled_path + '":' + data_variable + '_fill'
                iface.addRasterLayer(out, "DINEOF filled data")

        except Exception as e:
            self.log(e)


def _run():
    from qgis.core import QgsApplication
    app = QgsApplication([], True)
    app.initQgis()
    widget = DineofWidget()
    widget.show()
    app.exec_()


if __name__ == '__main__':
    _run()
