# ******************************************************
## Revision "$LastChangedDate: 2014-01-10 13:27:58 +0100 (Fri, 10 Jan 2014) $"
## Date "$LastChangedRevision: 444 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/P_model/trunk/tools/accuflux_retention.py $"
# ******************************************************

import os
import sys
import math

import ascraster
import accuflux

def temperature_factor(temp):
    '''
    Relation between temperature and retention (in rivers).
    With a temperature of 20.0 degrees Celcius, 1.0 is returned.
    The base number 1.06 (Marce and Armengol, 2009)
    We take air temperature as an approximation for water temperature.
    So air temperature has only not negative value
    '''
    return (1.06)**(max(0.0,temp)-20.0)

def retention_lake(params,tau):

    # Use model 4 from Brett and Benjamin (2008)
    # Tpout = Tpin/(1.0 + k4*(tau**x4)). Then retention_lake = (Tpin - Tpout)/Tpin = 1 - Tpout/Tpin = 1 - (1/(1.0 + k4*(tau**x4))
    # retention_lake = 1.0 - (1.0/(1.0 + k4*(tau**x4)))
    # with tau the residence time of the water in the lake, and k4 and x4 
    # calibrated factors.
    try:
        return 1.0 - (1.0/(1.0 + params.k4*(tau**params.x4))) 
    except ZeroDivisionError:
        return 0.0

def calculate_retention(params,P_vf,hydraulic_load,water_body,\
                        temperature,discharge=0.0,load=0.0):
    '''
    Argument discharge and load are not used here, but are needed to call this function from the generalcode 
    module subgrid_retention.
    '''

    # Calculate retention for a cell.
    # Retention is 1.0 - exp(-vf/Hl)
    if (water_body > 0.0):
        if (params.luse_temperature):
            temp_fac = temperature_factor(temperature)
        else:
            temp_fac = 1.0

        vf_temp_PO4 = P_vf * temp_fac
        try:
            # print "ret: ",vf_temp_PO4,hydraulic_load,(1.0 - math.exp(- vf_temp_PO4 / hydraulic_load))
            return (1.0 - math.exp(- vf_temp_PO4 / hydraulic_load))
        except ZeroDivisionError:
            # when hydraulic load is zero, retention of the water body is equal to one.
            return 1.0
        except OverflowError:
            print  - vf_temp_PO4,hydraulic_load,P_vf ,temp_fac
            raise OverflowError
    else:
        return 0.0

            
