<reponame>amirsarikhani/NEST
#-------------------------------------------------------------------------
#Name:        Water Balance module
#Purpose:
#1.) check if water balance per time step is ok ( = 0)
#2.) produce an annual overview - income, outcome storage
#Author:      PB
#
#Created:     22/08/2016
#Copyright:   (c) PB 2016
#-------------------------------------------------------------------------

from management_modules.data_handling import *

class waterbalance(object):
    """
    WATER BALANCE
    
    * check if water balnace per time step is ok ( = 0)
    * produce an annual overview - income, outcome storage
    """

    def __init__(self, waterbalance_variable):
        self.var = waterbalance_variable

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------

    def initial(self):
        """
        Initial part of the water balance module
        """

        if checkOption('calcWaterBalance'):
            i = 1
            self.var.nonIrrReturnFlow = 0
            self.var.localQW = 0
            self.var.channelStorageBefore = 0

        """ store the initial storage volume of snow, soil etc.
        """
        if checkOption('sumWaterBalance'):
            # variables of storage
            self.var.sum_balanceStore = ['SnowCover','sum_interceptStor','sum_topWaterLayer']

            # variable of fluxes
            self.var.sum_balanceFlux = ['Precipitation','SnowMelt','Rain','sum_interceptEvap','actualET']

            #for variable in self.var.sum_balanceStore:
                # vars(self.var)["sumup_" + variable] =  vars(self.var)[variable]
            for variable in self.var.sum_balanceFlux:
                vars(self.var)["sumup_" + variable] =  globals.inZero.copy()
            i =1


# --------------------------------------------------------------------------
# --------------------------------------------------------------------------

    def waterBalanceCheck(self, fluxesIn, fluxesOut, preStorages, endStorages, processName, printTrue=False):
        """
        Dynamic part of the water balance module
        Returns the water balance for a list of input, output, and storage map files

        :param fluxesIn: income
        :param fluxesOut: this goes out
        :param preStorages:  this was in before
        :param endStorages:  this was in afterwards
        :param processName:  name of the process
        :param printTrue: calculate it?
        :return:
        """

        if printTrue:
            minB =0
            maxB = 0
            maxBB = 0
            income =  0
            out = 0
            store = 0

            for fluxIn in fluxesIn:   income += fluxIn
            for fluxOut in fluxesOut:
                out += fluxOut
            for preStorage in preStorages: store += preStorage
            for endStorage in endStorages: store -= endStorage
            balance =  income + store - out
            #balance = endStorages
            #if balance is not empty
            if balance.size:
                minB = np.amin(balance)
                maxB = np.amax(balance)
                maxBB = np.maximum(np.abs(minB),np.abs(maxB))
            #meanB = np.average(balance, axis=0)
            #meanB = 0.0

            print "     %s %10.8f " % (processName, maxBB),
            #if maxBB > 0.00000001:
            #    sys.exit()
            """
            if maxBB > 0.0000000001:
                print "     %s %10.8f %10.8f" % (processName, minB,maxB),
                if (minB < -0.00001) or (maxB > 0.00001):
                   i=11111
            """



    def waterBalanceCheckSum(self, fluxesIn, fluxesOut, preStorages, endStorages, processName, printTrue=False):
        """
        Returns the water balance for a list of input, output, and storage map files
        and sums it up for a catchment

        :param fluxesIn: income
        :param fluxesOut: this goes out
        :param preStorages:  this was in before
        :param endStorages:  this was in afterwards
        :param processName:  name of the process
        :param printTrue: calculate it?
        :return: Water balance as output on the screen
        """

        if printTrue:
            minB =0
            maxB = 0
            maxBB = 0

            income =  0
            out = 0
            store =  0

            for fluxIn in fluxesIn:
                if not(isinstance(fluxIn,np.ndarray)) : fluxIn = globals.inZero
                income = income + np.bincount(self.var.catchmentAll,weights=fluxIn)
            for fluxOut in fluxesOut:
                if not (isinstance(fluxOut, np.ndarray)): fluxOut = globals.inZero
                out = out + np.bincount(self.var.catchmentAll,weights=fluxOut)
            for preStorage in preStorages:
                if not (isinstance(preStorage, np.ndarray)): preStorage = globals.inZero
                store = store + np.bincount(self.var.catchmentAll,weights=preStorage)
            for endStorage in endStorages:
                if not (isinstance(endStorage, np.ndarray)): endStorage = globals.inZero
                store = store - np.bincount(self.var.catchmentAll,weights=endStorage)
            balance =  income + store - out
            #balance = endStorages
            #if balance.size:
                #minB = np.amin(balance)
                #maxB = np.amax(balance)
                #maxBB = np.maximum(np.abs(minB),np.abs(maxB))
                #meanB = np.average(balance, axis=0)
                #meanB = 0.0
            #no = self.var.catchmentNo
            no = 0


            #print "     %s %10.8f " % (processName, maxBB),
            #print "     %s %10.8f %10.8f" % (processName, minB,maxB),
            #print "     %s %10.8f" % (processName, balance[no]),
            print "     %s %10.8f" % (processName, balance[no]),

            #avgArea = npareaaverage(self.var.cellArea, self.var.catchmentAll)
            #dis = balance[no] * avgArea[0] / self.var.DtSec
            #print "     %s %10.8f" % (processName, dis),





