# ******************************************************
## Revision "$LastChangedDate: 2014-01-03 14:08:37 +0100 (Fri, 03 Jan 2014) $"
## Date "$LastChangedRevision: 439 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/P_model/trunk/tools/p_wetlands.py $"
# ******************************************************

import os
import math
import ascraster
import aggregate
import accuflux_retention
 
def calculate(params,mask,mouth_dict,basin):
    '''
    Calculate the P load from wetlands.
    The only P input to wetlands is allochtoneous material. We use as an estimation the gnpp for this.
    The output of the wetlands are calculated with the spiraling method.
    '''
    if params.ldebug: print "Start calculation of wetland loads."

    wetland_fraction = ascraster.Asciigrid(ascii_file=params.wetland_fraction,\
                                           mask=mask,numtype=float)
    gnpp             = ascraster.Asciigrid(ascii_file=params.gnpp,\
                                           mask=mask,numtype=float)
    cellarea         = ascraster.Asciigrid(ascii_file=params.cellarea,\
                                           mask=mask,numtype=float)
    temperature      = ascraster.Asciigrid(ascii_file=params.water_temperature,\
                                           mask=mask,numtype=float)
    wetland_runoff   = ascraster.Asciigrid(ascii_file=params.wetland_runoff,\
                                           mask=mask,numtype=float)  

    # Create new P grid with P load of gnpp falling in the wetlands
    P_gnpp_in = ascraster.duplicategrid(cellarea)

    # Make all cells zero.
    P_gnpp_in.nodata_value = -1
    P_gnpp_in.add_values(P_gnpp_in.length*[0.0])
    P_gnpp_out = ascraster.duplicategrid(P_gnpp_in)
    wetland_area = ascraster.duplicategrid(P_gnpp_in)

    P_vf = params.P_vf_wetlands

    for icell in xrange(wetland_fraction.length):
        wetland_fraction_cell = wetland_fraction.get_data(icell,0.0)
        if (wetland_fraction_cell > 0.0):
            gnpp_cell = gnpp.get_data(icell,0.0)
            area_cell = cellarea.get_data(icell,0.0)
            # Calculate the amount of C in the cell (NPP = gC /m2)
            # Area_cell is in m2. So 0.001 is from gr to kg.
            totC_cell = gnpp_cell * area_cell * 0.001
            # Calculate the amount of P in the wetland.
            totP_cell = params.fr_gnpp_to_surface_water * totC_cell * min(wetland_fraction_cell,1.0) / params.CP_ratio_gnpp
            P_gnpp_in.set_data(icell,totP_cell)
            # Calculate the hydraulic load of the wetland (depth in m, runoff in mm/yr, so residence time in yr)
            # Hl = depth/residence_time = depth/(volume/discharge) = depth*discharge/volume = depth*discharge/depth*area =
            # discharge/area = runoff in m/yr! So Hl = 0.001 * wetland_runoff to convert from mm/yr to m/yr.
            hydraulic_load = 0.001 * wetland_runoff.get_data(icell,0.0)
    
            # Calculate retention for the wetlands in a cell.
            # Retention is 1.0 - exp(-vf/Hl)
            temp = temperature.get_data(icell, 20.0)
            # We set waterbody to 1,0 in the argument list of the retention, because we check here whether it is a water body.
            retention = accuflux_retention.calculate_retention(params,P_vf,hydraulic_load,1.0,temp)
            P_gnpp_out.set_data(icell,totP_cell*(1.0-retention))
            # Calculate the wetland area in km2
            area_wetland = area_cell * min(wetland_fraction_cell,1.0) * 1.0e-6
            wetland_area.set_data(icell,area_wetland)

    # Aggregate rasters
    aggregate.aggregate_grid(basin,wetland_area,mouth_dict,item="wetland_area")
    aggregate.aggregate_grid(basin,P_gnpp_in,mouth_dict,item="p_wetland_in")
    aggregate.aggregate_grid(basin,P_gnpp_out,mouth_dict,item="p_wetland_out")

    # Write to output file
    P_gnpp_in.write_ascii_file(os.path.join(params.outputdir,"P_wetland_in.asc"))
    P_gnpp_out.write_ascii_file(os.path.join(params.outputdir,"P_wetland_out.asc"))

    return P_gnpp_out
