# ******************************************************
## Revision "$LastChangedDate: 2015-05-21 09:31:08 +0200 (Thu, 21 May 2015) $"
## Date "$LastChangedRevision: 507 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/N_model/trunk/tools/n_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 N_uptake
    N_uptake = ascraster.Asciigrid(ascii_file=params.N_uptake,mask=mask,numtype=float)
    

    N_uptake_arable  = ascraster.duplicategrid(N_uptake)
    N_uptake_grass   = ascraster.duplicategrid(N_uptake)
    
    for icell in range(N_uptake.length):
        uptake_cell = N_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 N 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
               
            N_uptake_grass.set_data(icell,fr_grass * uptake_cell)
            N_uptake_arable.set_data(icell,fr_arable * uptake_cell)

    return  N_uptake_grass,N_uptake_arable 
    
def calculate(params,mask,grassarea,croparea,natarea,N_bal_grass,\
              N_bal_arable,N_bal_natural,fQsro,pnet,\
              mouth_dict,basin):
    '''
    Calculates the nitrogen surface runoff component.
    It returns the ascraster object with the nitrogen surface runoff in kg N.
    The nitrogen sdurface components are set in the mouth_dict dictionary and written to ascii output files.
    The balances (grass,arable and natural) are changed by substracting the surface runoff load. 
    Also mismatch in data is corrected,
    when area does not match with the balances, than the balances are made zero.
    '''
    if params.ldebug: print "Start with calculation of surface runoff."

    # Produce error messages when balance is above or below this errorlevel (kg N)
    errorlevel = 10

    # Read the N input of manure, fertilizer, biological N fixation and atmospheric deposition (kg/ha year).
    # Losses of NH3 are already substracted from this input. 

    if (params.luptake_tot):
        # Split one total uptake file into uptake_grass and uptake_arable 
        N_uptake_grass,N_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 N_uptake
        N_uptake_grass   = ascraster.Asciigrid(ascii_file=params.N_uptake_grass,mask=mask,numtype=float)
        N_uptake_arable  = ascraster.Asciigrid(ascii_file=params.N_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
    N_sro       = ascraster.duplicategrid(N_uptake_grass)
    if (N_sro.nodata_value == None):
        N_sro.nodata_value = -1 
    N_sro_crop  = ascraster.duplicategrid(N_sro)
    N_sro_grass = ascraster.duplicategrid(N_sro)
    N_sro_nat   = ascraster.duplicategrid(N_sro)
    
    # Calculate the surface runoff per landuse type
    for icell in range(N_sro.length):
        uptake_grass_cell  = N_uptake_grass.get_data(icell,0.0)
        uptake_arable_cell = N_uptake_arable.get_data(icell,0.0)
        fQsro_cell = fQsro.get_data(icell,0.0)
        grass_cell = grassarea.get_data(icell,0.0)
        crop_cell  = croparea.get_data(icell,0.0)
        nat_cell   = natarea.get_data(icell,0.0)
        pnet_cell = pnet.get_data(icell,0.0)        

        # Default values
        Nsro_cell_grass = 0.0    
        Nsro_cell_crop  = 0.0
        Nsro_cell_nat   = 0.0
        
        N_bal_cell = N_bal_grass.get_data(icell,0.0)
        if (grass_cell > 0.0):
            # There is grass
            Ntot_cell = max(0.0,N_bal_cell + uptake_grass_cell)
            # 0.25 is the reduction of water for grass. 0.5 is the calibration factor (GvD, 2010).
            Nsro_cell_grass = params.N_calibration_factor * params.water_reduction_grass * fQsro_cell * Ntot_cell
        elif (abs(N_bal_cell) > 0.0):
            if (abs(N_bal_cell) > errorlevel):
                print "Grass area and grasbalance do not belong to each other. Gras area = 0.0 and balance_grass = " + str(N_bal_cell)
            N_bal_cell = 0.0
            Nsro_cell_grass = 0.0
        
        # Substracted surface runoff of the balance available for denitrification (must be positive)
        N_bal_grass.set_data(icell,max(0.0,N_bal_cell - Nsro_cell_grass))        
        
        N_bal_cell = N_bal_arable.get_data(icell,0.0)           
        if (crop_cell > 0.0):
            # There is arable land
            # Here assume that there is always enough water for surface runoff,for example irrigation.
            Ntot_cell = max(0.0,N_bal_cell + uptake_arable_cell)
            # 1.0 is the reduction of water for crops. 0.5 is the calibration factor (GvD, 2010).
            Nsro_cell_crop = params.N_calibration_factor * params.water_reduction_crop * fQsro_cell * Ntot_cell
        elif (abs(N_bal_cell) > 0.0):
            if (abs(N_bal_cell) > errorlevel):
                print "Crop area and cropbalance do not belong to each other. Crop area = 0.0 and balance_crop = " + str(N_bal_cell)
            N_bal_cell = 0.0
            Nsro_cell_crop = 0.0
        
        # Substracted surface runoff of the balance available for denitrification (must be positive)
        N_bal_arable.set_data(icell,max(0.0,N_bal_cell - Nsro_cell_crop))

        N_bal_cell = N_bal_natural.get_data(icell,0.0)
        if (nat_cell > 0.0):
            # There is natural land
            # 0.125 is the reduction of water for natural systems. 0.5 is the calibration factor (GvD, 2010).
            Nsro_cell_nat = params.N_calibration_factor * params.water_reduction_natural * fQsro_cell * max(0.0,N_bal_cell)
        elif (abs(N_bal_cell) > 0.0):
            if (abs(N_bal_cell) > errorlevel):
                print "Natural area and natural balance do not belong to each other. Natural area = 0.0 and balance_nat = " + str(N_bal_cell)
            N_bal_cell = 0.0
            Nsro_cell_nat = 0.0
        
        # Substracted surface runoff of the balance available for denitrification (must be positive)
        N_bal_natural.set_data(icell,max(0.0,N_bal_cell - Nsro_cell_nat))

        # Put result in raster
        total_sro = Nsro_cell_grass + Nsro_cell_crop + Nsro_cell_nat
        N_sro.set_data(icell,total_sro)
        N_sro_crop.set_data(icell,Nsro_cell_crop)
        N_sro_grass.set_data(icell,Nsro_cell_grass)
        N_sro_nat.set_data(icell,Nsro_cell_nat)
        
   
    # Write N load to surface runoff to output file
    N_sro.write_ascii_file(os.path.join(params.outputdir,"N_sro.asc"))
    N_sro_crop.write_ascii_file(os.path.join(params.outputdir,"N_sro_arable.asc"))
    N_sro_grass.write_ascii_file(os.path.join(params.outputdir,"N_sro_grs.asc"))
    N_sro_nat.write_ascii_file(os.path.join(params.outputdir,"N_sro_nat.asc")) 
    
    # Delete the data which is not used anymore
    del N_uptake_arable
    del N_uptake_grass

    # Aggregate the surface runoff
    aggregate.aggregate_grid(basin,N_sro,mouth_dict,item="nsro")
    aggregate.aggregate_grid(basin,N_sro_crop,mouth_dict,item="nsro_arable")
    aggregate.aggregate_grid(basin,N_sro_grass,mouth_dict,item="nsro_grs")
    aggregate.aggregate_grid(basin,N_sro_nat,mouth_dict,item="nsro_nat")
  
    return N_sro


def sro_per_country(params,mask):
    '''
    Aggregate the N surface runoff on country base and write to output file
    '''
    import aggregate
    import general_class
    
    # Read country grid.
    isogrid = ascraster.Asciigrid(ascii_file=r"/home/arthur/globalnutrients/scen_input/history_100/gcountry.map",mask=mask,numtype=int)
    isogrid.write_ascii_file(os.path.join(params.outputdir,"gcountry.asc"))
    # Make dictionary with isocode as keys and different erosion components as attributes.

    G = general_class.General()
    G.add_item("isocode",None)
    G.add_item("n_sro_tot",None)
    G.add_item("n_sro_arable",None)
    G.add_item("n_sro_grass",None)
    G.add_item("n_sro_natural",None)
    
    # Make a dictionary with isocodes as key.
    dict_out = {}
    for icell in range(isogrid.length):
        id = isogrid.get_data(icell,-1)
        if (id > 0):
            try:
                qq = dict_out[str(int(id))]
            except KeyError:
                dict_out[str(int(id))] = G.copy()
                dict_out[str(int(id))].set_val("isocode",int(id))
    print dict_out.keys()

    # Read erosion grids and aggregate them.
    grid = ascraster.Asciigrid(ascii_file=os.path.join(params.outputdir,"N_sro.asc"),mask=mask,numtype=float)
    aggregate.aggregate_grid(isogrid,grid,dict_out,"n_sro_tot")
    grid = ascraster.Asciigrid(ascii_file=os.path.join(params.outputdir,"N_sro_arable.asc"),mask=mask,numtype=float)
    aggregate.aggregate_grid(isogrid,grid,dict_out,"n_sro_arable")
    grid = ascraster.Asciigrid(ascii_file=os.path.join(params.outputdir,"N_sro_grs.asc"),mask=mask,numtype=float)
    aggregate.aggregate_grid(isogrid,grid,dict_out,"n_sro_grass")
    grid = ascraster.Asciigrid(ascii_file=os.path.join(params.outputdir,"N_sro_nat.asc"),mask=mask,numtype=float)
    aggregate.aggregate_grid(isogrid,grid,dict_out,"n_sro_natural")

    # Write to outputfile
    fp = open(os.path.join(params.outputdir,"nsro_country.csv"),"w")
    lheader = True
    for key in sorted(dict_out):
        dict_out[key].write(fp,sep=params.sep,lheader=lheader,NoneValue="0.0")
        lheader = False
    fp.close()
