# ******************************************************
## 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/N_model/trunk/tools/nretention.py $"
# ******************************************************

import os
import ascraster
import accuflux_retention
import accuflux

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

    # Calculate the total Nload (accuload) with retention that reaches the river mouth    
    accuNload,retention = accuflux_retention.N_accuflux(params,mask,lake_icell_dict,lddmap,Nload,hydraulic_load,\
                                water_body,discharge)

    accuNload.write_ascii_file(os.path.join(params.outputdir,"Nload.asc"))
    retention.write_ascii_file(os.path.join(params.outputdir,"retention.asc"))

    # Put the nload at the mouth of the river in the mouth database
    for key in mouth_dict.keys():
        mouth_dict[key].nload_mouth = accuNload.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].nconc_mouth = 1.0e-6 * mouth_dict[key].nload_mouth/mouth_dict[key].discharge_mouth
        else:
            mouth_dict[key].nconc_mouth = 0.0
            mouth_dict[key].nload_mouth = 0.0
        try:
            mouth_dict[key].retention = 1.0 - (mouth_dict[key].nload_mouth/mouth_dict[key].nload_river)
        except ZeroDivisionError:
            mouth_dict[key].retention = 0.0

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

        load = accuNload.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
                    nconc.set_data(icell,1.0e-6*load/flux)
                else:    
                    nconc.set_data(icell,0.0)
        else:
            nconc.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 = nconc.get_data(icell)
            for icell_lake in lake_icell_dict[icell].cells:
                nconc.set_data(icell_lake,conc)

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

        load = Nload.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
                    nconc.set_data(icell,1.0e-6 * (1.0 - reten) * load/flux)        
                else:    
                    nconc.set_data(icell,0.0)
        else:
            nconc.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 = nconc.get_data(icell)
            for icell_lake in lake_icell_dict[icell].cells:
                nconc.set_data(icell_lake,conc)

    nconc.write_ascii_file(os.path.join(params.outputdir,"Nconc_local.asc"))

    return  accuNload,retention

