# -*- coding: utf-8 -*-

"""
/***************************************************************************
 ProcessingUMEP
                                 A QGIS plugin
 UMEP for processing toolbox
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2020-04-02
        copyright            : (C) 2020 by Fredrik Lindberg
        email                : fredrikl@gvc.gu.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     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""

__author__ = 'Fredrik Lindberg'
__date__ = '2020-04-02'
__copyright__ = '(C) 2020 by Fredrik Lindberg'

# This will get replaced with a git SHA1 when you do a git archive

__revision__ = '$Format:%H$'

from qgis.PyQt.QtCore import QCoreApplication, QVariant
# from qgis.PyQt.QtWidgets import QMessageBox
from qgis.core import (QgsProcessing,
                       QgsFeatureSink,
                       QgsProcessingAlgorithm,
                       QgsProcessingParameterFile,
                       QgsProcessingParameterFolderDestination,
                       QgsProcessingParameterEnum,
                       QgsProcessingParameterNumber,
                       QgsProcessingParameterBoolean,
                       QgsProcessingException)

# try:
#     import supy as sp
#     from supy import __version__ as ver_supy
# except:
#     pass
from pathlib import Path
from ..util import f90nml
import sys, os
from qgis.PyQt.QtGui import QIcon
import inspect
from pathlib import Path

class ProcessingSuewsAlgorithm(QgsProcessingAlgorithm):
    """
    This is a processing algorithm for the SUEWS model
    """

    OUTPUT_DIR = 'OUTPUT_DIR'
    INPUT_DIR = 'INPUT_DIR'
    ANTHRO = 'ANTHRO'
    NET = 'NET'
    STORAGE = 'STORAGE'
    OHM = 'OHM'
    Z0 = 'Z0'
    SMD = 'SMD'
    STAB = 'STAB'
    WU = 'WU'
    AERO = 'AERO'
    SNOW = 'SNOW'
    SPINUP = 'SPINUP'
    TIMERESOUT = 'TIMERESOUT'

    def initAlgorithm(self, config):

        self.anthro = ((self.tr('Observed data'), '0'),
                      (self.tr('Loridan et al. 2011 (NOT RECOMMENDED)'), '1'),
                      (self.tr('Järvi et al. 2011 (Default)'), '2'))
        self.net = ((self.tr('Observed data'), '0'),
                   (self.tr('Modelled but Ldown observed'), '1'),
                   (self.tr('Modelled, Ldown from cloud cover'), '2'),
                   (self.tr('Modelled, Ldown from Ta and RH (Default)'), '3'))
        self.storage = ((self.tr('OHM - Objective hysteresis model (Default)'), '1'),
                       (self.tr('Observed data'), '2'),
                       (self.tr('Analytical Objective hysterisis model (AnOHM)'), '3'),
                       (self.tr('Element Surface Temperature Model'), '4'))
        self.ohm = ((self.tr('From Q* (Default)'), '0'),
                   (self.tr('[1] From Q* + Qf'), '1'))
        self.z0 = ((self.tr('as 0.1 * roughness length for momentum'), '1'),
                  (self.tr('Kawai et al. 2009 (Default)'), '2'),
                  (self.tr('Voogt and Grimmond 2000'), '3'),
                  (self.tr('Kanda et al. 2007'), '4'))
        self.smd = ((self.tr('Modelled (Default)'), '0'),
                   (self.tr('Measured, volumetric data'), '1'),
                   (self.tr('Measured, gravimetric data'), '2'))
        self.stab = ((self.tr('Dyer 1974 etc.'), '2'),
                      (self.tr('Campbell & Norman 1998 etc. (Default)'), '3'),
                      (self.tr('Businger et al. 1971'), '4'))
        self.wu = ((self.tr('Modelled data (Default)'), '0'),
                      (self.tr('Observed data'), '1'))
        self.aero = ((self.tr('Observed data (Default)'), '1'),
                      (self.tr('Rule of Thumb, Grimmond and Oke 1999'), '2'),
                      (self.tr('Macdonald et al. 1998'), '3'))

        self.addParameter(QgsProcessingParameterFile(self.INPUT_DIR,
                                                     'Input folder',
                                                     QgsProcessingParameterFile.Folder))
        self.addParameter(QgsProcessingParameterEnum(self.ANTHRO,
                                                     self.tr('Net radiation method'),
                                                     options=[i[0] for i in self.anthro],
                                                     defaultValue=2))
        self.addParameter(QgsProcessingParameterEnum(self.NET,
                                                     self.tr('Anthropogenic heat flux method'),
                                                     options=[i[0] for i in self.net],
                                                     defaultValue=3))
        self.addParameter(QgsProcessingParameterEnum(self.STORAGE,
                                                     self.tr('Storage heat flux method'),
                                                     options=[i[0] for i in self.storage],
                                                     defaultValue=0))
        self.addParameter(QgsProcessingParameterEnum(self.OHM,
                                                     self.tr('OHM option'),
                                                     options=[i[0] for i in self.ohm],
                                                     defaultValue=0))
        self.addParameter(QgsProcessingParameterEnum(self.Z0,
                                                     self.tr('Roughness length for heat method'),
                                                     options=[i[0] for i in self.z0],
                                                     defaultValue=1))
        self.addParameter(QgsProcessingParameterEnum(self.SMD,
                                                     self.tr('Soil moisture deficit method'),
                                                     options=[i[0] for i in self.smd],
                                                     defaultValue=0))
        self.addParameter(QgsProcessingParameterEnum(self.STAB,
                                                     self.tr('Atmospheric stability method'),
                                                     options=[i[0] for i in self.stab],
                                                     defaultValue=1))
        self.addParameter(QgsProcessingParameterEnum(self.WU,
                                                     self.tr('External water use method'),
                                                     options=[i[0] for i in self.wu],
                                                     defaultValue=0))
        self.addParameter(QgsProcessingParameterEnum(self.AERO,
                                                     self.tr('Aerodynamic properties:'),
                                                     options=[i[0] for i in self.aero],
                                                     defaultValue=0))
        self.addParameter(QgsProcessingParameterBoolean(self.SNOW,
                                                        self.tr("Use snow module"),
                                                        defaultValue=False))
        # self.addParameter(QgsProcessingParameterBoolean(self.SPINUP,
        #                                                 self.tr("Apply spin-up using existing meteorological data (only possible if one full year of data is used)"),
        #                                                 defaultValue=False))
        self.addParameter(QgsProcessingParameterNumber(self.TIMERESOUT, 
                                                       self.tr("Output time resolution (minutes)"),
                                                       QgsProcessingParameterNumber.Integer,
                                                       QVariant(60),
                                                       minValue=1))                                                                                       
        self.addParameter(QgsProcessingParameterFolderDestination(self.OUTPUT_DIR,
                                                     'Output folder'))


    def processAlgorithm(self, parameters, context, feedback):
        try:
            import supy as sp
            from supy import __version__ as ver_supy
        except:
            raise QgsProcessingException('This plugin requires the supy package '
                        'to be installed OR upgraded. Please consult the FAQ in the manual '
                        'for further information on how to install missing python packages.')
            # QMessageBox.critical(None, 'Error', 'This plugin requires the supy package '
            #             'to be installed OR upgraded. Please consult the FAQ in the manual '
            #             'for further information on how to install missing python packages.')
            # return
        feedback.setProgressText('SuPy version: ' + ver_supy)
        self.supylib = sys.modules["supy"].__path__[0]
        feedback.setProgressText(self.supylib)
        infolder = self.parameterAsString(parameters, self.INPUT_DIR, context)
        outfolder = self.parameterAsString(parameters, self.OUTPUT_DIR, context)
        if self.parameterAsBool(parameters, self.SNOW, context):
            usesnow = 1
        else:
            usesnow = 0
        net = self.parameterAsString(parameters, self.NET, context)
        qf = self.parameterAsString(parameters, self.ANTHRO, context)
        ohm = self.parameterAsString(parameters, self.OHM, context)
        stab = self.parameterAsString(parameters, self.STAB, context)
        qs = self.parameterAsString(parameters, self.STORAGE, context)
        aeroD = self.parameterAsString(parameters, self.AERO, context)
        z0 = self.parameterAsString(parameters, self.Z0, context)
        smd = self.parameterAsString(parameters, self.SMD, context)
        wu = self.parameterAsString(parameters, self.WU, context)
        outputRes = self.parameterAsInt(parameters, self.TIMERESOUT, context)
        # spinup = self.parameterAsBool(parameters, self.SPINUP, context)

        feedback.setProgressText(self.supylib)
        
        feedback.setProgressText("Creating RunControl.nml")

        # Create modified RunControl.nml
        filenamemetdata = None
        for file in os.listdir(infolder):
            if 'data' in file:
                filenamemetdata = file

        underscorePos = ([pos for pos, char in enumerate(filenamemetdata) if char == '_'])
        if (underscorePos[1] - underscorePos[0]) == 1:
            addunderscore = '_'
        else:
            addunderscore = ''
        numunderscores = underscorePos.__len__()
        inputRes = filenamemetdata[underscorePos[numunderscores - 1] + 1:filenamemetdata.find('.')]
        filecode = filenamemetdata[0:underscorePos[0]] + addunderscore
        
        # nml = f90nml.read(self.model_dir + '/BaseFiles/RunControl.nml')
        nml = f90nml.read(self.supylib + '/sample_run/RunControl.nml')
        # nml['runcontrol']['CBLuse'] = int(usecbl)
        nml['runcontrol']['SnowUse'] = int(usesnow)
        nml['runcontrol']['NetRadiationMethod'] = int(net)
        nml['runcontrol']['EmissionsMethod'] = int(qf)
        nml['runcontrol']['OHMIncQF'] = int(ohm)
        nml['runcontrol']['StabilityMethod'] = int(stab) + 2
        nml['runcontrol']['StorageHeatMethod'] = int(qs) + 1
        nml['runcontrol']['RoughLenMomMethod'] = int(aeroD) + 1
        nml['runcontrol']['RoughLenHeatMethod'] = int(z0) + 1
        nml['runcontrol']['SMDMethod'] = int(smd)
        nml['runcontrol']['WaterUseMethod'] = int(wu)
        nml['runcontrol']['fileinputpath'] = str(infolder) + "/"
        nml['runcontrol']['fileoutputpath'] = str(outfolder) + "/"
        nml['runcontrol']['fileCode'] = str(filecode)
        nml['runcontrol']['ResolutionFilesOut'] = int(int(outputRes) * 60.)
        nml['runcontrol']['ResolutionFilesIn'] = int(int(inputRes) * 60.)
        
        nml.write(Path(str(infolder) + '/RunControl.nml'), force=True)
        
        #####################################################################################
        # SuPy

        # SuPy initialisation
        path_runcontrol = Path(infolder) / 'RunControl.nml'
        feedback.setProgressText("Initiating model")
        df_state_init = sp.init_supy(path_runcontrol)
        grid = df_state_init.index[0]
        feedback.setProgressText("Loading forcing data")
        df_forcing = sp.load_forcing_grid(path_runcontrol, grid)

        # SuPy simulation
        feedback.setProgressText("Running model (QGIS not responsive)")
        df_output, df_state_final = sp.run_supy(df_forcing,
                                                df_state_init,
                                                check_input=True,
                                                serial_mode=True,
                                                )

        # resampling SuPy results for plotting
        # df_output_suews = df_output.loc[grid, 'SUEWS']
        # df_output_suews_rsmp = df_output_suews.resample('1h').mean()

        # use SuPy function to save results
        feedback.setProgressText("Saving to disk")
        list_path_save = sp.save_supy(df_output,
                                    df_state_final,
                                    path_runcontrol=path_runcontrol)
        #####################################################################################

        feedback.setProgressText('Model finished')

        return {self.OUTPUT_DIR: outfolder}

    def name(self):
        return 'Urban Energy Balance: SUEWS'

    def displayName(self):
        return self.tr('Urban Energy Balance: SUEWS v2020a')

    def group(self):
        return self.tr(self.groupId())

    def groupId(self):
        return 'Processor'

    def shortHelpString(self):
        return self.tr('SUEWS - Surface Urban Energy and Water Balance Scheme (Järvi et al. 2011, 2014, Ward et al. 2016) simulates the urban radiation, '
                       'energy and water balances using commonly measured/modeled meteorological variables and '
                       'information about the surface cover. It utilizes an evaporation-interception approach '
                       '(Grimmond et al. 1991), similar to that used in forests, to model evaporation from urban surfaces.<br>'
                       '---------------\n'
                       'Järvi L, Grimmond CSB & Christen A (2011) The Surface Urban Energy and Water Balance Scheme (SUEWS): Evaluation in Los Angeles and Vancouver J. Hydrol. 411, 219-237.<br>'
                        '\n'
                        'Järvi L, Grimmond CSB, Taka M, Nordbo A, Setälä H &Strachan IB (2014) Development of the Surface Urban Energy and Water balance Scheme (SUEWS) for cold climate cities, Geosci. Model Dev. 7, 1691-1711, doi:10.5194/gmd-7-1691-2014.<br>'
                        '\n'
                        'Ward HC, L Järvi, S Onomura, F Lindberg, CSB Grimmond (2016a) SUEWS Manual: Version 2016a<br>'
                       '---------------\n'
                       'Full manual available via the <b>Help</b>-button.')

    def helpUrl(self):
        url = "https://umep-docs.readthedocs.io/en/latest/processor/Urban%20Energy%20Balance%20Urban%20Energy%20Balance%20(SUEWS.BLUEWS,%20advanced).html"
        return url
    
    def tr(self, string):
        return QCoreApplication.translate('Processing', string)

    def icon(self):
        cmd_folder = Path(os.path.split(inspect.getfile(inspect.currentframe()))[0]).parent
        icon = QIcon(str(cmd_folder) + "/icons/SuewsLogo.png")
        return icon

    def createInstance(self):
        return ProcessingSuewsAlgorithm()