# ----------------------------------------------------------------------------------------
# ----------------------------------------------------------------------------------------

    def checkWaterSoilGround(self):
        """
        Check water balance of snow, vegetation, soil, groundwater
        """

        if checkOption('calcWaterBalance'):


            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.sumirrConsumption],
                [self.var.sum_directRunoff, self.var.sum_interflow, self.var.sum_gwRecharge,self.var.totalET],  # Out
                [self.var.pretotalSto],  # prev storage
                [self.var.totalSto],
                "Soil_all", False)


            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.sumirrConsumption,self.var.sum_capRiseFromGW],
                [self.var.sum_directRunoff, self.var.sum_perc3toGW, self.var.sum_prefFlow, \
                 self.var.totalET],  # Out
                [self.var.pretotalSto],  # prev storage
                [self.var.totalSto],
                "Soil_all1", False)



            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.sumirrConsumption,self.var.sum_capRiseFromGW],
                [self.var.sum_directRunoff, self.var.sum_perc3toGW, self.var.sum_prefFlow, \
                 self.var.sum_actTransTotal, self.var.sum_actBareSoilEvap, self.var.sum_openWaterEvap, self.var.sum_interceptEvap, self.var.snowEvap],  # Out
                [self.var.pretotalSto],  # prev storage
                [self.var.totalSto],
                "Soil_all2", False)

            #if self.var.sumirrConsumption > 0:
            #    iii =1


            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.sumirrConsumption],
                [self.var.sum_directRunoff, self.var.sum_interflow, self.var.nonFossilGroundwaterAbs,self.var.baseflow, \
                 self.var.sum_actTransTotal, self.var.sum_actBareSoilEvap, self.var.sum_openWaterEvap, self.var.sum_interceptEvap, self.var.snowEvap],  # Out
                [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                [self.var.totalSto,self.var.storGroundwater],
                "Soil+G", False)

            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.unmetDemand,self.var.nonFossilGroundwaterAbs, self.var.actSurfaceWaterAbstract], # In
                [self.var.act_irrDemand,self.var.act_nonIrrDemand],     # Out
                [globals.inZero],
                [globals.inZero],
                "Waterdemand1", False)

            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.act_irrDemand],  # In
                [self.var.sumirrConsumption, self.var.returnflowIrr,self.var.addtoevapotrans],  # Out
                [globals.inZero],
                [globals.inZero],
                "Waterdemand2", False)



            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.Precipitation, self.var.act_irrDemand],
                [self.var.sum_directRunoff, self.var.sum_interflow, self.var.nonFossilGroundwaterAbs,self.var.baseflow, \
                 self.var.sum_actTransTotal, self.var.sum_actBareSoilEvap, self.var.sum_openWaterEvap, self.var.sum_interceptEvap, self.var.snowEvap, \
                 self.var.returnflowIrr, self.var.addtoevapotrans],  # Out
                [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                [self.var.totalSto,self.var.storGroundwater],
                "Soil+G+WD", False)

            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.act_irrDemand],
                [self.var.totalET, self.var.runoff, self.var.nonFossilGroundwaterAbs, self.var.returnflowIrr],  # Out
                [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                [self.var.totalSto,self.var.storGroundwater],
                "Soil+G+WD2", False)

            nonIrrReturn = self.var.nonIrrReturnFlowFraction * self.var.act_nonIrrDemand
            nonIrruse = self.var.act_nonIrrDemand - nonIrrReturn

            if checkOption('includeWaterBodies'):
                self.var.waterbalance_module.waterBalanceCheck(
                    [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                    [self.var.totalET, self.var.runoff, self.var.returnflowIrr, self.var.act_nonIrrDemand,  self.var.smallrunoffDiff ],  # Out
                    [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                    [self.var.totalSto,self.var.storGroundwater],
                    "Soil+G+WD31", True)

                self.var.waterbalance_module.waterBalanceCheck(
                    [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                    [self.var.totalET, self.var.runoff, self.var.returnFlow, nonIrruse, self.var.smallrunoffDiff],  # Out
                    [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                    [self.var.totalSto,self.var.storGroundwater ],
                    "Soil+G+WD32", True)


                if checkOption('calcWaterBalance'):
                    self.var.waterbalance_module.waterBalanceCheck(
                        [self.var.lakeResInflowM],  # In
                        [self.var.lakeResOutflowM, self.var.EvapWaterBodyM],  # Out  self.var.evapWaterBodyC
                        [self.var.prelakeResStorage / self.var.cellArea],  # prev storage
                        [self.var.lakeResStorage / self.var.cellArea],
                        "LR1", False)

                if checkOption('calcWaterBalance') and returnBool('useSmallLakes'):
                    self.var.waterbalance_module.waterBalanceCheck(
                        [self.var.smallLakeIn],  # In
                        [self.var.smallLakeout,  self.var.smallevapWaterBody],  # Out
                        [self.var.preSmalllakeStorage / self.var.cellArea],  # prev storage
                        [self.var.smalllakeStorage / self.var.cellArea],
                        "LR2", False)


                self.var.waterbalance_module.waterBalanceCheck(
                    [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                    [self.var.totalET, self.var.runoff, self.var.returnFlow, nonIrruse, self.var.lakeResOutflowM - self.var.lakeResInflowM, self.var.EvapWaterBodyM,
                     self.var.smallrunoffDiff],  # Out
                    [self.var.pretotalSto,self.var.prestorGroundwater, self.var.prelakeResStorage/ self.var.cellArea, ],  # prev storage
                    [self.var.totalSto,self.var.storGroundwater, self.var.lakeResStorage/ self.var.cellArea],
                    "Soil+G+WD+LR", False)

                if checkOption('calcWaterBalance') and returnBool('useSmallLakes'):
                    self.var.waterbalance_module.waterBalanceCheck(
                        [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                        [self.var.totalET, self.var.runoff, self.var.returnFlow, nonIrruse, self.var.lakeResOutflowM - self.var.lakeResInflowM ,
                         self.var.smallLakeout - self.var.smallLakeIn, self.var.EvapWaterBodyM,  self.var.smallevapWaterBody, self.var.smallrunoffDiff],  # Out
                        [self.var.pretotalSto,self.var.prestorGroundwater, self.var.prelakeResStorage/ self.var.cellArea, self.var.preSmalllakeStorage / self.var.cellArea],  # prev storage
                        [self.var.totalSto,self.var.storGroundwater, self.var.lakeResStorage/ self.var.cellArea, self.var.smalllakeStorage / self.var.cellArea ],
                        "Soil+G+WD+LR2", False)


                if checkOption('calcWaterBalance') and returnBool('useSmallLakes'):
                    self.var.waterbalance_module.waterBalanceCheck(
                        [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                        [self.var.totalET, self.var.runoff, self.var.returnFlow, nonIrruse,self.var.lakeResOutflowM - self.var.lakeResInflowM ,
                         self.var.EvapWaterBodyM,  self.var.smallevapWaterBody,  self.var.smallLakeDiff],  # Out
                        [self.var.pretotalSto,self.var.prestorGroundwater, self.var.prelakeResStorage/ self.var.cellArea, self.var.preSmalllakeStorage / self.var.cellArea],  # prev storage
                        [self.var.totalSto,self.var.storGroundwater, self.var.lakeResStorage/ self.var.cellArea, self.var.smalllakeStorage / self.var.cellArea ],
                        "Soil+G+WD+LR22", True)

#### IMPORTANT set Routingstep to 1 to test!

            DisOut = self.var.discharge * self.var.DtSec / self.var.cellArea
            DisOut = np.where(self.var.lddCompress == 5, DisOut, 0.)



            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.runoff, self.var.returnFlow],  # In
                [self.var.sumsideflow / self.var.cellArea, self.var.EvapoChannel / self.var.cellArea, self.var.actSurfaceWaterAbstract, ],  # Out
                [],  # prev storage
                [],
                "rout1", False)

            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.runoff, self.var.returnFlow],  # In
                [DisOut, self.var.EvapoChannel / self.var.cellArea, self.var.actSurfaceWaterAbstract],  # Out
                [self.var.prechannelStorage / self.var.cellArea],  # prev storage
                [self.var.channelStorage / self.var.cellArea],
                "rout4", False)

            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.Precipitation, self.var.unmetDemand],
                [self.var.totalET, self.var.EvapoChannel / self.var.cellArea, nonIrruse, DisOut ],  # Out
                [self.var.pretotalSto,self.var.prestorGroundwater, self.var.prechannelStorage / self.var.cellArea],  # prev storage
                [self.var.totalSto,self.var.storGroundwater, self.var.channelStorage / self.var.cellArea ],
                "Soil+G+WD+R", False)



            #print self.var.channelStorageBefore[10],
            # print "%10.8f %10.8f %10.8f " % (income[10], out[10], store[10]),

        ii =1

        # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++


        if checkOption('sumWaterBalance'):
            """
            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.nonFossilGroundwaterAbs, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],  # In
                [self.var.waterWithdrawal],  # Out
                [globals.inZero],
                [globals.inZero],
                "WaterdemandSum", False)

            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.Precipitation, self.var.surfaceWaterInf,self.var.sumIrrDemand,self.var.nonIrrReturnFlow],    # In
                [self.var.totalET,   self.var.localQW, self.var.riverbedExchange * self.var.InvCellArea, self.var.nonFossilGroundwaterAbs, self.var.actSurfaceWaterAbstract],
                [self.var.pretotalSto,self.var.prestorGroundwater,self.var.prechannelStorage * self.var.InvCellArea],       # prev storage
                [self.var.totalSto,self.var.storGroundwater,self.var.channelStorageBefore * self.var.InvCellArea],
                "S+G+Rout1Sum", False)
            """

            # discharge at outlets of catchments
            """
            DisOut = self.var.discharge * self.var.DtSec
            DisOut = np.where(self.var.outlets > 0, DisOut, 0.)

            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.Precipitation * self.var.cellArea, self.var.surfaceWaterInf * self.var.cellArea, self.var.sum_irrDemand * self.var.cellArea,self.var.nonIrrReturnFlow * self.var.cellArea],    # In
                [self.var.totalET  * self.var.cellArea,   self.var.localQW  * self.var.cellArea, self.var.riverbedExchange, self.var.nonFossilGroundwaterAbs * self.var.cellArea, self.var.sum_actSurfaceWaterAbstract * self.var.cellArea, DisOut],
                [self.var.prevSnowCover * self.var.cellArea, self.var.pretotalSoil * self.var.cellArea,self.var.prestorGroundwater * self.var.cellArea,self.var.prechannelStorage],       # prev storage
                [self.var.SnowCover * self.var.cellArea, self.var.totalSoil * self.var.cellArea,self.var.storGroundwater * self.var.cellArea,self.var.channelStorage],
                "S+G+Rout2Sum", False)
            """

        """
        if checkOption('budyko'):


            # in mm
            self.var.sumETRef =[]
            self.var.sumP = []
            self.var.sumETA = []
            for i in xrange(self.var.noOutpoints):
                # TODO change this to store in arrays not in maps for each gauge station!
                area = npareatotal(self.var.cellArea, self.var.evalCatch[i])
                self.var.sumETRef.append(1000 * npareatotal(self.var.ETRef * self.var.cellArea, self.var.evalCatch[i]) / area)
                self.var.sumP.append(1000 * npareatotal(self.var.Precipitation * self.var.cellArea, self.var.evalCatch[i]) / area)
                self.var.sumETA.append(1000 * npareatotal((self.var.totalET) * self.var.cellArea, self.var.evalCatch[i]) / area)

            ##self.var.area = npareatotal(self.var.cellArea, self.var.catchment)
            ##self.var.sumETRef = 1000 * npareatotal(self.var.ETRef * self.var.cellArea, self.var.catchment)/self.var.area
            ##self.var.sumP = 1000 * npareatotal(self.var.Precipitation * self.var.cellArea, self.var.catchment) / self.var.area
            #self.var.sumETA = 1000 * npareatotal((self.var.totalET + self.var.localQW) * self.var.cellArea , self.var.catchment)/self.var.area
            ##self.var.sumETA = 1000 * npareatotal((self.var.totalET) * self.var.cellArea, self.var.catchment) / self.var.area
            #self.var.sumRunoff = 1000 *  npareatotal((self.var.channelStorageBefore - self.var.prechannelStorage), self.var.catchment)/self.var.area
            #self.var.sumDelta1 = 1000 * npareatotal((self.var.SnowCover + self.var.totalSoil + self.var.storGroundwater) * self.var.cellArea , self.var.catchment)/self.var.area
            #self.var.sumDelta2 = 1000 * npareatotal((self.var.prevSnowCover + self.var.pretotalSoil + self.var.prestorGroundwater) * self.var.cellArea , self.var.catchment)/self.var.area
            #self.var.sumAll = self.var.sumDelta1- self.var.sumDelta2


            #report(decompress(self.var.sumRunoff), "C:\work\output3/sumRun.map")

            ii= 1
        """



    def dynamic(self):
        """
        Dynamic part of the water balance module
        If option **sumWaterBalance** sum water balance for certain variables
        """

        if checkOption('sumWaterBalance'):
            i = 1

        # sum up storage variables
        #for variable in self.var.sum_balanceStore:
         #   vars(self.var)["sumup_" + variable] =  vars(self.var)[variable].copy()


        # sum up fluxes variables
            for variable in self.var.sum_balanceFlux:
                vars(self.var)["sumup_" + variable] = vars(self.var)["sumup_" + variable] + vars(self.var)[variable]
#-------------------------------------------------------------------------
#Name:        Water Balance module
#Purpose:
#1.) check if water balance per time step is ok ( = 0)
#2.) produce an annual overview - income, outcome storage
#Author:      PB
#
#Created:     22/08/2016
#Copyright:   (c) PB 2016
#-------------------------------------------------------------------------

from management_modules.data_handling import *

class waterbalance(object):
    """
    WATER BALANCE
    
    * check if water balnace per time step is ok ( = 0)
    * produce an annual overview - income, outcome storage
    """

    def __init__(self, waterbalance_variable):
        self.var = waterbalance_variable

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------

    def initial(self):
        """
        Initial part of the water balance module
        """

        if checkOption('calcWaterBalance'):
            i = 1
            self.var.nonIrrReturnFlow = 0
            self.var.localQW = 0
            self.var.channelStorageBefore = 0

        """ store the initial storage volume of snow, soil etc.
        """
        if checkOption('sumWaterBalance'):
            # variables of storage
            self.var.sum_balanceStore = ['SnowCover','sum_interceptStor','sum_topWaterLayer']

            # variable of fluxes
            self.var.sum_balanceFlux = ['Precipitation','SnowMelt','Rain','sum_interceptEvap','actualET']

            #for variable in self.var.sum_balanceStore:
                # vars(self.var)["sumup_" + variable] =  vars(self.var)[variable]
            for variable in self.var.sum_balanceFlux:
                vars(self.var)["sumup_" + variable] =  globals.inZero.copy()
            i =1


# --------------------------------------------------------------------------
# --------------------------------------------------------------------------

    def waterBalanceCheck(self, fluxesIn, fluxesOut, preStorages, endStorages, processName, printTrue=False):
        """
        Dynamic part of the water balance module
        Returns the water balance for a list of input, output, and storage map files

        :param fluxesIn: income
        :param fluxesOut: this goes out
        :param preStorages:  this was in before
        :param endStorages:  this was in afterwards
        :param processName:  name of the process
        :param printTrue: calculate it?
        :return:
        """

        if printTrue:
            minB =0
            maxB = 0
            maxBB = 0
            income =  0
            out = 0
            store = 0

            for fluxIn in fluxesIn:   income += fluxIn
            for fluxOut in fluxesOut:
                out += fluxOut
            for preStorage in preStorages: store += preStorage
            for endStorage in endStorages: store -= endStorage
            balance =  income + store - out
            #balance = endStorages
            #if balance is not empty
            if balance.size:
                minB = np.amin(balance)
                maxB = np.amax(balance)
                maxBB = np.maximum(np.abs(minB),np.abs(maxB))
            #meanB = np.average(balance, axis=0)
            #meanB = 0.0

            print "     %s %10.8f " % (processName, maxBB),
            #if maxBB > 0.00000001:
            #    sys.exit()
            """
            if maxBB > 0.0000000001:
                print "     %s %10.8f %10.8f" % (processName, minB,maxB),
                if (minB < -0.00001) or (maxB > 0.00001):
                   i=11111
            """



    def waterBalanceCheckSum(self, fluxesIn, fluxesOut, preStorages, endStorages, processName, printTrue=False):
        """
        Returns the water balance for a list of input, output, and storage map files
        and sums it up for a catchment

        :param fluxesIn: income
        :param fluxesOut: this goes out
        :param preStorages:  this was in before
        :param endStorages:  this was in afterwards
        :param processName:  name of the process
        :param printTrue: calculate it?
        :return: Water balance as output on the screen
        """

        if printTrue:
            minB =0
            maxB = 0
            maxBB = 0

            income =  0
            out = 0
            store =  0

            for fluxIn in fluxesIn:
                if not(isinstance(fluxIn,np.ndarray)) : fluxIn = globals.inZero
                income = income + np.bincount(self.var.catchmentAll,weights=fluxIn)
            for fluxOut in fluxesOut:
                if not (isinstance(fluxOut, np.ndarray)): fluxOut = globals.inZero
                out = out + np.bincount(self.var.catchmentAll,weights=fluxOut)
            for preStorage in preStorages:
                if not (isinstance(preStorage, np.ndarray)): preStorage = globals.inZero
                store = store + np.bincount(self.var.catchmentAll,weights=preStorage)
            for endStorage in endStorages:
                if not (isinstance(endStorage, np.ndarray)): endStorage = globals.inZero
                store = store - np.bincount(self.var.catchmentAll,weights=endStorage)
            balance =  income + store - out
            #balance = endStorages
            #if balance.size:
                #minB = np.amin(balance)
                #maxB = np.amax(balance)
                #maxBB = np.maximum(np.abs(minB),np.abs(maxB))
                #meanB = np.average(balance, axis=0)
                #meanB = 0.0
            #no = self.var.catchmentNo
            no = 0


            #print "     %s %10.8f " % (processName, maxBB),
            #print "     %s %10.8f %10.8f" % (processName, minB,maxB),
            #print "     %s %10.8f" % (processName, balance[no]),
            print "     %s %10.8f" % (processName, balance[no]),

            #avgArea = npareaaverage(self.var.cellArea, self.var.catchmentAll)
            #dis = balance[no] * avgArea[0] / self.var.DtSec
            #print "     %s %10.8f" % (processName, dis),





# ----------------------------------------------------------------------------------------
# ----------------------------------------------------------------------------------------

    def checkWaterSoilGround(self):
        """
        Check water balance of snow, vegetation, soil, groundwater
        """

        if checkOption('calcWaterBalance'):


            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.sumirrConsumption],
                [self.var.sum_directRunoff, self.var.sum_interflow, self.var.sum_gwRecharge,self.var.totalET],  # Out
                [self.var.pretotalSto],  # prev storage
                [self.var.totalSto],
                "Soil_all", False)


            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.sumirrConsumption,self.var.sum_capRiseFromGW],
                [self.var.sum_directRunoff, self.var.sum_perc3toGW, self.var.sum_prefFlow, \
                 self.var.totalET],  # Out
                [self.var.pretotalSto],  # prev storage
                [self.var.totalSto],
                "Soil_all1", False)



            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.sumirrConsumption,self.var.sum_capRiseFromGW],
                [self.var.sum_directRunoff, self.var.sum_perc3toGW, self.var.sum_prefFlow, \
                 self.var.sum_actTransTotal, self.var.sum_actBareSoilEvap, self.var.sum_openWaterEvap, self.var.sum_interceptEvap, self.var.snowEvap],  # Out
                [self.var.pretotalSto],  # prev storage
                [self.var.totalSto],
                "Soil_all2", False)

            #if self.var.sumirrConsumption > 0:
            #    iii =1


            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.sumirrConsumption],
                [self.var.sum_directRunoff, self.var.sum_interflow, self.var.nonFossilGroundwaterAbs,self.var.baseflow, \
                 self.var.sum_actTransTotal, self.var.sum_actBareSoilEvap, self.var.sum_openWaterEvap, self.var.sum_interceptEvap, self.var.snowEvap],  # Out
                [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                [self.var.totalSto,self.var.storGroundwater],
                "Soil+G", False)

            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.unmetDemand,self.var.nonFossilGroundwaterAbs, self.var.actSurfaceWaterAbstract], # In
                [self.var.act_irrDemand,self.var.act_nonIrrDemand],     # Out
                [globals.inZero],
                [globals.inZero],
                "Waterdemand1", False)

            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.act_irrDemand],  # In
                [self.var.sumirrConsumption, self.var.returnflowIrr,self.var.addtoevapotrans],  # Out
                [globals.inZero],
                [globals.inZero],
                "Waterdemand2", False)



            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.Precipitation, self.var.act_irrDemand],
                [self.var.sum_directRunoff, self.var.sum_interflow, self.var.nonFossilGroundwaterAbs,self.var.baseflow, \
                 self.var.sum_actTransTotal, self.var.sum_actBareSoilEvap, self.var.sum_openWaterEvap, self.var.sum_interceptEvap, self.var.snowEvap, \
                 self.var.returnflowIrr, self.var.addtoevapotrans],  # Out
                [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                [self.var.totalSto,self.var.storGroundwater],
                "Soil+G+WD", False)

            self.var.waterbalance_module.waterBalanceCheck(
                [self.var.Precipitation, self.var.act_irrDemand],
                [self.var.totalET, self.var.runoff, self.var.nonFossilGroundwaterAbs, self.var.returnflowIrr],  # Out
                [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                [self.var.totalSto,self.var.storGroundwater],
                "Soil+G+WD2", False)

            nonIrrReturn = self.var.nonIrrReturnFlowFraction * self.var.act_nonIrrDemand
            nonIrruse = self.var.act_nonIrrDemand - nonIrrReturn

            if checkOption('includeWaterBodies'):
                self.var.waterbalance_module.waterBalanceCheck(
                    [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                    [self.var.totalET, self.var.runoff, self.var.returnflowIrr, self.var.act_nonIrrDemand,  self.var.smallrunoffDiff ],  # Out
                    [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                    [self.var.totalSto,self.var.storGroundwater],
                    "Soil+G+WD31", True)

                self.var.waterbalance_module.waterBalanceCheck(
                    [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                    [self.var.totalET, self.var.runoff, self.var.returnFlow, nonIrruse, self.var.smallrunoffDiff],  # Out
                    [self.var.pretotalSto,self.var.prestorGroundwater],  # prev storage
                    [self.var.totalSto,self.var.storGroundwater ],
                    "Soil+G+WD32", True)


                if checkOption('calcWaterBalance'):
                    self.var.waterbalance_module.waterBalanceCheck(
                        [self.var.lakeResInflowM],  # In
                        [self.var.lakeResOutflowM, self.var.EvapWaterBodyM],  # Out  self.var.evapWaterBodyC
                        [self.var.prelakeResStorage / self.var.cellArea],  # prev storage
                        [self.var.lakeResStorage / self.var.cellArea],
                        "LR1", False)

                if checkOption('calcWaterBalance') and returnBool('useSmallLakes'):
                    self.var.waterbalance_module.waterBalanceCheck(
                        [self.var.smallLakeIn],  # In
                        [self.var.smallLakeout,  self.var.smallevapWaterBody],  # Out
                        [self.var.preSmalllakeStorage / self.var.cellArea],  # prev storage
                        [self.var.smalllakeStorage / self.var.cellArea],
                        "LR2", False)


                self.var.waterbalance_module.waterBalanceCheck(
                    [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                    [self.var.totalET, self.var.runoff, self.var.returnFlow, nonIrruse, self.var.lakeResOutflowM - self.var.lakeResInflowM, self.var.EvapWaterBodyM,
                     self.var.smallrunoffDiff],  # Out
                    [self.var.pretotalSto,self.var.prestorGroundwater, self.var.prelakeResStorage/ self.var.cellArea, ],  # prev storage
                    [self.var.totalSto,self.var.storGroundwater, self.var.lakeResStorage/ self.var.cellArea],
                    "Soil+G+WD+LR", False)

                if checkOption('calcWaterBalance') and returnBool('useSmallLakes'):
                    self.var.waterbalance_module.waterBalanceCheck(
                        [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                        [self.var.totalET, self.var.runoff, self.var.returnFlow, nonIrruse, self.var.lakeResOutflowM - self.var.lakeResInflowM ,
                         self.var.smallLakeout - self.var.smallLakeIn, self.var.EvapWaterBodyM,  self.var.smallevapWaterBody, self.var.smallrunoffDiff],  # Out
                        [self.var.pretotalSto,self.var.prestorGroundwater, self.var.prelakeResStorage/ self.var.cellArea, self.var.preSmalllakeStorage / self.var.cellArea],  # prev storage
                        [self.var.totalSto,self.var.storGroundwater, self.var.lakeResStorage/ self.var.cellArea, self.var.smalllakeStorage / self.var.cellArea ],
                        "Soil+G+WD+LR2", False)


                if checkOption('calcWaterBalance') and returnBool('useSmallLakes'):
                    self.var.waterbalance_module.waterBalanceCheck(
                        [self.var.Precipitation, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],
                        [self.var.totalET, self.var.runoff, self.var.returnFlow, nonIrruse,self.var.lakeResOutflowM - self.var.lakeResInflowM ,
                         self.var.EvapWaterBodyM,  self.var.smallevapWaterBody,  self.var.smallLakeDiff],  # Out
                        [self.var.pretotalSto,self.var.prestorGroundwater, self.var.prelakeResStorage/ self.var.cellArea, self.var.preSmalllakeStorage / self.var.cellArea],  # prev storage
                        [self.var.totalSto,self.var.storGroundwater, self.var.lakeResStorage/ self.var.cellArea, self.var.smalllakeStorage / self.var.cellArea ],
                        "Soil+G+WD+LR22", True)

#### IMPORTANT set Routingstep to 1 to test!

            DisOut = self.var.discharge * self.var.DtSec / self.var.cellArea
            DisOut = np.where(self.var.lddCompress == 5, DisOut, 0.)



            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.runoff, self.var.returnFlow],  # In
                [self.var.sumsideflow / self.var.cellArea, self.var.EvapoChannel / self.var.cellArea, self.var.actSurfaceWaterAbstract, ],  # Out
                [],  # prev storage
                [],
                "rout1", False)

            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.runoff, self.var.returnFlow],  # In
                [DisOut, self.var.EvapoChannel / self.var.cellArea, self.var.actSurfaceWaterAbstract],  # Out
                [self.var.prechannelStorage / self.var.cellArea],  # prev storage
                [self.var.channelStorage / self.var.cellArea],
                "rout4", False)

            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.Precipitation, self.var.unmetDemand],
                [self.var.totalET, self.var.EvapoChannel / self.var.cellArea, nonIrruse, DisOut ],  # Out
                [self.var.pretotalSto,self.var.prestorGroundwater, self.var.prechannelStorage / self.var.cellArea],  # prev storage
                [self.var.totalSto,self.var.storGroundwater, self.var.channelStorage / self.var.cellArea ],
                "Soil+G+WD+R", False)



            #print self.var.channelStorageBefore[10],
            # print "%10.8f %10.8f %10.8f " % (income[10], out[10], store[10]),

        ii =1

        # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++


        if checkOption('sumWaterBalance'):
            """
            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.nonFossilGroundwaterAbs, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],  # In
                [self.var.waterWithdrawal],  # Out
                [globals.inZero],
                [globals.inZero],
                "WaterdemandSum", False)

            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.Precipitation, self.var.surfaceWaterInf,self.var.sumIrrDemand,self.var.nonIrrReturnFlow],    # In
                [self.var.totalET,   self.var.localQW, self.var.riverbedExchange * self.var.InvCellArea, self.var.nonFossilGroundwaterAbs, self.var.actSurfaceWaterAbstract],
                [self.var.pretotalSto,self.var.prestorGroundwater,self.var.prechannelStorage * self.var.InvCellArea],       # prev storage
                [self.var.totalSto,self.var.storGroundwater,self.var.channelStorageBefore * self.var.InvCellArea],
                "S+G+Rout1Sum", False)
            """

            # discharge at outlets of catchments
            """
            DisOut = self.var.discharge * self.var.DtSec
            DisOut = np.where(self.var.outlets > 0, DisOut, 0.)

            self.var.waterbalance_module.waterBalanceCheckSum(
                [self.var.Precipitation * self.var.cellArea, self.var.surfaceWaterInf * self.var.cellArea, self.var.sum_irrDemand * self.var.cellArea,self.var.nonIrrReturnFlow * self.var.cellArea],    # In
                [self.var.totalET  * self.var.cellArea,   self.var.localQW  * self.var.cellArea, self.var.riverbedExchange, self.var.nonFossilGroundwaterAbs * self.var.cellArea, self.var.sum_actSurfaceWaterAbstract * self.var.cellArea, DisOut],
                [self.var.prevSnowCover * self.var.cellArea, self.var.pretotalSoil * self.var.cellArea,self.var.prestorGroundwater * self.var.cellArea,self.var.prechannelStorage],       # prev storage
                [self.var.SnowCover * self.var.cellArea, self.var.totalSoil * self.var.cellArea,self.var.storGroundwater * self.var.cellArea,self.var.channelStorage],
                "S+G+Rout2Sum", False)
            """

        """
        if checkOption('budyko'):


            # in mm
            self.var.sumETRef =[]
            self.var.sumP = []
            self.var.sumETA = []
            for i in xrange(self.var.noOutpoints):
                # TODO change this to store in arrays not in maps for each gauge station!
                area = npareatotal(self.var.cellArea, self.var.evalCatch[i])
                self.var.sumETRef.append(1000 * npareatotal(self.var.ETRef * self.var.cellArea, self.var.evalCatch[i]) / area)
                self.var.sumP.append(1000 * npareatotal(self.var.Precipitation * self.var.cellArea, self.var.evalCatch[i]) / area)
                self.var.sumETA.append(1000 * npareatotal((self.var.totalET) * self.var.cellArea, self.var.evalCatch[i]) / area)

            ##self.var.area = npareatotal(self.var.cellArea, self.var.catchment)
            ##self.var.sumETRef = 1000 * npareatotal(self.var.ETRef * self.var.cellArea, self.var.catchment)/self.var.area
            ##self.var.sumP = 1000 * npareatotal(self.var.Precipitation * self.var.cellArea, self.var.catchment) / self.var.area
            #self.var.sumETA = 1000 * npareatotal((self.var.totalET + self.var.localQW) * self.var.cellArea , self.var.catchment)/self.var.area
            ##self.var.sumETA = 1000 * npareatotal((self.var.totalET) * self.var.cellArea, self.var.catchment) / self.var.area
            #self.var.sumRunoff = 1000 *  npareatotal((self.var.channelStorageBefore - self.var.prechannelStorage), self.var.catchment)/self.var.area
            #self.var.sumDelta1 = 1000 * npareatotal((self.var.SnowCover + self.var.totalSoil + self.var.storGroundwater) * self.var.cellArea , self.var.catchment)/self.var.area
            #self.var.sumDelta2 = 1000 * npareatotal((self.var.prevSnowCover + self.var.pretotalSoil + self.var.prestorGroundwater) * self.var.cellArea , self.var.catchment)/self.var.area
            #self.var.sumAll = self.var.sumDelta1- self.var.sumDelta2


            #report(decompress(self.var.sumRunoff), "C:\work\output3/sumRun.map")

            ii= 1
        """



    def dynamic(self):
        """
        Dynamic part of the water balance module
        If option **sumWaterBalance** sum water balance for certain variables
        """

        if checkOption('sumWaterBalance'):
            i = 1

        # sum up storage variables
        #for variable in self.var.sum_balanceStore:
         #   vars(self.var)["sumup_" + variable] =  vars(self.var)[variable].copy()


        # sum up fluxes variables
            for variable in self.var.sum_balanceFlux:
                vars(self.var)["sumup_" + variable] = vars(self.var)["sumup_" + variable] + vars(self.var)[variable]
