# ******************************************************
## Revision "$LastChangedDate: 2013-12-19 17:49:14 +0100 (Thu, 19 Dec 2013) $"
## Date "$LastChangedRevision: 433 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/P_model/trunk/tools/erosion.py $"
# ******************************************************

import os
import ascraster
import interpolate
import math
import my_sys
import make_fraction
import aggregate
from error import *

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

    # 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_calibration(params,mask,grassarea,croparea,natarea):
    '''
    Use this function to calibrate the soil loss. This can only be done for Europe (Cerda et al 2010).
    The outcome of this routine are three output files on the output directory: erodibility_arable.csv, erodibility_grass.csv and
    erodibility_nat.asc.
    Read these file and copy the content in the erosion_table.xls on this directory.
    The calibration takes care that the soil loss for europe is:
    Nat area 0.15 t/ha, grass area 0.4 t/ha and croparea 3.6t/ha soil loss
    Parameters for in the parameter_inputfile are in columns U,W and Y on row 49 in excel file erosion_table.xls
    '''
    
    # 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)
    erodibility.write_ascii_file(os.path.join(params.outputdir,"erodibility.asc"))
    # Read country file to aggregate the results
    gcnt = ascraster.Asciigrid(ascii_file="/home/arthur/globalnutrients/pointsources/input_Anne/gcountry.map",mask=mask,numtype=int)

    # Make a combination of landuse and erodibility to get a soil loss per hectare.
    grid = ascraster.duplicategrid(croparea)
    for icell in range(grid.length):
        val = grid.get_data(icell)
        if (val != None):
            erod_cell = erodibility.get_data(icell,0.0)
            # Multiply with area
            erod_cell *= val
            grid.set_data(icell,erod_cell)
    
    grid.write_ascii_file(os.path.join(params.outputdir,"erodibility_arable.asc"))
    dic = {}
    for icell in range(gcnt.length):
        icnt = gcnt.get_data(icell,-1)
        if (icnt > 0):
            try:
                qq = dic[int(icnt)]
                dic[int(icnt)][0] += croparea.get_data(icell,0.0)
                dic[int(icnt)][1] += grid.get_data(icell,0.0)
            except KeyError:
                dic[int(icnt)] = [croparea.get_data(icell,0.0),grid.get_data(icell,0.0)]
    fp = open(os.path.join(params.outputdir,"erodibility_arable.csv"),"w")
    fp.write("isocode" + ";" + "croparea" + ";" + "erosion" + "\n")
    for key in sorted(dic):
        fp.write(str(key) + ";" + str(dic[key][0]) + ";" + str(dic[key][1]) + "\n")
    fp.close()

    # Make a combination of landuse and erodibility to get a soil loss per hectare.
    grid = ascraster.duplicategrid(grassarea)
    for icell in range(grid.length):
        val = grid.get_data(icell)
        if (val != None):
            erod_cell = erodibility.get_data(icell,0.0)
            # Multiply with area
            erod_cell *= val
            grid.set_data(icell,erod_cell)
    
    grid.write_ascii_file(os.path.join(params.outputdir,"erodibility_grass.asc"))
    dic = {}
    for icell in range(gcnt.length):
        icnt = gcnt.get_data(icell,-1)
        if (icnt > 0):
            try:
                qq = dic[int(icnt)]
                dic[int(icnt)][0] += grassarea.get_data(icell,0.0)
                dic[int(icnt)][1] += grid.get_data(icell,0.0)
            except KeyError:
                dic[int(icnt)] = [grassarea.get_data(icell,0.0),grid.get_data(icell,0.0)]
    fp = open(os.path.join(params.outputdir,"erodibility_grass.csv"),"w")
    fp.write("isocode" + ";" + "grassarea" + ";" + "erosion" + "\n")
    for key in sorted(dic):
        fp.write(str(key) + ";" + str(dic[key][0]) + ";" + str(dic[key][1]) + "\n")
    fp.close()

    # Make a combination of landuse and erodibility to get a soil loss per hectare.
    grid = ascraster.duplicategrid(natarea)
    for icell in range(grid.length):
        val = grid.get_data(icell)
        if (val != None):
            erod_cell = erodibility.get_data(icell,0.0)
            # Multiply with area
            erod_cell *= val
            grid.set_data(icell,erod_cell)
    
    grid.write_ascii_file(os.path.join(params.outputdir,"erodibility_nat.asc"))
    dic = {}
    for icell in range(gcnt.length):
        icnt = gcnt.get_data(icell,-1)
        if (icnt > 0):
            try:
                qq = dic[int(icnt)]
                dic[int(icnt)][0] += natarea.get_data(icell,0.0)
                dic[int(icnt)][1] += grid.get_data(icell,0.0)
            except KeyError:
                dic[int(icnt)] = [natarea.get_data(icell,0.0),grid.get_data(icell,0.0)]
    fp = open(os.path.join(params.outputdir,"erodibility_nat.csv"),"w")
    fp.write("isocode" + ";" + "natarea" + ";" + "erosion" + "\n")
    for key in sorted(dic):
        fp.write(str(key) + ";" + str(dic[key][0]) + ";" + str(dic[key][1]) + "\n")
    fp.close()

