# ******************************************************
## Revision "$LastChangedDate: 2013-12-16 15:52:10 +0100 (Mon, 16 Dec 2013) $"
## Date "$LastChangedRevision: 429 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/N_model/trunk/tools/n_erosion.py $"
# ******************************************************

import os
import ascraster
import math
import my_sys
import aggregate

def calculate_erodibility(params,mask):
    '''
    The erodibility of a soil is dependent of texture and slope. Both are time independent.
    Same function for N and P!
    '''

    # Read median slope in m/km
    slope = ascraster.Asciigrid(ascii_file=params.slope,mask=mask,numtype=float)
    # Read texture
    txt = ascraster.Asciigrid(ascii_file=params.s_txt,mask=mask,numtype=int)

    erodibility = ascraster.duplicategrid(txt)
    erodibility.nodata_value = -9999
    erodibility.add_values(erodibility.length*[erodibility.nodata_value])


    # Read vat table
    lines = my_sys.my_readfile(params.texture_erosion_file,my_sys_ldebug=0,fatal=1)
    # Make dictionairy of vat table
    dict_soil = {}
    for line in range(len(lines)):
        if (len(lines[line]) == 0):
            # Empty line
            continue
        if (lines[line][0] == "*" or lines[line][0] == "#"):
            # Comment line, so skip
            continue
        fields = lines[line].split()
        if (len(fields) == 2):
            dict_soil[int(fields[0])] = float(fields[1])
               
    # Calculate erodibility.
    #max_sum = 0.0
    for icell in xrange(slope.length):
        cell_slope = slope.get_data(icell,0.0)
        # In the formula the angle is used, this is equal slope/1000 (m/km)
        cell_slope = -1.5 + (17/(1.0 + math.exp(2.3 - 6.1*cell_slope*0.001)))

        cell_text = txt.get_data(icell)
        if not (cell_text == None):
            cell_slope = cell_slope * dict_soil[int(cell_text)]
            erodibility.set_data(icell,cell_slope)
        
        #if (cell_slope > max_sum): 
            #print  slope.get_data(icell,0.0), txt.get_data(icell),cell_slope
            #max_sum = cell_slope

    return erodibility