#------------------------------------------------------------------------- #Name:        Water Balance module #Purpose: #1.) check if water balance per time step is ok ( = 0) #2.) produce an annual overview - income, outcome storage #Author:      PB # #Created:     22/08/2016 #Copyright:   (c) PB 2016 #------------------------------------------------------------------------- WATER BALANCE          * check if water balnace per time step is ok ( = 0)     * produce an annual overview - income, outcome storage # -------------------------------------------------------------------------- # -------------------------------------------------------------------------- Initial part of the water balance module store the initial storage volume of snow, soil etc. # variables of storage # variable of fluxes #for variable in self.var.sum_balanceStore: # vars(self.var)["sumup_" + variable] =  vars(self.var)[variable] # -------------------------------------------------------------------------- # -------------------------------------------------------------------------- Dynamic part of the water balance module         Returns the water balance for a list of input, output, and storage map files          :param fluxesIn: income         :param fluxesOut: this goes out         :param preStorages:  this was in before         :param endStorages:  this was in afterwards         :param processName:  name of the process         :param printTrue: calculate it?         :return: #balance = endStorages #if balance is not empty #meanB = np.average(balance, axis=0) #meanB = 0.0 #if maxBB > 0.00000001: #    sys.exit() if maxBB > 0.0000000001:                 print "     %s %10.8f %10.8f" % (processName, minB,maxB),                 if (minB < -0.00001) or (maxB > 0.00001):                    i=11111 Returns the water balance for a list of input, output, and storage map files         and sums it up for a catchment          :param fluxesIn: income         :param fluxesOut: this goes out         :param preStorages:  this was in before         :param endStorages:  this was in afterwards         :param processName:  name of the process         :param printTrue: calculate it?         :return: Water balance as output on the screen #balance = endStorages #if balance.size: #minB = np.amin(balance) #maxB = np.amax(balance) #maxBB = np.maximum(np.abs(minB),np.abs(maxB)) #meanB = np.average(balance, axis=0) #meanB = 0.0 #no = self.var.catchmentNo #print "     %s %10.8f " % (processName, maxBB), #print "     %s %10.8f %10.8f" % (processName, minB,maxB), #print "     %s %10.8f" % (processName, balance[no]), #avgArea = npareaaverage(self.var.cellArea, self.var.catchmentAll) #dis = balance[no] * avgArea[0] / self.var.DtSec #print "     %s %10.8f" % (processName, dis), # ---------------------------------------------------------------------------------------- # ---------------------------------------------------------------------------------------- Check water balance of snow, vegetation, soil, groundwater # Out # prev storage # Out # prev storage # Out # prev storage #if self.var.sumirrConsumption > 0: #    iii =1 # Out # prev storage # In # Out # In # Out # Out # prev storage # Out # prev storage # Out # prev storage # Out # prev storage # In # Out  self.var.evapWaterBodyC # prev storage # In # Out # prev storage # Out # prev storage # Out # prev storage # Out # prev storage #### IMPORTANT set Routingstep to 1 to test! # In # Out # prev storage # In # Out # prev storage # Out # prev storage #print self.var.channelStorageBefore[10], # print "%10.8f %10.8f %10.8f " % (income[10], out[10], store[10]), # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ self.var.waterbalance_module.waterBalanceCheckSum(                 [self.var.nonFossilGroundwaterAbs, self.var.unmetDemand, self.var.actSurfaceWaterAbstract],  # In                 [self.var.waterWithdrawal],  # Out                 [globals.inZero],                 [globals.inZero],                 "WaterdemandSum", False)              self.var.waterbalance_module.waterBalanceCheckSum(                 [self.var.Precipitation, self.var.surfaceWaterInf,self.var.sumIrrDemand,self.var.nonIrrReturnFlow],    # In                 [self.var.totalET,   self.var.localQW, self.var.riverbedExchange * self.var.InvCellArea, self.var.nonFossilGroundwaterAbs, self.var.actSurfaceWaterAbstract],                 [self.var.pretotalSto,self.var.prestorGroundwater,self.var.prechannelStorage * self.var.InvCellArea],       # prev storage                 [self.var.totalSto,self.var.storGroundwater,self.var.channelStorageBefore * self.var.InvCellArea],                 "S+G+Rout1Sum", False) # discharge at outlets of catchments DisOut = self.var.discharge * self.var.DtSec             DisOut = np.where(self.var.outlets > 0, DisOut, 0.)              self.var.waterbalance_module.waterBalanceCheckSum(                 [self.var.Precipitation * self.var.cellArea, self.var.surfaceWaterInf * self.var.cellArea, self.var.sum_irrDemand * self.var.cellArea,self.var.nonIrrReturnFlow * self.var.cellArea],    # In                 [self.var.totalET  * self.var.cellArea,   self.var.localQW  * self.var.cellArea, self.var.riverbedExchange, self.var.nonFossilGroundwaterAbs * self.var.cellArea, self.var.sum_actSurfaceWaterAbstract * self.var.cellArea, DisOut],                 [self.var.prevSnowCover * self.var.cellArea, self.var.pretotalSoil * self.var.cellArea,self.var.prestorGroundwater * self.var.cellArea,self.var.prechannelStorage],       # prev storage                 [self.var.SnowCover * self.var.cellArea, self.var.totalSoil * self.var.cellArea,self.var.storGroundwater * self.var.cellArea,self.var.channelStorage],                 "S+G+Rout2Sum", False) if checkOption('budyko'):               # in mm             self.var.sumETRef =[]             self.var.sumP = []             self.var.sumETA = []             for i in xrange(self.var.noOutpoints):                 # TODO change this to store in arrays not in maps for each gauge station!                 area = npareatotal(self.var.cellArea, self.var.evalCatch[i])                 self.var.sumETRef.append(1000 * npareatotal(self.var.ETRef * self.var.cellArea, self.var.evalCatch[i]) / area)                 self.var.sumP.append(1000 * npareatotal(self.var.Precipitation * self.var.cellArea, self.var.evalCatch[i]) / area)                 self.var.sumETA.append(1000 * npareatotal((self.var.totalET) * self.var.cellArea, self.var.evalCatch[i]) / area)              ##self.var.area = npareatotal(self.var.cellArea, self.var.catchment)             ##self.var.sumETRef = 1000 * npareatotal(self.var.ETRef * self.var.cellArea, self.var.catchment)/self.var.area             ##self.var.sumP = 1000 * npareatotal(self.var.Precipitation * self.var.cellArea, self.var.catchment) / self.var.area             #self.var.sumETA = 1000 * npareatotal((self.var.totalET + self.var.localQW) * self.var.cellArea , self.var.catchment)/self.var.area             ##self.var.sumETA = 1000 * npareatotal((self.var.totalET) * self.var.cellArea, self.var.catchment) / self.var.area             #self.var.sumRunoff = 1000 *  npareatotal((self.var.channelStorageBefore - self.var.prechannelStorage), self.var.catchment)/self.var.area             #self.var.sumDelta1 = 1000 * npareatotal((self.var.SnowCover + self.var.totalSoil + self.var.storGroundwater) * self.var.cellArea , self.var.catchment)/self.var.area             #self.var.sumDelta2 = 1000 * npareatotal((self.var.prevSnowCover + self.var.pretotalSoil + self.var.prestorGroundwater) * self.var.cellArea , self.var.catchment)/self.var.area             #self.var.sumAll = self.var.sumDelta1- self.var.sumDelta2               #report(decompress(self.var.sumRunoff), "C:\work\output3/sumRun.map")              ii= 1 Dynamic part of the water balance module         If option **sumWaterBalance** sum water balance for certain variables # sum up storage variables #for variable in self.var.sum_balanceStore: #   vars(self.var)["sumup_" + variable] =  vars(self.var)[variable].copy() # sum up fluxes variables