# ******************************************************
## Revision "$LastChangedDate: 2015-05-21 09:30:05 +0200 (Thu, 21 May 2015) $"
## Date "$LastChangedRevision: 506 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/P_model/trunk/tools/p_surface_runoff.py $"
# ******************************************************

import os
import ascraster
import aggregate

def split_uptake(params,mask,grassarea,croparea):

    if params.ldebug: print "Start with split uptake into grass and arable."

    # Read P_uptake
    P_uptake = ascraster.Asciigrid(ascii_file=params.P_uptake,mask=mask,numtype=float)
    

    P_uptake_arable  = ascraster.duplicategrid(P_uptake)
    P_uptake_grass   = ascraster.duplicategrid(P_uptake)
    
    for icell in range(P_uptake.length):
        uptake_cell = P_uptake.get_data(icell)
        if (not uptake_cell == None):
            grass_cell  = grassarea.get_data(icell,0.0)
            arable_cell = croparea.get_data(icell,0.0)
            # Determine the area fraction of arable and grass for distribution of the P uptake. 
            try:
                fr_grass = grass_cell/(grass_cell+arable_cell)
                fr_arable  = 1.0 - fr_grass 
            except ZeroDivisionError:
                fr_grass = 0.0
                fr_arable  = 0.0
               
            P_uptake_grass.set_data(icell,fr_grass * uptake_cell)
            P_uptake_arable.set_data(icell,fr_arable * uptake_cell)

    return  P_uptake_grass,P_uptake_arable 

def calculate(params,mask,grassarea,croparea,P_bal_grass,\
              P_bal_arable,fQsro,pnet,\
              mouth_dict,basin):
    '''
    Calculates the Phosphorus surface runoff component.
    It returns the ascraster object with the Phosphorus surface runoff
    '''
    
    if params.ldebug: print "Start with calculation of surface runoff."

    if (params.luptake_tot):
        # Split one total uptake file into uptake_grass and uptake_arable 
        P_uptake_grass,P_uptake_arable = split_uptake(params,mask,grassarea,croparea)
        if params.ldebug: print "One uptake file (combined grass and arable) is read in calculation of surface runoff."
    else:  
        # Read P_uptake
        P_uptake_grass   = ascraster.Asciigrid(ascii_file=params.P_uptake_grass,mask=mask,numtype=float)
        P_uptake_arable  = ascraster.Asciigrid(ascii_file=params.P_uptake_arable,mask=mask,numtype=float)
        if params.ldebug: print "Uptake (grass and arable) files are read in calculation of surface runoff."
        
    # Make a copy of the return values
    P_sro       = ascraster.duplicategrid(P_uptake_grass)
    if (P_sro.nodata_value == None):
        P_sro.nodata_value = -1 
    P_sro_crop  = ascraster.duplicategrid(P_sro)
    P_sro_grass = ascraster.duplicategrid(P_sro)
    
    # Change uptake term from P2O5 into P:
    P2O5_to_P = 62.0/142.0

    for icell in range(P_sro.length):
        uptake_grass_cell = P_uptake_grass.get_data(icell,0.0) * P2O5_to_P
        uptake_arable_cell = P_uptake_arable.get_data(icell,0.0) * P2O5_to_P
        fQsro_cell = fQsro.get_data(icell,0.0)
        grass_cell = grassarea.get_data(icell,0.0)
        crop_cell = croparea.get_data(icell,0.0)
        
        # Default values
        Psro_cell_grass = 0.0    
        Psro_cell_crop  = 0.0
        
        if (grass_cell > 0.0):
            # There is grass
            pnet_cell = pnet.get_data(icell,0.0)
            if (pnet_cell < params.minimal_soil_infiltration):
                # There is not enough water to transport P
                Psro_cell_grass = 0.0
            else:
                P_bal_cell = P_bal_grass.get_data(icell,0.0)
                Ptot_cell = max(0.0,P_bal_cell + uptake_grass_cell)
                # 0.25 is the reduction of water for grass.
                Psro_cell_grass = params.P_calibration_factor * params.water_reduction_grass * fQsro_cell * Ptot_cell
                   
        if (crop_cell > 0.0):
            # There is arable land
            # For arable land, we assume irrigation. So P is also removed when pnet < params.minimal_soil_infiltration.
            P_bal_cell = P_bal_arable.get_data(icell,0.0)
            Ptot_cell = max(0.0,P_bal_cell + uptake_arable_cell)
            # 1.0 is the reduction of water for crops.
            Psro_cell_crop = params.P_calibration_factor * params.water_reduction_crop * fQsro_cell * Ptot_cell
        
        # Put result in raster
        total_sro = Psro_cell_grass + Psro_cell_crop
        P_sro.set_data(icell,total_sro)
        P_sro_crop.set_data(icell,Psro_cell_crop)
        P_sro_grass.set_data(icell,Psro_cell_grass)
        
    # Write P load to surface runoff to output file
    P_sro.write_ascii_file(os.path.join(params.outputdir,"P_sro.asc"))
    P_sro_crop.write_ascii_file(os.path.join(params.outputdir,"P_sro_arable.asc"))
    P_sro_grass.write_ascii_file(os.path.join(params.outputdir,"P_sro_grs.asc"))
    
    # Delete the data which is not used anymore
    del P_uptake_arable
    del P_uptake_grass

    # Aggregate the surface runoff
    aggregate.aggregate_grid(basin,P_sro,mouth_dict,item="psro")
    aggregate.aggregate_grid(basin,P_sro_crop,mouth_dict,item="psro_arable")
    aggregate.aggregate_grid(basin,P_sro_grass,mouth_dict,item="psro_grs")
  
    return P_sro
