# ******************************************************
## Revision "$LastChangedDate: 2014-01-09 23:55:12 +0100 (Thu, 09 Jan 2014) $"
## Date "$LastChangedRevision: 442 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/N_model/trunk/tools/accuflux_retention.py $"
# ******************************************************

import os
import sys
import math

import ascraster
import interpolate
import accuflux

icold = [0,-1,0,1,-1,0,1,-1,0,1]
irowd = [0,1,1,1,0,0,0,-1,-1,-1]

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.0717734625 is equal to power(2;0.1). So doubling in 10 degrees.
    To account for air temperature / water temperature differences, we assume that temperature is positive.
    '''
    return (1.0717734625)**(max(0.0,temp)-20.0)

def N_conc(params,discharge,Nload):

    try:
        # Nload is in kg and discharge is in km3/yr.
        conc = 1.0e-6*Nload/discharge
        # Function is:
        # below 0.01 mg N/l then factor = 3.0
        # above 10. mg N/l then factor = 0.1
        # Between these values interpolation is used.
        return interpolate.interpolate(conc,params.nconc_list,params.reductionfactor_list,extrapol=1)
    except ZeroDivisionError:
        return 0.0    


def set_conc_relation_retention(params):
    '''
    To calculate the retention reduction due to the high or low NO3 concentration
    This function puts the given information from the parameters input file into one structure.
    All parameters of params.nconc_list* and params.reductionfactor_list* (with * a integer number) are join together in a list.
    The joined information is put back in the params object.
    This is the same as for example:
    params.nconc_list = map(float,[params.nconc_list1,params.nconc_list2,params.nconc_list3])
    params.reductionfactor_list = map(float,[params.reductionfactor_list1,params.reductionfactor_list2,params.reductionfactor_list3])
    '''
    if (params.luse_concentration):
        try:
            # Check whether this parameter is already set. 
            # When it is set, then do nothing.
            qq = params.nconc_list
        except AttributeError:
            counter = 1
            params.nconc_list = []
            params.reductionfactor_list = []
            lfound = True
            while (lfound):
                try:
                    xval = float(getattr(params,"nconc_list"+str(counter)))
                    yval = float(getattr(params,"reductionfactor_list"+str(counter)))
                    params.nconc_list.append(xval)
                    params.reductionfactor_list.append(yval)
                    print "nconc_list"+str(counter)+" and "+"reductionfactor_list"+str(counter)+" were found."
                    counter += 1
                except AttributeError:
                    lfound = False
                     
            # Check whether there is information found.
            if (params.nconc_list == []):
                raise MyError("There is no information found like nconc_list1 and reductionfactor_list1 in the parameters file.")
            elif (len(params.nconc_list)  != len(params.reductionfactor_list)):
                raise MyError("Information to compile nconc_list and reductionfactor_list in the parameters file is not correct.",\
                              "Number of entries for nconc_list = "+str(len(params.nconc_list)),\
                              "Number of entries for reductionfactor_list = "+str(len(params.reductionfactor_list)))


def calculate_retention(params,vf,hydraulic_load,water_body,\
                        temperature,discharge,Nload):

    # 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
        if (params.luse_concentration):
            conc_fac = N_conc(params,discharge,Nload)
        else:
            conc_fac = 1.0

        vf_temp_no3 = vf * temp_fac * conc_fac
        try:
            return (1.0 - math.exp(- vf_temp_no3 / 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_no3,hydraulic_load,params.vf ,temp_fac,conc_fac
            raise OverflowError
    else:
        return 0.0

            
def N_accuflux(params,mask,lake_icell_dict,lddmap,Nload,hydraulic_load,\
               water_body,discharge):

    '''
    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
    '''

    #retention_output_per_step = 1
    # Set the relation between concentration and retention.
    set_conc_relation_retention(params)

    # Make output rasters
    accuNload = ascraster.duplicategrid(Nload)
    retention = ascraster.duplicategrid(Nload)

    # 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 and water_fraction
    temperature = ascraster.Asciigrid(ascii_file=params.water_temperature,\
                                      mask=mask,numtype=float) 

    fraction_water_grid = ascraster.Asciigrid(ascii_file=params.fraction_water,\
                                  mask=mask,numtype=float)  

    # Make preparations for saving the retention process.
    #order_number = 0
    #order_number_prev = 0
    #if not os.path.isdir(os.path.join(params.outputdir,"rivload")):
    #    os.mkdir(os.path.join(params.outputdir,"rivload"))

    # Loop over all cells:    
    for item in range(len(pointer1)):
        icell = pointer1[item][1]
        # Get mass flux at the cell.
        Nload_cell = accuNload.get_data(icell)
        if (Nload_cell == None):
            # No data value
            continue
        if Nload_cell < 0.0:
            print "Negative mass flow of " + str(Nload_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",Nload_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.vf_river,hydraulic_load_cell,\
                                  water_body_cell,\
                                  temperature_cell,discharge_cell,Nload_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 (lake_icell_dict[icell].id < 10000):
                        # It is a reservoir
                        vf = params.vf_reservoir
                    else:
                        # It is a reservoir
                        vf = params.vf_lake
                    ret_cell = calculate_retention(params,vf,hydraulic_load_cell,\
                                  water_body_cell,\
                                  temperature_cell,discharge_cell,Nload_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,Nload_cell
                raise OverflowError
        else:
            ret_cell = 0.0


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

        # Store this retention step into the output file
        #if (retention_output_per_step == 1):
        #    order_number = pointer1[item][0]
        #    if (order_number != order_number_prev):
        #        accuNload.write_ascii_file(os.path.join(os.path.join(params.outputdir,"rivload"),"Nload_"+str(order_number)+".asc"))
        #        order_number_prev = order_number
        #        print order_number
            

        # 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 = accuNload.get_data(icell_loop,0.0)
                accuNload.set_data(icell_loop,prev + Nload_cell)

    return accuNload,retention
