# ******************************************************
## Revision "$LastChangedDate: 2014-01-26 12:25:00 +0100 (Sun, 26 Jan 2014) $"
## Date "$LastChangedRevision: 456 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/P_model/trunk/tools/pretention.py $"
# ******************************************************

import os
import ascraster
import accuflux_retention
import accuflux

def calculate(params,mask,mouth_dict,lake_icell_dict,lddmap,Pload,hydraulic_load,\
              water_body,pnet,discharge,residence_time_grid):
    '''
    This function calculates the retention of and reduces the river load by this.
    '''

    # Calculate the total Pload (accuload) with retention that reaches the river mouth    
    accuPload,retention = accuflux_retention.P_accuflux(params,mask,lake_icell_dict,lddmap,Pload,hydraulic_load,\
                                water_body,discharge,residence_time_grid)

    accuPload.write_ascii_file(os.path.join(params.outputdir,"Pload.asc"))
    retention.write_ascii_file(os.path.join(params.outputdir,"retention.asc"))

    # Put the Pload at the mouth of the river in the mouth database
    for key in mouth_dict.keys():
        mouth_dict[key].pload_mouth = accuPload.get_data(mouth_dict[key].index_grid,0.0)
        if (mouth_dict[key].discharge_mouth > params.minimal_discharge):
            # 1.0e-6 is used to convert from kg/km3 to g/m3 or mg/l
            mouth_dict[key].pconc_mouth = 1.0e-6 * mouth_dict[key].pload_mouth/mouth_dict[key].discharge_mouth
        else:
            mouth_dict[key].pconc_mouth = 0.0
            mouth_dict[key].pload_mouth = 0.0
        try:
            mouth_dict[key].retention = 1.0 - (mouth_dict[key].pload_mouth/mouth_dict[key].pload_river)
        except ZeroDivisionError:
            mouth_dict[key].retention = 0.0

    # Make two extra maps with concentrations
    # Concentration of accumulated flux
    pconc = ascraster.duplicategrid(accuPload)
    pconc.add_values(pconc.length*[pconc.nodata_value])
    for icell in range(accuPload.length):
        # Check whether this cell is an outflow of a lake/reservoir
        try:            
            loutflow = lake_icell_dict[icell].outcell
        except KeyError:
            loutflow = -1

        load = accuPload.get_data(icell,0.0)
        if (load > 0.0):
            if (loutflow == -1 or loutflow == 1):
                flux =  discharge.get_data(icell,0.0)
                # Here we use a constraint to avoid high concentrations in the case that there is almost no discharge.
                if (flux > params.minimal_discharge):
                    # 1.0e-6 is used to convert from kg/km3 to g/m3 or mg/l
                    pconc.set_data(icell,1.0e-6*load/flux)
                else:    
                    pconc.set_data(icell,0.0)
        else:
            pconc.set_data(icell,0.0)

        if (loutflow == 1):
            # This is a outflow point of a water body:
            # Set all the cells of this water body to the same concentration.
            conc = pconc.get_data(icell)
            for icell_lake in lake_icell_dict[icell].cells:
                pconc.set_data(icell_lake,conc)

    pconc.write_ascii_file(os.path.join(params.outputdir,"Pconc.asc"))      
    
    # Local concentration
    pconc = ascraster.duplicategrid(accuPload)
    pconc.add_values(pconc.length*[pconc.nodata_value])    
    for icell in range(Pload.length):
        # Check whether this cell is an outflow of a lake/reservoir
        try:            
            loutflow = lake_icell_dict[icell].outcell
        except KeyError:
            loutflow = -1

        load = Pload.get_data(icell,0.0)
        if (load > 0.0):
            if (loutflow == -1 or loutflow == 1):
                reten = retention.get_data(icell,0.0)
                flux =  pnet.get_data(icell,0.0)
                if (flux >= params.minimal_soil_infiltration):
                    # 1.0e-6 is used to convert from kg/km3 to g/m3 or mg/l
                    pconc.set_data(icell,1.0e-6 * (1.0 - reten) * load/flux)        
                else:    
                    pconc.set_data(icell,0.0)
        else:
            pconc.set_data(icell,0.0)

        if (loutflow == 1):
            # This is a outflow point of a water body:
            # Set all the cells of this water body to the same concentration.
            conc = pconc.get_data(icell)
            for icell_lake in lake_icell_dict[icell].cells:
                pconc.set_data(icell_lake,conc)

    pconc.write_ascii_file(os.path.join(params.outputdir,"Pconc_local.asc"))       

    return  accuPload,retention

