# -*- 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,
                       QgsProcessingAlgorithm,
                       QgsProcessingParameterFile,
                       QgsProcessingParameterFolderDestination,
                       QgsProcessingParameterEnum,
                       QgsProcessingParameterNumber,
                       QgsProcessingParameterBoolean,
                       QgsProcessingParameterDefinition,
                       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
import yaml
import numpy as np


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

    OUTPUT_DIR = 'OUTPUT_DIR'
    INPUT_FILE = 'INPUT_FILE'
    ANTHRO = 'ANTHRO'
    NET = 'NET'
    STORAGE = 'STORAGE'
    OHM = 'OHM'
    Z0 = 'Z0'
    SMD = 'SMD'
    STAB = 'STAB'
    WU = 'WU'
    AERO = 'AERO'
    SNOW = 'SNOW'
    SPINUP = 'SPINUP'
    # TIMERESOUT = 'TIMERESOUT'
    CHUNKBOOL = 'CHUNKBOOL'
    CHUNK = 'CHUNK'

    def initAlgorithm(self, config):

        self.anthro = ((self.tr('0. Observed data'), '0'),
                      (self.tr('1. Loridan et al. 2011 (NOT RECOMMENDED)'), '1'),
                      (self.tr('2. Järvi et al. 2011 (Default)'), '2'),
                      (self.tr('3. Updated Loridan et al. 2011'), '3'),
                      (self.tr('4. Järvi et al. 2019 no CO2 Emission calculated'), '4'),
                      (self.tr('46. Järvi et al. 2019 CO2 Emission calculated'), '46'))
        self.net = ((self.tr('0. Observed data'), '0'),
                   (self.tr('1. Modelled but Ldown observed'), '1'),
                   (self.tr('2. Modelled, Ldown from cloud cover'), '2'),
                   (self.tr('3. Modelled, Ldown from Ta and RH (Default)'), '3'),
                   (self.tr('11. Same as 1 but with Lup modelled using surface temperature (Not recommended in this version)'), '11'),
                   (self.tr('12. Same as 2 but with Lup modelled using surface temperature (Not recommended in this version)'), '12'),
                   (self.tr('13. Same as 3 but with Lup modelled using surface temperature (Not recommended in this version)'), '13'),
                   (self.tr('100. Same as 1 but SSss_YYYY_NARPOut.txt file produced. (Not recommended in this version)'), '100'),
                   (self.tr('200. Same as 2 but SSss_YYYY_NARPOut.txt file produced. (Not recommended in this version)'), '200'),
                   (self.tr('300. Same as 3 but SSss_YYYY_NARPOut.txt file produced. (Not recommended in this version)'), '300'),
                   (self.tr('1001. Modelled with SPARTACUS-Surface (SS) but with Ldown modelled as in 1 (Experimental in this version)'), '1001'),
                   (self.tr('1002. Modelled with SPARTACUS-Surface (SS) but with Ldown modelled as in 1 (Experimental in this version)'), '1002'),
                   (self.tr('1003. Modelled with SPARTACUS-Surface (SS) but with Ldown modelled as in 1 (Experimental in this version)'), '1003'),)
        self.storage = ((self.tr('0. Observed data provided in meteorological  file'), '0'),
                        (self.tr('1. OHM - Objective hysteresis model (Default)'), '1'),
                       (self.tr('3. Analytical Objective hysterisis model (AnOHM) (Not recommended in this version)'), '3'),
                       (self.tr('4. Element Surface Temperature Model (Not recommended in this version)'), '4'),
                       (self.tr('6. OHM as of Liu et al, (2026) (experimental in this version)'), '6'))
        self.ohm = ((self.tr('0. From Q* (Default)'), '0'),
                   (self.tr('1. From Q* + Qf'), '1'))
        self.z0 = ((self.tr('1. as 0.1 * roughness length for momentum'), '1'),
                  (self.tr('2. Kawai et al. 2009'), '2'),
                  (self.tr('3. Voogt and Grimmond 2000'), '3'),
                  (self.tr('4. Kanda et al. 2007'), '4'),
                  (self.tr('5. Adaptively using z0m based on pervious coverage: if fully pervious, use method 1; otherwise, use method 2. (Recommended in this version)'), '5'))
        self.smd = ((self.tr('0. Modelled (Default)'), '0'),
                   (self.tr('1. Measured, volumetric data'), '1'),
                   (self.tr('2. Measured, gravimetric data'), '2'))
        self.stab = ((self.tr('2. Dyer 1974 etc.'), '2'),
                      (self.tr('3. Campbell & Norman 1998 etc. (Recommended in this version)'), '3'),
                      (self.tr('4. Businger et al. 1971'), '4'))
        self.wu = ((self.tr('0. Modelled data (Default)'), '0'),
                      (self.tr('1. Observed data'), '1'))
        self.aero = ((self.tr('1. Observed data (Default)'), '1'),
                      (self.tr('2. Rule of Thumb, Grimmond and Oke 1999'), '2'),
                      (self.tr('3. Macdonald et al. 1998'), '3'))
        
        self.addParameter(QgsProcessingParameterFile(self.INPUT_FILE,
            self.tr('Input yaml file(.yml)'), extension='yml'))
    
        self.addParameter(QgsProcessingParameterEnum(self.NET,
                                                     self.tr('Net radiation method'),
                                                     options=[i[0] for i in self.net],
                                                     defaultValue=3))
        self.addParameter(QgsProcessingParameterEnum(self.ANTHRO,
                                                     self.tr('Anthropogenic heat flux method'),
                                                     options=[i[0] for i in self.anthro],
                                                     defaultValue=2))
        self.addParameter(QgsProcessingParameterEnum(self.STORAGE,
                                                     self.tr('Storage heat flux method'),
                                                     options=[i[0] for i in self.storage],
                                                     defaultValue=1))
        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=1))
        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'))
        
        #Advanced parameters
        chunkBool = QgsProcessingParameterBoolean(self.CHUNKBOOL,
                                                    self.tr("Devide calculation in chunks to reduce issues with memory running low on your computer."),
                                                    defaultValue=False)
        chunkBool.setFlags(chunkBool.flags() | QgsProcessingParameterDefinition.FlagAdvanced)
        self.addParameter(chunkBool)

        chunk = QgsProcessingParameterNumber(self.CHUNK, self.tr('Number of chunks'),
                QgsProcessingParameterNumber.Integer,
                QVariant(2), optional=True, minValue=0, maxValue=1000)
        chunk.setFlags(chunk.flags() | QgsProcessingParameterDefinition.FlagAdvanced)
        self.addParameter(chunk)


    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)
        infile = self.parameterAsString(parameters, self.INPUT_FILE, 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) #issue 766
        # spinup = self.parameterAsBool(parameters, self.SPINUP, context)
        chunkBool = self.parameterAsBool(parameters, self.CHUNKBOOL, context)
        noOfChunks = self.parameterAsInt(parameters, self.CHUNK, context)

        feedback.setProgressText(self.supylib)
        
        feedback.setProgressText("Reading and updating Yaml-file")

        with open(infile, 'r') as f:
            yaml_dict = yaml.load(f, Loader=yaml.SafeLoader)
        
         # nml['runcontrol']['CBLuse'] = int(usecbl)
        yaml_dict['model']['physics']['snowuse']['value'] = int(usesnow)
        yaml_dict['model']['physics']['netradiationmethod']['value'] = int(self.net[int(net)][1])
        yaml_dict['model']['physics']['emissionsmethod']['value'] = int(self.anthro[int(qf)][1])
        yaml_dict['model']['physics']['ohmincqf']['value'] = int(self.ohm[int(ohm)][1])
        yaml_dict['model']['physics']['stabilitymethod']['value'] = int(self.stab[int(stab)][1])
        yaml_dict['model']['physics']['storageheatmethod']['value'] = int(self.storage[int(qs)][1])
        yaml_dict['model']['physics']['roughlenmommethod']['value'] = int(self.aero[int(aeroD)][1])
        yaml_dict['model']['physics']['roughlenheatmethod']['value'] = int(self.z0[int(z0)][1])
        yaml_dict['model']['physics']['smdmethod']['value'] = int(self.smd[int(smd)][1])
        yaml_dict['model']['physics']['waterusemethod']['value'] = int(self.wu[int(wu)][1])
        # yaml_dict['model']['physics']['roughlenheatmethod'] = str(infolder) + "/"
        yaml_dict['model']['control']['output_file'] = str(outfolder) + "/"
        # yaml_dict['model']['control']['forcing_file'] = str(filecode)
        #yaml_dict['model']['physics']['tstep'] = int(int(outputRes) * 60.)
        # yaml_dict['model']['physics']['tstep'] = int(int(inputRes) * 60.)
        
        # nml.write(Path(str(infolder) + '/RunControl.nml'), force=True)
        
        #####################################################################################
        # SuPy
        with open(infile, 'w') as file:
            yaml.dump(yaml_dict, file, sort_keys = False)

        print(yaml_dict['model']['physics'])
        # SuPy initialisation
        # path_runcontrol = Path(infolder) / 'RunControl.nml'
        feedback.setProgressText("Initiating model")
        config = sp.data_model.init_config_from_yaml(infile)
        df_state_init = config.to_df_state()
        
        feedback.setProgressText("Loading forcing data")
        
        grid = df_state_init.index[0]
        df_forcing = sp.load_forcing_grid(infile, grid, df_state_init=df_state_init)

        if chunkBool:
            noOfDays = (df_forcing.index.max() - df_forcing.index.min()).days
            chunkDay = np.ceil(noOfDays / noOfChunks)
            feedback.setProgressText("Model run divided into " + str(int(chunkDay)) + ' day period')
        else:
            chunkDay = 3660
       
        # SuPy simulation
        feedback.setProgressText("Running model (QGIS not responsive)")
        df_output, df_state_final = sp.run_supy(df_forcing,
                                                df_state_init,
                                                chunk_day=chunkDay
                                                )

        # use SuPy function to save results
        feedback.setProgressText("Saving to disk")
        sp.save_supy(
            df_output,
            df_state_final,
            path_dir_save = yaml_dict['model']['control']['output_file'] )
        #####################################################################################

        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 v2025.7.9.dev0')

    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()