def calculate_erosion(params,mask,grassarea,croparea,natarea,P_bal_total,\
                      basin,mouth_dict):
    '''
    The historical delivery of 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 Pcontent of the soil is given in g/m2.
    The change in Pcontent is calculated with the surplus/depletion of the soil budget and the soilloss.
    '''
    # Read the bulkdensity of the soil in t/m3 or g/cm3
    bulkdensity = ascraster.Asciigrid(ascii_file=params.bulkdensity_file,mask=mask,numtype=float)
    # Fill nodata of bulkdensity with 1.3 kg/cm3
    bulkdensity.nodata_value = 1.3
    # Remove nodata value, so 1.3 is used as a value!
    bulkdensity.nodata_value = None
    # Make sure the bulkdensity is above one.
    bulkdensity.multiply(1.0,minimum=1.0)

    # Read land area im km2
    landarea = ascraster.Asciigrid(ascii_file=params.landarea,mask=mask,numtype=float)

    # Get the last Pcontent grid known with the year (gr P per m2) for top soil layer of 30 cm 
    year,Pcontent,Psoilloss_agri_prev,P_bal_total_prev = find_last_Psoil(params,mask)

    # Pdensity_soil = (Pcontent * A (m2)/(A(m2) * 0.3 (m)) = 0.001 * Pcontent/0.3  = (0.001/0.3) * Pcontent = kg P/m3
    Pdensity_soil = ascraster.duplicategrid(Pcontent)
    Pdensity_soil.multiply(0.001/0.3)
    Pdensity_soil.write_ascii_file(os.path.join(params.outputdir,"Pdensity_soil.asc"),output_nodata_value=0.0)

    # Pdensity_soil of the initial condition = (Pcontent * A (m2)/(A(m2) * 0.3 (m)) = 0.001 * Pcontent/0.3  = (0.001/0.3) * Pcontent = kg P/m3
    # Read the initial Pcontent of the soil
    filename = os.path.join(params.P_histdir,"Pcontent_1900.asc")
    if (os.path.isfile(filename)):
        Pcontent_init = ascraster.Asciigrid(ascii_file=filename,mask=mask,numtype=float)
    else:
        Pcontent_init = ascraster.duplicategrid(Pcontent)
    Pdensity_soil_init = Pcontent_init
    Pdensity_soil_init.multiply(0.001/0.3)

    # Pstore in the soil
    Pstore = ascraster.duplicategrid(Pcontent)
    # Convert Pcontent from g/m2 to kg/km2
    Pstore.multiply(1000.0)
    # Multiply by area to get the full P storage in kg P
    Pstore.multiply(landarea,default_nodata_value = -1)
    Pstore.write_ascii_file(os.path.join(params.outputdir,"Pstore.asc"))

    # 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"))

    # Volume soilloss per gridcell in m3
    soilloss_arable.divide(bulkdensity)
    soilloss_grass.divide(bulkdensity)
    soilloss_nat.divide(bulkdensity)

    # Psoilloss = Pdensity_soil * Vol(m3) = kg P per year
    Psoilloss_arable = ascraster.duplicategrid(Pdensity_soil)  
    Psoilloss_arable.multiply(soilloss_arable,default_nodata_value = -1)
    aggregate.aggregate_grid(basin,Psoilloss_arable,mouth_dict,item="p_erosion_arable")
    Psoilloss_arable.write_ascii_file(os.path.join(params.outputdir,"Psoilloss_arable.asc"),output_nodata_value=0.0)

    Psoilloss_grass = ascraster.duplicategrid(Pdensity_soil)  
    Psoilloss_grass.multiply(soilloss_grass,default_nodata_value = -1)
    aggregate.aggregate_grid(basin,Psoilloss_grass,mouth_dict,item="p_erosion_grass")
    Psoilloss_grass.write_ascii_file(os.path.join(params.outputdir,"Psoilloss_grass.asc"),output_nodata_value=0.0)

    Psoilloss_nat = ascraster.duplicategrid(Pdensity_soil)  
    Psoilloss_nat.multiply(soilloss_nat,default_nodata_value = -1)
    aggregate.aggregate_grid(basin,Psoilloss_nat,mouth_dict,item="p_erosion_nat")
    Psoilloss_nat.write_ascii_file(os.path.join(params.outputdir,"Psoilloss_nat.asc"),output_nodata_value=0.0)

    # Calculate the new P storage in the soil. The Ploss on natural land is not included in the Psoil change!
    # Calculate the total loss on agricultural land from the previous time untill now with this current loss
    Psoilloss_tot = ascraster.duplicategrid(Psoilloss_arable)
    Psoilloss_tot.add(Psoilloss_grass)
    Psoilloss_agri = ascraster.duplicategrid(Psoilloss_tot)
    Psoilloss_tot.add(Psoilloss_agri_prev)
    # Multiply the average Psoilloss_agri by the number of years and 0.5 because average is 0.5 * (current + previous)
    Psoilloss_tot.multiply(0.5 * float(params.year - year))
    Pstore.substract(Psoilloss_tot,default_nodata_value = -1,minimum=0)
    # Take the addition of P surplus/depletion and add this to the compartment.
    # Average of the change on previous time and the current change.
    P_bal_total_tot = ascraster.duplicategrid(P_bal_total)
    P_bal_total_tot.add(P_bal_total_prev)
    # Multiply the average P_bal_total by the number of years and 0.5 because average is 0.5 * (current + previous)
    P_bal_total_tot.multiply(0.5 * float(params.year - year))
    Pstore.add(P_bal_total_tot,default_nodata_value = -1,minimum=0)
    # Add at the bottom of the soil box the same volume of soilloss with the initial Pcontent.
    Psoiladdition = Psoilloss_tot
    # Make volume of soil of total P flux.
    Psoiladdition.divide(Pdensity_soil)
    # Calculate the P flux with the initial Pcontent.
    Psoiladdition.multiply(Pdensity_soil_init)
    # Add this to the soil (from below).
    Pstore.add(Psoiladdition,default_nodata_value = -1)

    # Make a new P content
    # Conversion from kg P to kg P per km2.
    Pstore.divide(landarea,default_nodata_value = -1.0)
    # Convert from kg/km2 to g/m2
    Pstore.multiply(0.001)

    # Write the three files to output file which can be used on the p_histdir directory.
    Pstore.write_ascii_file(os.path.join(params.outputdir,"Pcontent.asc"),output_nodata_value=0.0)
    # Write Psoilloss to output file
    Psoilloss_agri.write_ascii_file(os.path.join(params.outputdir,"Psoilloss_agri.asc"),output_nodata_value=0.0)
    # Write P_bal_total to output file
    P_bal_total.write_ascii_file(os.path.join(params.outputdir,"P_bal_total.asc"),output_nodata_value=0.0)    

    # Calculate the difference with the starting map of Pcontent.
    for i in range(1900,params.year):
        filename = os.path.join(params.P_histdir,"Pcontent_"+str(i)+".asc")
        if (os.path.isfile(filename)):
            Pcontent = ascraster.Asciigrid(ascii_file=filename,mask=mask,numtype=float)
            Pstore.divide(Pcontent,default_nodata_value = -1)
            Pstore.write_ascii_file(os.path.join(params.outputdir,"Pcontent_change.asc"),output_nodata_value=0.0)
            break

    # Add natural and agricultural Psoilloss for transport to the river
    Psoilloss = Psoilloss_agri
    Psoilloss.add(Psoilloss_nat,default_nodata_value = -1)
    aggregate.aggregate_grid(basin,Psoilloss,mouth_dict,item="p_erosion")
    Psoilloss.write_ascii_file(os.path.join(params.outputdir,"Psoilloss.asc"),output_nodata_value=0.0)

    return Psoilloss
    

