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

import os
import ascraster
import aggregate
 
def calculate(params,mask,mouth_dict,basin):
    '''
    Calculates the direct nitrogen deposition on lakes and reservoirs. 
    '''

    # Read the fraction of the cell which is an waterbody
    fraction_water_grid = ascraster.Asciigrid(ascii_file=params.fraction_water,\
                                  mask=mask,numtype=float)

    # Make temporary N depostion file
    #print "*****************************************************************"
    #print "Fake depostion is used"
    # Read deposition in mg N/m2/yr
    ndep_grid = ascraster.Asciigrid(ascii_file=params.N_deposition,\
                                    mask=mask,numtype=float)
    #ndep_grid = ascraster.duplicategrid(fraction_water_grid)
    #ndep_grid.add_values(ndep_grid.length*[5.0])  # Set deposition to 5 kg/ha
    #print "*****************************************************************"

    cellarea = ascraster.Asciigrid(ascii_file=params.cellarea,\
                                   mask=mask,numtype=float)
    
    # Create new N deposition grid with N deposition load on water
    N_deposition = ascraster.duplicategrid(cellarea)

    # Make all cells zero.
    N_deposition.add_values(N_deposition.length*[0.0])
    N_deposition.nodata_value = -1

    for icell in xrange(fraction_water_grid.length):
        fraction_water_grid_cell = fraction_water_grid.get_data(icell,0.0)
        if (fraction_water_grid_cell > 0.0):
            ndep_cell = ndep_grid.get_data(icell,0.0)
            area_cell = cellarea.get_data(icell,0.0)
            # Area_cell is in m2 and deposition is in mg N/m2. So 1.0e-6 is from mg to kg.
            ndep_water = ndep_cell * area_cell * fraction_water_grid_cell * 1.0e-6
            N_deposition.set_data(icell,ndep_water)

    # Aggregate rasters
    aggregate.aggregate_grid(basin,N_deposition,mouth_dict,item="n_deposition_water")

    # Write to output file
    N_deposition.write_ascii_file(os.path.join(params.outputdir,"N_deposition_water.asc"))

    return N_deposition         
        