def P_accuflux(params,mask,lake_icell_dict,lddmap,Pload,hydraulic_load,\
               water_body,discharge,residence_time_grid):

    '''
    ACCUFLUX
    Based on Gerard van Drecht, RIVM, Bilthoven, February 2002
    April 2002: accu fraction flux and accu fraction state added
    
    PURPOSE:
    Accumulate flow of mass, given by a material map [grid],
    along a flowpath, given by a local drain direction(LDD) map [grid].
    The output map contains accumulated flux [grid].
    All files are ascii-grid formatted.
    
    Lddmap: network of local drain directions
    flow_direction can be coded as follows:
    UNH/GRDC                       PCraster(LDD)
    32 64 128   meaning: NW N NE       7 8 9
    16  -   1             W -  E       4 5 6
     8  4   2            SW S SE       1 2 3
    We use the PCraster LDD-format; negative and zero values are assumed to be nodata.
    '''
    # Make output rasters
    accuPload = ascraster.duplicategrid(Pload)
    retention = ascraster.duplicategrid(Pload)

    # Make nodata on all cells.
    retention.add_values(retention.length*[retention.nodata_value])

    # Make order of all the cells.
    ordermap = accuflux.determine_order(lddmap)

    # Make a pointer for the calculation order
    pointer1 = []
    for icell in xrange(ordermap.length):
        pointer1.append([int(ordermap.get_data(icell)),icell])
    # Sort the pointer array so that lowest order is first calculated.
    pointer1.sort()

    # Read water temperature, flooding fraction and water_fraction
    temperature = ascraster.Asciigrid(ascii_file=params.water_temperature,\
                                      mask=mask,numtype=float) 
    flooding_fraction = ascraster.Asciigrid(ascii_file=params.flooding_fraction,\
                                    mask=mask,numtype=float)
    fraction_water_grid = ascraster.Asciigrid(ascii_file=params.fraction_water,\
                                  mask=mask,numtype=float) 

    # Loop over all cells:    
    for item in range(len(pointer1)):
        icell = pointer1[item][1]
        # Get mass flux at the cell.
        Pload_cell = accuPload.get_data(icell)
        if (Pload_cell == None):
            # No data value
            continue
        if Pload_cell < 0.0:
            print "Negative mass flow of " + str(Pload_cell) + " at cell: " + str(icell)           
            # Go to next cell.
            continue
        
        # Retention in cell
        water_body_cell = water_body.get_data(icell,0)
        if (water_body_cell > 0):
            hydraulic_load_cell = hydraulic_load.get_data(icell,0.0)
            # Nodata value for temperature is 20 degrees Celcius, because then the temperature does not influence the retention
            temperature_cell = temperature.get_data(icell,20.0)
            discharge_cell = discharge.get_data(icell,0.0)
            # Check whether this cell is an outflow of a lake/reservoir
            try:            
                loutflow = lake_icell_dict[icell].outcell
                if (loutflow == 1): setattr(lake_icell_dict[icell],"load_in",Pload_cell)
            except KeyError:
                # No lake or reservoir
                loutflow = -1
            try:
                if (loutflow == -1):
                    # Normal cell, no lake or reservoir.
                    ret_cell = calculate_retention(params,params.P_vf_river,hydraulic_load_cell,\
                                                   water_body_cell,\
                                                   temperature_cell)

                elif (loutflow == 0):
                    # Lake or reservoir, but no outflow point.
                    ret_cell = 0.0
                else:
                    # Cell is outflow of lake or reservoir.
                    # Recalculate the fraction water and temperature
                    frac = 0.0
                    temp = 0.0
                    number_of_cells = len(lake_icell_dict[icell].cells)

                    for item in range(number_of_cells):
                        icell_lake = lake_icell_dict[icell].cells[item]
                        frac      += fraction_water_grid.get_data(icell_lake,0.0)
                        temp      += temperature.get_data(icell_lake,0.0)
                    # Take average over the cells.
                    fraction_cell = frac/float(number_of_cells)
                    if (fraction_cell > 1.0): fraction_cell = 1.0
                    temperature_cell = temp/float(number_of_cells)

                    if (params.lake_retention == 1):
                        # This is a lake/reservoir and and an outflow cell.
                        # Calculation of retention with formula of Brett.
                        residence_time_cell = lake_icell_dict[icell].residence_time
                        ret_cell = retention_lake(params,residence_time_cell)
                    else:
                        if (lake_icell_dict[icell].id < 10000):
                            # It is a reservoir
                            P_vf = params.P_vf_reservoir
                        else:
                            # It is a reservoir
                            P_vf = params.P_vf_lake
                        ret_cell = calculate_retention(params,P_vf,hydraulic_load_cell,\
                                           water_body_cell,\
                                           temperature_cell)

                    if (lake_icell_dict[icell].lendolake == 1): ret_cell = 1.0
                    setattr(lake_icell_dict[icell],"retention",ret_cell)
                    setattr(lake_icell_dict[icell],"fraction_water",fraction_cell)
                    setattr(lake_icell_dict[icell],"temperature",temperature_cell)

            except OverflowError:
                print hydraulic_load_cell,water_body_cell,temperature_cell,discharge_cell,Pload_cell
                raise OverflowError
        else:
            ret_cell = 0.0


        # Fill retention grid with retention value
        retention.set_data(icell,ret_cell)
        Pload_cell = (1.0 - ret_cell) * Pload_cell
        if (loutflow == 1): setattr(lake_icell_dict[icell],"load_out",Pload_cell)
        accuPload.set_data(icell,Pload_cell)

        # Get direction of the ldd 
        direction = int(lddmap.get_data(icell,-1))
        if (direction < 1 or direction > 9 or direction == 5 ):
            pass
        else:
            icell_loop = accuflux.goto_nextcell(icell,direction,lddmap.ncols,mask=lddmap.mask)
            if (icell_loop > -1):
                prev = accuPload.get_data(icell_loop,0.0)
                accuPload.set_data(icell_loop,prev + Pload_cell)

    return accuPload,retention