def find_last_Psoil(params,mask):
    '''
    Returns the last known Psoil raster from the P_histdir directory
    '''

    lfound = False
    # Look for the year in the filename
    for i in range(params.year-1,1899,-1):
        filename = os.path.join(params.P_histdir,"Pcontent_"+str(i)+".asc")
        if (os.path.isfile(filename)):
            # Found file
            lfound = True
            print "Historical file found: " + filename
            # Read grid
            Pcontent = ascraster.Asciigrid(ascii_file=os.path.join(params.P_histdir,filename),mask=mask,numtype=float)
            filename = os.path.join(params.P_histdir,"Psoilloss_agri_"+str(i)+".asc")
            Psoilloss_agri = ascraster.Asciigrid(ascii_file=os.path.join(params.P_histdir,filename),mask=mask,numtype=float)
            print "Historical file found: " + filename
            filename = os.path.join(params.P_histdir,"P_bal_total_"+str(i)+".asc")
            P_bal_total = ascraster.Asciigrid(ascii_file=os.path.join(params.P_histdir,filename),mask=mask,numtype=float)
            print "Historical file found: " + filename            
            # Store year
            year = i   
            # Stop searching
            break
    if not (lfound):
        print "Historical file NOT found." 
        # There is no previous Pcontent, so take the fixed Pcontent file.
        Pcontent = ascraster.Asciigrid(ascii_file=params.Pcontent_file,mask=mask,numtype=float)
        # Fill the missing data on land with the average of the Pcontent world wide
        b = [x for x in Pcontent.values if x > Pcontent.nodata_value] 
        avg_value = sum(b)/len(b)
        print "Missing data in Pcontent is filled with ",avg_value
        Pcontent.nodata_value = avg_value
        Pcontent.nodata_value = None
        # Change soil depth (50 cm) to (30 cm)
        Pcontent.multiply(0.6)
        # Assume no flux at this time
        Psoilloss_agri = ascraster.duplicategrid(Pcontent)
        Psoilloss_agri.add_values(Psoilloss_agri.length * [0.0])
        P_bal_total = ascraster.duplicategrid(Pcontent)
        P_bal_total.add_values(P_bal_total.length * [0.0])
        
        year = 1899
    
    return year,Pcontent,Psoilloss_agri,P_bal_total