def calculate_erosion(params,mask,grassarea,croparea,natarea,\
                      basin,mouth_dict):
    '''
    The historical delivery of N and P is modelled as a semi erosion equation.
    First the erodibility of the soil (time independent) is calculated based on texture and slope.
    Then with landuse the erodibility is converted into a soil loss [ton per year]
    The volume of the soiloss is calculated with bulkdensity.

    The Ncontent of the soilloss is given as C. So conversion from C to N, which is landuse dependent is calculated.
    For N there is no accumulation/depletion of the soil!
    '''
    if params.ldebug: print "Start with calculation of N erosion."

    # Get the percentage of SOC  
    s_oc = ascraster.Asciigrid(ascii_file=params.s_oc,mask=mask,numtype=float)
    if params.ldebug: print "S_OC is read in calculation of n_erosion."
    # Make fraction of s_oc
    s_oc.multiply(0.01)

    # Calculate the erodibility of the soil based on slope and texture (European study of Cerda et al, 2010, Geomorphology 122 (2010) 167-177)     
    erodibility = calculate_erodibility(params,mask)

    # Calculate the soil loss for each grid cell ton soil per square kilometer.
    soilloss_arable = ascraster.duplicategrid(erodibility)
    soilloss_arable.multiply(params.soilloss_arable_factor)
    # Calculate the soil loss for each grid cell in ton soil.
    area = ascraster.duplicategrid(croparea)
    area.nodata_value = 0.0
    area.nodata_value = None
    soilloss_arable.multiply(area)
    # Soilloss_arable in ton per year
    soilloss_arable.write_ascii_file(os.path.join(params.outputdir,"soilloss_arable.asc"))

    # Calculate the soil loss for each grid cell ton soil per square kilometer.
    soilloss_grass = ascraster.duplicategrid(erodibility)
    soilloss_grass.multiply(params.soilloss_grass_factor)
    # Calculate the soil loss for each grid cell in ton soil.
    area = ascraster.duplicategrid(grassarea)
    area.nodata_value = 0.0
    area.nodata_value = None
    soilloss_grass.multiply(area)
    # Soilloss_grass in ton per year
    soilloss_grass.write_ascii_file(os.path.join(params.outputdir,"soilloss_grass.asc"))

    # Calculate the soil loss for each grid cell ton soil per square kilometer.
    soilloss_nat = ascraster.duplicategrid(erodibility)
    soilloss_nat.multiply(params.soilloss_nat_factor)
    # Calculate the soil loss for each grid cell in ton soil.
    area = ascraster.duplicategrid(natarea)
    area.nodata_value = 0.0
    area.nodata_value = None
    soilloss_nat.multiply(area)
    # Soilloss_nat in ton per year
    soilloss_nat.write_ascii_file(os.path.join(params.outputdir,"soilloss_nat.asc"))

    # Nsoilloss = conversion_tonnes_kg [1000] * (SOC [fraction] * soilloss [ton])/ CNratio [weight ratio] = kg N per year
    Nsoilloss_arable = ascraster.duplicategrid(s_oc)  
    Nsoilloss_arable.multiply(soilloss_arable,default_nodata_value = -1)
    Nsoilloss_arable.multiply(1000./params.CN_ratio_erosion_arable,default_nodata_value = -1)
    aggregate.aggregate_grid(basin,Nsoilloss_arable,mouth_dict,item="n_erosion_arable")
    Nsoilloss_arable.write_ascii_file(os.path.join(params.outputdir,"Nsoilloss_arable.asc"),output_nodata_value=0.0)

    Nsoilloss_grass = ascraster.duplicategrid(s_oc)  
    Nsoilloss_grass.multiply(soilloss_grass,default_nodata_value = -1)
    Nsoilloss_grass.multiply(1000./params.CN_ratio_erosion_grass,default_nodata_value = -1)
    aggregate.aggregate_grid(basin,Nsoilloss_grass,mouth_dict,item="n_erosion_grs")
    Nsoilloss_grass.write_ascii_file(os.path.join(params.outputdir,"Nsoilloss_grass.asc"),output_nodata_value=0.0)

    Nsoilloss_nat = ascraster.duplicategrid(s_oc)  
    Nsoilloss_nat.multiply(soilloss_nat,default_nodata_value = -1)
    Nsoilloss_nat.multiply(1000./params.CN_ratio_erosion_natural,default_nodata_value = -1)
    aggregate.aggregate_grid(basin,Nsoilloss_nat,mouth_dict,item="n_erosion_nat")
    Nsoilloss_nat.write_ascii_file(os.path.join(params.outputdir,"Nsoilloss_nat.asc"),output_nodata_value=0.0)

    # Calculate the total N soil erosion.
    Nsoilloss_tot = ascraster.duplicategrid(Nsoilloss_arable)
    Nsoilloss_tot.add(Nsoilloss_grass)
    Nsoilloss_tot.add(Nsoilloss_nat)
    aggregate.aggregate_grid(basin,Nsoilloss_tot,mouth_dict,item="n_erosion")
    Nsoilloss_tot.write_ascii_file(os.path.join(params.outputdir,"Nsoilloss.asc"),output_nodata_value=0.0)

    return Nsoilloss_tot

def erosion_per_country(params,mask):
    '''
    Aggregate the N erosion 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)

    # 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_erosion_tot",None)
    G.add_item("n_erosion_arable",None)
    G.add_item("n_erosion_grass",None)
    G.add_item("n_erosion_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))

    # Read erosion grids and aggregate them.
    grid = ascraster.Asciigrid(ascii_file=os.path.join(params.outputdir,"Nsoilloss.asc"),mask=mask,numtype=float)
    aggregate.aggregate_grid(isogrid,grid,dict_out,"n_erosion_tot")
    grid = ascraster.Asciigrid(ascii_file=os.path.join(params.outputdir,"Nsoilloss_arable.asc"),mask=mask,numtype=float)
    aggregate.aggregate_grid(isogrid,grid,dict_out,"n_erosion_arable")
    grid = ascraster.Asciigrid(ascii_file=os.path.join(params.outputdir,"Nsoilloss_grass.asc"),mask=mask,numtype=float)
    aggregate.aggregate_grid(isogrid,grid,dict_out,"n_erosion_grass")
    grid = ascraster.Asciigrid(ascii_file=os.path.join(params.outputdir,"Nsoilloss_nat.asc"),mask=mask,numtype=float)
    aggregate.aggregate_grid(isogrid,grid,dict_out,"n_erosion_natural")

    # Write to outputfile
    fp = open(os.path.join(params.outputdir,"erosion_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()
