#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# PCR-GLOBWB (PCRaster Global Water Balance) Global Hydrological Model
#
# Copyright (C) 2016, Ludovicus P. H. (Rens) van Beek, Edwin H. Sutanudjaja, Yoshihide Wada,
# Joyce H. C. Bosmans, Niels Drost, Inge E. M. de Graaf, Kor de Jong, Patricia Lopez Lopez,
# Stefanie Pessenteiner, Oliver Schmitz, Menno W. Straatsma, Niko Wanders, Dominik Wisser,
# and Marc F. P. Bierkens,
# Faculty of Geosciences, Utrecht University, Utrecht, The Netherlands
#
# 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
# (at your option) 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
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os

import pcraster as pcr
from . import virtualOS as vos

from wflow.wflow_lib import configsection
from wflow.wflow_lib import configget

class SoilAndTopoParameters(object):

    def __init__(self, iniItems, landmask, Dir, cloneMap, tmpDir):
        object.__init__(self)

        # cloneMap, tmpDir, inputDir based on the configuration/setting given in the ini/configuration file
        self.cloneMap =  cloneMap #iniItems.cloneMap
        self.tmpDir   = tmpDir #iniItems.tmpDir
        self.inputDir = Dir #iniItems.globalOptions['inputDir']
        self.landmask = landmask
        
        # How many soil layers (excluding groundwater):
        self.numberOfLayers = int(configget(iniItems,"landSurfaceOptions","numberOfUpperSoilLayers","2"))

    def read(self, iniItems, optionDict = None):
		
        self.readTopo(iniItems, optionDict)
        self.readSoil(iniItems, optionDict)

    def readTopo(self, iniItems, optionDict):

        # a dictionary/section of options that will be used
        if optionDict == None: optionDict = iniItems._sections["landSurfaceOptions"]

        # maps of elevation attributes: 
        topoParams = ['tanslope','slopeLength','orographyBeta']
        if optionDict['topographyNC'] == str(None):
            for var in topoParams:
                input = configget(iniItems,"landSurfaceOptions",str(var),"None")
                vars(self)[var] = vos.readPCRmapClone(input,self.cloneMap,
                                                self.tmpDir,self.inputDir)
                if var != "slopeLength": vars(self)[var] = pcr.cover(vars(self)[var], 0.0)
        else:
            topoPropertiesNC = vos.getFullPath(\
                               optionDict['topographyNC'],
                                                self.inputDir)

            for var in topoParams:
                vars(self)[var] = vos.netcdf2PCRobjCloneWithoutTime(\
                                    topoPropertiesNC,var, \
                                    cloneMapFileName = self.cloneMap)
                if var != "slopeLength": vars(self)[var] = pcr.cover(vars(self)[var], 0.0)

        #~ self.tanslope = pcr.max(self.tanslope, 0.00001)              # In principle, tanslope can be zero. Zero tanslope will provide zero TCL (no interflow)

        # covering slopeLength with its maximum value 
        self.slopeLength = pcr.cover(self.slopeLength, pcr.mapmaximum(self.slopeLength))
        
        # maps of relative elevation above flood plains 
        dzRel = ['dzRel0001','dzRel0005',
                 'dzRel0010','dzRel0020','dzRel0030','dzRel0040','dzRel0050',
                 'dzRel0060','dzRel0070','dzRel0080','dzRel0090','dzRel0100']
        if optionDict['topographyNC'] == str(None):
            for i in range(0, len(dzRel)):
                var = dzRel[i]
                input = optionDict[str(var)]
                vars(self)[var] = vos.readPCRmapClone(input,self.cloneMap,
                                                self.tmpDir,self.inputDir)
                vars(self)[var] = pcr.cover(vars(self)[var], 0.0)
                if i > 0: vars(self)[var] = pcr.max(vars(self)[var], vars(self)[dzRel[i-1]])
        else:
            for i in range(0, len(dzRel)):
                var = dzRel[i]
                vars(self)[var] = vos.netcdf2PCRobjCloneWithoutTime(\
                                    topoPropertiesNC,var, \
                                    cloneMapFileName = self.cloneMap)
                vars(self)[var] = pcr.cover(vars(self)[var], 0.0)
                if i > 0: vars(self)[var] = pcr.max(vars(self)[var], vars(self)[dzRel[i-1]])

    def readSoilMapOfFAO(self, iniItems, optionDict = None):

        # a dictionary/section of options that will be used
        if optionDict == None: optionDict = iniItems._sections["landSurfaceOptions"] #iniItems.landSurfaceOptions
        
        # soil variable names given either in the ini or netCDF file:
        soilParameters = ['airEntryValue1','airEntryValue2',       
                          'poreSizeBeta1','poreSizeBeta2',        
                          'resVolWC1','resVolWC2',            
                          'satVolWC1','satVolWC2',
                          'KSat1','KSat2',
                          'percolationImp']
        if optionDict['soilPropertiesNC'] == str(None):
            for var in soilParameters:
                input = optionDict[str(var)]
                vars(self)[var] = \
                               vos.readPCRmapClone(input,self.cloneMap,\
                                             self.tmpDir,self.inputDir)
                vars(self)[var] = pcr.scalar(vars(self)[var])

                if input == "percolationImp": vars(self)[var] = pcr.cover(vars(self)[var], 0.0)
                
                # extrapolation 
                # - TODO: Make a general extrapolation option as a function in the virtualOS.py 
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 0.75))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))
                vars(self)[var] = pcr.cover(vars(self)[var], 0.0)

        else:
            soilPropertiesNC = vos.getFullPath(\
                               optionDict['soilPropertiesNC'],
                                                self.inputDir)
            for var in soilParameters:
                vars(self)[var] = vos.netcdf2PCRobjCloneWithoutTime(\
                                    soilPropertiesNC,var, \
                                    cloneMapFileName = self.cloneMap)

                if var == "percolationImp": vars(self)[var] = pcr.cover(vars(self)[var], 0.0)

                # extrapolation 
                # - TODO: Make a general extrapolation option as a function in the virtualOS.py 
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 0.75))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))
                vars(self)[var] = pcr.cover(vars(self)[var],
                                  pcr.windowaverage(vars(self)[var], 1.00))

                vars(self)[var] = pcr.cover(vars(self)[var], 0.01)
            
        # make sure that resVolWC1 <= satVolWC1
        self.resVolWC1 = pcr.min(self.resVolWC1, self.satVolWC1)
        self.resVolWC2 = pcr.min(self.resVolWC2, self.satVolWC2)
        
        if self.numberOfLayers == 2:
            self.satVolMoistContUpp = self.satVolWC1                         # saturated volumetric moisture content (m3.m-3)
            self.satVolMoistContLow = self.satVolWC2
            self.resVolMoistContUpp = self.resVolWC1                         # residual volumetric moisture content (m3.m-3)
            self.resVolMoistContLow = self.resVolWC2
            self.airEntryValueUpp   = self.airEntryValue1                    # air entry value (m) according to soil water retention curve of Clapp & Hornberger (1978)
            self.airEntryValueLow   = self.airEntryValue2
            self.poreSizeBetaUpp    = self.poreSizeBeta1                     # pore size distribution parameter according to Clapp & Hornberger (1978) 
            self.poreSizeBetaLow    = self.poreSizeBeta2
            self.kSatUpp            = self.KSat1                             # saturated hydraulic conductivity (m.day-1) 
            self.kSatLow            = self.KSat2

        if self.numberOfLayers == 3:
            self.satVolMoistContUpp000005 = self.satVolWC1     
            self.satVolMoistContUpp005030 = self.satVolWC1     
            self.satVolMoistContLow030150 = self.satVolWC2
            self.resVolMoistContUpp000005 = self.resVolWC1     
            self.resVolMoistContUpp005030 = self.resVolWC1     
            self.resVolMoistContLow030150 = self.resVolWC2
            self.airEntryValueUpp000005   = self.airEntryValue1
            self.airEntryValueUpp005030   = self.airEntryValue1
            self.airEntryValueLow030150   = self.airEntryValue2
            self.poreSizeBetaUpp000005    = self.poreSizeBeta1 
            self.poreSizeBetaUpp005030    = self.poreSizeBeta1 
            self.poreSizeBetaLow030150    = self.poreSizeBeta2
            self.kSatUpp000005            = self.KSat1         
            self.kSatUpp005030            = self.KSat1         
            self.kSatLow030150            = self.KSat2

        self.percolationImp = pcr.cover(self.percolationImp, 0.0)            # fractional area where percolation to groundwater store is impeded (dimensionless)

        # soil thickness and storage variable names 
        # as given either in the ini or netCDF file:
        soilStorages = ['firstStorDepth',      'secondStorDepth',      
                        'soilWaterStorageCap1','soilWaterStorageCap2'] 
        if optionDict['soilPropertiesNC'] == str(None):
            for var in soilStorages:
                input = optionDict[str(var)]
                temp = str(var)+'Inp'
                vars(self)[temp] = vos.readPCRmapClone(input,\
                                            self.cloneMap,
                                            self.tmpDir,self.inputDir)

                # extrapolation 
                # - TODO: Make a general extrapolation option as a function in the virtualOS.py 
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 0.75))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp], 0.0)

        else:
            soilPropertiesNC = vos.getFullPath(\
                               optionDict['soilPropertiesNC'],
                                                self.inputDir)
            for var in soilStorages:
                temp = str(var)+'Inp'
                vars(self)[temp] = vos.netcdf2PCRobjCloneWithoutTime(\
                                     soilPropertiesNC,var, \
                                     cloneMapFileName = self.cloneMap)
                # extrapolation 
                # - TODO: Make a general extrapolation option as a function in the virtualOS.py 
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 0.75))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp],
                                   pcr.windowaverage(vars(self)[temp], 1.05))
                vars(self)[temp] = pcr.cover(vars(self)[temp], 0.0)

        # layer thickness
        if self.numberOfLayers == 2:
            self.thickUpp = (0.30/0.30)*self.firstStorDepthInp
            self.thickLow = (1.20/1.20)*self.secondStorDepthInp
        if self.numberOfLayers == 3:
            self.thickUpp000005 = (0.05/0.30)*self.firstStorDepthInp
            self.thickUpp005030 = (0.25/0.30)*self.firstStorDepthInp
            self.thickLow030150 = (1.20/1.20)*self.secondStorDepthInp

        # soil storage
        if self.numberOfLayers == 2:
            #~ self.storCapUpp = (0.30/0.30)*self.soilWaterStorageCap1Inp
            #~ self.storCapLow = (1.20/1.20)*self.soilWaterStorageCap2Inp                     # 22 Feb 2014: We can calculate this based on thickness and porosity. 
            self.storCapUpp = self.thickUpp * \
                             (self.satVolMoistContUpp - self.resVolMoistContUpp)
            self.storCapLow = self.thickLow * \
                             (self.satVolMoistContLow - self.resVolMoistContLow)
            self.rootZoneWaterStorageCap = self.storCapUpp + \
                                           self.storCapLow                                    # This is called as WMAX in the original pcrcalc script. 
        if self.numberOfLayers == 3:
            self.storCapUpp000005 = self.thickUpp000005 * \
                             (self.satVolMoistContUpp000005 - self.resVolMoistContUpp000005)
            self.storCapUpp005030 = self.thickUpp005030 * \
                             (self.satVolMoistContUpp005030 - self.resVolMoistContUpp005030)
            self.storCapLow030150 = self.thickLow030150 * \
                             (self.satVolMoistContLow030150 - self.resVolMoistContLow030150)
            self.rootZoneWaterStorageCap = self.storCapUpp000005 + \
                                           self.storCapUpp005030 + \
                                           self.storCapLow030150

    def readSoil(self, iniItems, optionDict = None):

        # a dictionary/section of options that will be used
        if optionDict == None: optionDict = iniItems._sections["landSurfaceOptions"] #iniItems.landSurfaceOptions

        # default values of soil parameters that are constant/uniform for the entire domain:
        self.clappAddCoeff   = pcr.scalar(3.0)        # dimensionless
        self.matricSuctionFC = pcr.scalar(1.0)        # unit: m
        #~ self.matricSuction50 = pcr.scalar(10./3.)  # unit: m
        self.matricSuction50 = pcr.scalar(3.33)       # unit: m
        self.matricSuctionWP = pcr.scalar(156.0)      # unit: m
        self.maxGWCapRise    = pcr.scalar(5.0)        # unit: m
        #  
        # values defined in the ini/configuration file:
        soilParameterConstants = ['clappAddCoeff',        
                                  'matricSuctionFC',      
                                  'matricSuction50',      
                                  'matricSuctionWP',      
                                  'maxGWCapRise']
        for var in soilParameterConstants:
            if var in configsection(iniItems,"landSurfaceOptions"):
                input = configget(iniItems,"landSurfaceOptions",str(var),"None")
                vars(self)[var] = vos.readPCRmapClone(input,self.cloneMap,\
                                                            self.tmpDir,self.inputDir)
        
        # read soil parameter based on the FAO soil map:
        self.readSoilMapOfFAO(iniItems, optionDict)                                 
       
        # assign Campbell's (1974) beta coefficient, as well as degree 
        # of saturation at field capacity and corresponding unsaturated hydraulic conductivity
        #
        if self.numberOfLayers == 2:

            self.campbellBetaUpp = self.poreSizeBetaUpp*2.0 + \
                                      self.clappAddCoeff                       # Campbell's (1974) coefficient ; Rens's line: BCB = 2*BCH + BCH_ADD
            self.campbellBetaLow = self.poreSizeBetaLow*2.0 + \
                                      self.clappAddCoeff                

            self.effSatAtFieldCapUpp = \
                     (self.matricSuctionFC / self.airEntryValueUpp)**\
                                        (-1.0/ self.poreSizeBetaUpp )          # saturation degree at field capacity       : THEFF_FC = (PSI_FC/PSI_A)**(-1/BCH)
            self.effSatAtFieldCapUpp = pcr.cover(self.effSatAtFieldCapUpp, 1.0)
            
            self.effSatAtFieldCapLow = \
                     (self.matricSuctionFC / self.airEntryValueLow)**\
                                        (-1.0/ self.poreSizeBetaLow )
            self.effSatAtFieldCapLow = pcr.cover(self.effSatAtFieldCapLow, 1.0)

            self.kUnsatAtFieldCapUpp = pcr.max(0., \
             (self.effSatAtFieldCapUpp ** self.campbellBetaUpp) * self.kSatUpp)  # unsaturated conductivity at field capacity: KTHEFF_FC = max(0,THEFF_FC[TYPE]**BCB*KS1)
            self.kUnsatAtFieldCapLow = pcr.max(0., \
             (self.effSatAtFieldCapLow ** self.campbellBetaLow) * self.kSatLow)
        #
        if self.numberOfLayers == 3:

            self.campbellBetaUpp000005 = self.poreSizeBetaUpp000005*2.0 + \
                                            self.clappAddCoeff
            self.campbellBetaUpp005030 = self.poreSizeBetaUpp005030*2.0 + \
                                            self.clappAddCoeff                
            self.campbellBetaLow030150 = self.poreSizeBetaLow030150*2.0 + \
                                            self.clappAddCoeff                

            self.effSatAtFieldCapUpp000005 = \
                     (self.matricSuctionFC / self.airEntryValueUpp000005)**\
                                        (-1.0/ self.poreSizeBetaUpp000005)
            self.effSatAtFieldCapUpp005030 = \
                     (self.matricSuctionFC / self.airEntryValueUpp005030)**\
                                        (-1.0/ self.poreSizeBetaUpp005030)
            self.effSatAtFieldCapLow030150 = \
                     (self.matricSuctionFC / self.airEntryValueLow030150)**\
                                        (-1.0/ self.poreSizeBetaLow030150)

            self.kUnsatAtFieldCapUpp000005 = pcr.max(0., \
             (self.effSatAtFieldCapUpp000005 ** self.campbellBetaUpp000005) * self.kSatUpp000005)
            self.kUnsatAtFieldCapUpp005030 = pcr.max(0., \
             (self.effSatAtFieldCapUpp005030 ** self.campbellBetaUpp005030) * self.kSatUpp005030)
            self.kUnsatAtFieldCapLow030150 = pcr.max(0., \
             (self.effSatAtFieldCapLow030150 ** self.campbellBetaLow030150) * self.kSatLow030150)

        # calculate degree of saturation at which transpiration is halved (50) 
        # and at wilting point
        #
        if self.numberOfLayers == 2:
            self.effSatAt50Upp = (self.matricSuction50/self.airEntryValueUpp)**\
                                                    (-1.0/self.poreSizeBetaUpp)
            self.effSatAt50Upp = pcr.cover(self.effSatAt50Upp, 1.0)                                        
            self.effSatAt50Low = (self.matricSuction50/self.airEntryValueLow)**\
                                                    (-1.0/self.poreSizeBetaLow)
            self.effSatAt50Low = pcr.cover(self.effSatAt50Low, 1.0)                                        
            self.effSatAtWiltPointUpp = pcr.cover(\
                                 (self.matricSuctionWP/self.airEntryValueUpp)**\
                                                    (-1.0/self.poreSizeBetaUpp), 1.0)
            self.effSatAtWiltPointLow = pcr.cover(\
                                 (self.matricSuctionWP/self.airEntryValueLow)**\
                                                    (-1.0/self.poreSizeBetaLow), 1.0)
        if self.numberOfLayers == 3:
            self.effSatAt50Upp000005 = (self.matricSuction50/self.airEntryValueUpp000005)**\
                                                          (-1.0/self.poreSizeBetaUpp000005)
            self.effSatAt50Upp005030 = (self.matricSuction50/self.airEntryValueUpp005030)**\
                                                          (-1.0/self.poreSizeBetaUpp005030)
            self.effSatAt50Low030150 = (self.matricSuction50/self.airEntryValueLow030150)**\
                                                          (-1.0/self.poreSizeBetaLow030150)
            self.effSatAtWiltPointUpp000005 = \
                                       (self.matricSuctionWP/self.airEntryValueUpp000005)**\
                                                          (-1.0/self.poreSizeBetaUpp000005)
            self.effSatAtWiltPointUpp005030 = \
                                       (self.matricSuctionWP/self.airEntryValueUpp005030)**\
                                                          (-1.0/self.poreSizeBetaUpp005030)
            self.effSatAtWiltPointLow030150 = \
                                       (self.matricSuctionWP/self.airEntryValueLow030150)**\
                                                          (-1.0/self.poreSizeBetaLow030150)

        # calculate interflow parameter (TCL): 
        #
        if self.numberOfLayers == 2:
            self.interflowConcTime = (self.kSatLow * self.tanslope*2.0) / \
                     (self.slopeLength * (1.- self.effSatAtFieldCapLow) * \
                    (self.satVolMoistContLow - self.resVolMoistContLow))    # TCL = Duration*(2*KS2*TANSLOPE)/(LSLOPE*(1-THEFF2_FC)*(THETASAT2-THETARES2))
        #
        if self.numberOfLayers == 3:
            self.interflowConcTime = (self.kSatLow030150 * self.tanslope*2.0) / \
                     (self.slopeLength * (1.-self.effSatAtFieldCapLow030150) * \
             (self.satVolMoistContLow030150 - self.resVolMoistContLow030150))
        
        self.interflowConcTime = pcr.max(0.0, pcr.cover(self.interflowConcTime, 0.0))