# ******************************************************
## Revision "$LastChangedDate: 2015-05-21 16:47:35 +0200 (Thu, 21 May 2015) $"
## Date "$LastChangedRevision: 510 $"
## Author "$LastChangedBy: arthurbeusen $"
## URL "$HeadURL: https://pbl.sliksvn.com/globalnutrients/P_model/trunk/tools/P-model.py $"
# ******************************************************

import sys
import os
import traceback

# Import Global ("generalcode") modules
# Generic code files ideally should be placed in the directory "generalcode"
# subfolder, but they may also be placed at the same base folder
__general = os.path.join(r"../../..", 'generalcode/trunk')
if os.path.exists(__general): 
    sys.path.insert(0, __general)
    print __general + " is added to the python search path for modules." 

# Import general python modules
import optparse
import profile

# Import own general modules
import aggregate
import ascraster
import cmd_options_general
import discharge
from error import MyError
import get_versioninfo
import hydraulic_load
import my_logging
import lu_area
#import make_balance_total
import make_fraction
import make_dict_lake
import mouth
import my_sys
import residence_time
import subgrid_retention
import travel_load
import water_specs
import write_dict

#Import own modules
import cmd_options_p
import erosion
import pmouth
import pbackground
import p_gnpp
import pretention
import p_surface_runoff
import p_wetlands 

def run_P_model(args):
    '''
    Run P model main routine
    @param listargs: list of arguments passed trough command-line or other script
    Must look like sys.argv so make sure is starts with scriptname.
    '''
    # Parse command-line arguments and set parameters for script
    try:
        param = cmd_options_p.InputP(args)
        params = param.options
        params_var = param.options_var
    except SystemExit:
        raise MyError("Error has occured in the reading of the commandline options.")
    

    # Start timer and logging
    s = my_sys.SimpleTimer()
    log = my_logging.Log(params.outputdir,"%s_%i.log" % (params.scenarioname,params.year))
    print "Log will be written to %s" % log.logFile
    
    # If no arguments are provided do a run with defaults in params
    if len(args) == 0:
        log.write_and_print("No arguments provided: starting default run...")
        
    # time start of run
    log.write_and_print("Starting run....")
    log.write("# Parameters used:",print_time=False,lcomment=False)
    for option in str(params_var).split(", "):
        log.write("--%s = %s" % (option.split(": ")[0].strip("{").strip("'"),
                                 option.split(": ")[1].strip("}").strip("'")),
                  print_time=False,lcomment=False)
    log.write("All parameters used:")
    for option in str(params).split(", "):
        log.write("# %s = %s" % (option.split(": ")[0].strip("{").strip("'"),
                                 option.split(": ")[1].strip("}").strip("'")),
                  print_time=False,lcomment=False)
    log.write("# End of all parameters used.",print_time=False,lcomment=False)                  
        
    # Check whether there are command-line arguments which are not used
    if (len(param.args) > 0):
        txt = "The following command line arguments will not be used:"
        log.write_and_print(txt + str(param.args))
        
    # Write svn information of input and scripts to log file.
    log.write("******************************************************",print_time=False,lcomment=True)
    log.write("Version information:",print_time=False,lcomment=True)
    log.write("Version information main script:",print_time=False,lcomment=True)
    log.write("Revision $LastChangedDate: 2015-05-21 16:47:35 +0200 (Thu, 21 May 2015) $",print_time=False,lcomment=True)
    log.write("Date $LastChangedRevision: 510 $",print_time=False,lcomment=True)

    message = get_versioninfo.get_versioninfo(params.water_inputdir,params.outputdir)
    message.extend(get_versioninfo.get_versioninfo(params.scenariodir,params.outputdir))
    message.extend(get_versioninfo.get_versioninfo(params.P_histdir,params.outputdir))
    message.extend(get_versioninfo.get_versioninfo(params.fix_input,params.outputdir))
    message.extend(get_versioninfo.get_versioninfo("tools",params.outputdir))
    for item in range(len(message)):
        log.write(str(message[item]),print_time=False,lcomment=True)
    log.write("******************************************************",print_time=False,lcomment=True)

    # Read mask of the input grids.
    # The mask takes care to do the calculations on an efficient way. Only
    # the grid cells which are in the mask are calculated and stored.
    if (params.lmask):
        mask = ascraster.create_mask(params.file_mask, 0.0,'GT',numtype=float)
        log.write_and_print(s.interval("Reading mask"))
    else:    
        mask = None
        log.write_and_print(s.interval("No mask is used for this simulation."))

    # Read river basin map
    basin = ascraster.Asciigrid(ascii_file=params.basin,mask=mask,numtype=int)
    if params.ldebug: print params.basin  + " has been read."         

    # Read the mouth coordinates of each river basin.
    lines = my_sys.my_readfile(params.mouth,my_sys_ldebug=0,fatal=1)

    # Make a library with river mouth information for each river basin.
    mouth_dict = {}
    # Put lat and lon coordinates in a library
    for line in lines:
        fields = line.split()
        if (len(fields) > 2):
            mouth_dict[int(fields[0])] = pmouth.PMouth(basinid=int(fields[0]),lon=fields[1],lat=fields[2])
            # Determin the index of the mouth in the mask grid. Does not matter which grid you take here
            mouth_dict[int(fields[0])].set_index_grid(basin)
            if (mouth_dict[int(fields[0])].index_grid == -1):
                print "River basin " + str(int(fields[0])) + " has been removed because mouth is not in mask."
                # Remove all river basins which have a mounth which is not in the mask of this simulation
                del mouth_dict[int(fields[0])]
            else:
                # Set index_grid for complete grid (not masked)
                mouth_dict[int(fields[0])].index_grid_full = mask[mouth_dict[int(fields[0])].index_grid]
                
    # Delete the info of params.mouth file
    del lines
    
    try:
        # Here we choose whether we want the whole terrertial part calculated or read the river load from grid.
        # If the label read_river_load_file is presented in the parameters file, then the river load is read from grid.
        # Otherwise all the compartments are calcuated. The value of the parameter read_river_load_file must be 
        # the filename of the grid file with the river load. The same holds for the read_diffuse_load_file. This file is used
        # to calculate the subgrid retention.

        filename = params.read_river_load_file
        if (not os.path.isfile(filename)):
            raise MyError("File '%s' does not exist (label read_river_load_file)." % filename)

        P_rivers = ascraster.Asciigrid(ascii_file=filename,mask=mask,numtype=float)
        aggregate.aggregate_grid(basin,P_rivers,mouth_dict,item="pload_river")
        log.write_and_print(s.interval("Total load to rivers is read from file " + filename + ". NO CALCULATION."))  

        filename = params.read_diffuse_load_file
        if (not os.path.isfile(filename)):
            raise MyError("File '%s' does not exist (label read_diffuse_load_file)." % filename)
        P_diffuse = ascraster.Asciigrid(ascii_file=filename,mask=mask,numtype=float)
        log.write_and_print(s.interval("Total diffuse load to rivers is read from file " + filename + ". NO CALCULATION.")) 

        # Read the net precipitation or the runoff
        pnet = ascraster.Asciigrid(ascii_file=params.pnet,mask=mask,numtype=float)
        if params.ldebug: print "Netto precipitation is read in P model."

    except AttributeError:            
        # Read grass, crop and natural area of the world and make km2 as unit.
        grassarea,croparea,natarea = lu_area.calculate(params,mask,mouth_dict,basin)
        log.write_and_print(s.interval("Calculating areas"))
        # Next statement is to calibrate the soil loss for each land use type in Europe.
        #erosion.calculate_erosion_callibration(params,mask,grassarea,croparea,natarea)
        #sys.exit(12)

        # Read P_balance for grass, arable and natural areas
        P_bal_grass   = ascraster.Asciigrid(ascii_file=params.P_bal_grass,mask=mask,numtype=float)
        if params.ldebug: print "P_bal_grass is read in calculation of P model."
        P_bal_arable  = ascraster.Asciigrid(ascii_file=params.P_bal_arable,mask=mask,numtype=float)
        if params.ldebug: print "P_bal_arable is read in calculation of P model."
        P_bal_natural = ascraster.Asciigrid(ascii_file=params.P_bal_natural,mask=mask,numtype=float)
        if params.ldebug: print "P_bal_natural is read in calculation of P model."
        log.write_and_print(s.interval("Reading balance files"))
        
        # Change budget terms from P2O5 into P:
        P2O5_to_P = 62.0/142.0
        P_bal_grass.multiply(P2O5_to_P)
        P_bal_arable.multiply(P2O5_to_P)
        P_bal_natural.multiply(P2O5_to_P)

        total_balance = ascraster.duplicategrid(P_bal_grass)       
        for icell in xrange(P_bal_grass.length):
            P_bal_grs_cell = P_bal_grass.get_data(icell,0.0)
            P_bal_arable_cell = P_bal_arable.get_data(icell,0.0)
            total_load = P_bal_grs_cell + P_bal_arable_cell
            total_balance.set_data(icell,total_load)

        aggregate.aggregate_grid(basin,total_balance,mouth_dict,item="pbudget")
        log.write_and_print(s.interval("Adding balances to a total balance."))
        P_bal_total = ascraster.duplicategrid(total_balance)

        # Read the fraction surface runoff
        fQsro = ascraster.Asciigrid(ascii_file=params.fQsro,mask=mask,numtype=float)
        if params.ldebug: print "Fraction surface runoff is read in P model."

        # Read the net precipitation or the runoff
        pnet = ascraster.Asciigrid(ascii_file=params.pnet,mask=mask,numtype=float)
        if params.ldebug: print "Netto precipitation is read in P model."

        # Calculate the P load of the surface runoff
        P_sro = p_surface_runoff.calculate(params,mask,grassarea,croparea,\
                                           P_bal_grass,P_bal_arable,fQsro,pnet,\
                                           mouth_dict,basin)                                    
        log.write_and_print(s.interval("Calculating P surface runoff"))

        # Substract the P_sro of the P_bal_total to get the rest.
        P_bal_total.substract(P_sro,default_nodata_value = -1.0)

        P_erosion = erosion.calculate_erosion(params,mask,grassarea,croparea,natarea,P_bal_total,\
                                              basin,mouth_dict)
        log.write_and_print(s.interval("Calculation of the historical P erosion"))
        # Next statement only for calibration purposes.
        #erosion.calculate_erosion_calibration(params,mask,grassarea,croparea,natarea)
        #raise MyError("Klaar")
        
        # Calculate the total P load to the rivers of diffuse source and the point source.
        P_point = ascraster.Asciigrid(ascii_file=params.P_point,mask=mask,numtype=float)
        if params.ldebug: print "P point source  is read in calculation of P model."
        aggregate.aggregate_grid(basin,P_point,mouth_dict,item="ppoint")

        # Read the total P load to the rivers of aquaculture.
        if os.path.isfile(params.P_aquaculture):
            P_aquaculture = ascraster.Asciigrid(ascii_file=params.P_aquaculture,mask=mask,numtype=float)
            if params.ldebug: print "P aquaculture is read in calculation of P model."
        else:
            # There is no aquaculture input file
            P_aquaculture = ascraster.duplicategrid(P_point)
            P_aquaculture.add_values(P_aquaculture.length * [0.0])
            if params.ldebug: print "P aquaculture is ZERO in P model."
        aggregate.aggregate_grid(basin,P_aquaculture,mouth_dict,item="p_aquaculture")

        # Read land area
        landarea = ascraster.Asciigrid(ascii_file=params.landarea,mask=mask,numtype=float)
        if params.ldebug: print params.landarea  + " has been read."
        
        # Calculate the background P (weathering)
        P_background = pbackground.calculate_P(params,mask,pnet,mouth_dict,basin,landarea)

        # Calculate the P from NPP in flooded areas
        P_gnpp = p_gnpp.calculate(params,mask,mouth_dict,basin)        

        # Calculate the P from NPP in wetlands
        P_wetland = p_wetlands.calculate(params,mask,mouth_dict,basin)        

        # Add all the diffuse sources.
        P_diffuse = P_sro
        P_diffuse.add(P_background,default_nodata_value = -1)
        P_diffuse.add(P_erosion,default_nodata_value = -1)
        P_diffuse.add(P_gnpp,default_nodata_value = -1.0)
        P_diffuse.add(P_wetland,default_nodata_value = -1.0)
        P_rivers = ascraster.duplicategrid(P_diffuse)
        P_diffuse.write_ascii_file(os.path.join(params.outputdir,"P_diffuse.asc"))
             
        # Aggregatie
        aggregate.aggregate_grid(basin,P_diffuse,mouth_dict,item="ptotal_diffuse")

        # Add the point sources to the diffuse sources to get the total load to the river
        P_rivers.add(P_point,default_nodata_value = -1.0)
        P_rivers.add(P_aquaculture,default_nodata_value = -1.0)

        # Write river load to output file and into the mouth_dict.
        #P_rivers.nodata_value = -9999.0
        P_rivers.write_ascii_file(os.path.join(params.outputdir,"P_load_river.asc"))
        aggregate.aggregate_grid(basin,P_rivers,mouth_dict,item="pload_river")
        log.write_and_print(s.interval("Calculating total load to rivers."))
        # Make nodata value -9999 (needed for running accuflux).
        P_rivers.nodata_value = -9999.0

    # Calculate accumulated water discharge in each cell or read data from file
    if (os.path.isfile(params.discharge)): 
        # Read discharge from file
        discharge_grid = ascraster.Asciigrid(ascii_file=params.discharge,mask=mask,numtype=float)
        discharge_grid.write_ascii_file(os.path.join(params.outputdir,"discharge.asc"))
        log.write_and_print(s.interval("Read discharge from file for each grid cell."))
    else:
        discharge_grid = discharge.calculate_discharge(params,mask,pnet)
    log.write_and_print(s.interval("Calculated discharge for each grid cell."))

    # Make a database for the lakes.
    lakeid_dict,lake_icell_dict = make_dict_lake.calculate(params,mask,discharge_grid)
              
    # Read the local drain direction map
    lddmap = ascraster.Asciigrid(ascii_file=params.ldd,mask=mask,numtype=int)

    # Calculate the the water area and the water body.
    water_area, water_body = water_specs.calculate(params,mask,mouth_dict,basin,pnet,discharge_grid)
    log.write_and_print(s.interval("Calculated water specs for each grid cell."))
 
    # Calculate the subgrid_retention for the diffuse sources which flow to the river.
    retentionload_subgrid,retention_subgrid = \
                   subgrid_retention.calculate_subgrid_retention(params,mask,mouth_dict,basin,lake_icell_dict,P_diffuse,\
                                                                 water_body,pnet,params.P_vf_river)
    # Reduce the river load by the retentionload on the subgrid
    P_rivers.substract(retentionload_subgrid)
    log.write_and_print(s.interval("Calculated retention on the subgrid."))

    #residence_time_grid,traveltime,retention,hydraulic_load_grid = \
    #                   retention_all.calculate(params,mask,mouth_dict,lake_icell_dict,basin,discharge_grid,lddmap,P_rivers)
    #log.write_and_print(s.interval("Calculated retention for each grid cell."))

    # Calculate the residence and traveltime to the mouth of the river in a grid cell
    residence_time_grid,traveltime = \
                 residence_time.calculate(params,mask,mouth_dict,lake_icell_dict,basin,discharge_grid,lddmap)
    log.write_and_print(s.interval("Calculated residence time and traveltime for each grid cell."))

    # Calculate the hydraulic load in a grid cell
    hydraulic_load_grid = \
                 hydraulic_load.calculate(params,mask,residence_time_grid,\
                 lake_icell_dict)
    log.write_and_print(s.interval("Calculated hydraulic load for each grid cell."))

    # Calculate the index as a combination of travel time and the river load 
    travel_load.calculate(mouth_dict,basin,P_rivers,traveltime,"sum_Ptraveltime",\
                          "avg_Ptraveltime","pload_river")
    log.write_and_print(s.interval("Calculated travel load for each grid cell."))

    # For testing purpose!
    temperature = ascraster.Asciigrid(ascii_file=params.water_temperature,mask=mask,numtype=float)
    # Calculate the index as a combination of travel time and the river load 
    travel_load.calculate(mouth_dict,basin,P_rivers,temperature,"sum_Ptemperature",\
                          "avg_Ptemperature","pload_river")
    travel_load.calculate(mouth_dict,basin,P_rivers,hydraulic_load_grid,"sum_Phydraulic_load",\
                          "avg_Phydraulic_load","pload_river")

    log.write_and_print(s.interval("Calculated travel load for each grid cell."))

   
    # Calculate the retention in the rivers.
    P_load_swr,retention = pretention.calculate(params,mask,mouth_dict,lake_icell_dict,lddmap,P_rivers,hydraulic_load_grid,\
                                      water_body,pnet,discharge_grid,residence_time_grid)


    # Add attributes to the lakeid_dict.
    for key in lakeid_dict:
        try:
            icell = lakeid_dict[key].outcell
            setattr(lakeid_dict[key],"temperature",lake_icell_dict[icell].temperature)
            setattr(lakeid_dict[key],"fraction_water",lake_icell_dict[icell].fraction_water)
            try:
                setattr(lakeid_dict[key],"load_in",lake_icell_dict[icell].load_in)
                setattr(lakeid_dict[key],"load_out",lake_icell_dict[icell].load_out)
                setattr(lakeid_dict[key],"retention_load",lake_icell_dict[icell].load_in - lake_icell_dict[icell].load_out)
                setattr(lakeid_dict[key],"retention",lake_icell_dict[icell].retention) 
            except AttributeError:
                setattr(lakeid_dict[key],"load_in",0.0)
                setattr(lakeid_dict[key],"load_out",0.0)
                setattr(lakeid_dict[key],"retention_load",0.0)
                setattr(lakeid_dict[key],"retention",0.0)
        except AttributeError:
            pass

    # Transfer lake attributes to the river output.
    for key in lakeid_dict:
        try:
            id = lakeid_dict[key].basinid
            if (int(key) > 10000):
                # Natural lake
                mouth_dict[id].number_of_lakes += 1
                mouth_dict[id].lakes_water_area += float(lakeid_dict[key].water_area)
                mouth_dict[id].lakes_volume += float(lakeid_dict[key].volume)
                mouth_dict[id].lakes_retention_load += float(lakeid_dict[key].retention_load)
                mouth_dict[id].lakes_load_in += float(lakeid_dict[key].load_in)
                mouth_dict[id].lakes_load_out += float(lakeid_dict[key].load_out)
            else:
                # Reservoir
                mouth_dict[id].number_of_reserv += 1
                mouth_dict[id].reserv_water_area += float(lakeid_dict[key].water_area)
                mouth_dict[id].reserv_volume += float(lakeid_dict[key].volume)
                mouth_dict[id].reserv_retention_load += float(lakeid_dict[key].retention_load)
                mouth_dict[id].reserv_load_in += float(lakeid_dict[key].load_in)
                mouth_dict[id].reserv_load_out += float(lakeid_dict[key].load_out)
        except (AttributeError,KeyError):
            print "************ River "+str(id) + " not found for lake " + str(key) 
            pass
    # Calculate the fraction lakes/reservoir retention of the total river retention and the lake/reservoir retention
    for key in mouth_dict:
        try:
            mouth_dict[key].fr_lakes_retention = mouth_dict[key].lakes_retention_load/(mouth_dict[key].retention * mouth_dict[key].pload_river)
        except ZeroDivisionError:
            mouth_dict[key].fr_lakes_retention = 0.0
        try:
            mouth_dict[key].fr_reserv_retention = mouth_dict[key].reserv_retention_load/(mouth_dict[key].retention * mouth_dict[key].pload_river)
        except ZeroDivisionError:
            mouth_dict[key].fr_reserv_retention = 0.0
        try:
            mouth_dict[key].lakes_retention = 1.0 - (mouth_dict[key].lakes_load_out/mouth_dict[key].lakes_load_in)
        except ZeroDivisionError:
            mouth_dict[key].lakes_retention = 0.0
        try:
            mouth_dict[key].reserv_retention = 1.0 - (mouth_dict[key].reserv_load_out/mouth_dict[key].reserv_load_in)
        except ZeroDivisionError:
            mouth_dict[key].reserv_retention = 0.0

    log.write_and_print(s.interval("Calculated retention for each river basin."))

    # Aggregation of the tabular information on another id-grid than basinid's
    if (params.aggregationgrid != None):
        for item in range(params.aggregationgrid_number):
            # Make output filename and do the aggregation
            gridname = getattr(params,"aggregationgrid"+str(item))
            qq = os.path.basename(gridname)
            filename = os.path.join(params.outputdir,params.scenarioname +"_" + str(params.year) + "_" +  qq + "_P.csv")
            mouth.aggregate(mouth_dict,gridname,mask,filename,params.sep)                       
    
    # Aggregate to global level
    for key in mouth_dict.keys():
        try:
            mouth_dict["Total"].sum(mouth_dict[key])
        except KeyError:
            mouth_dict["Total"] = mouth_dict[key].copy()   

    # Write all information of the river basin to output file:
    fp = open(os.path.join(params.outputdir,params.scenarioname +"_" + str(params.year) + "_P.csv"),"w")
    for key in mouth_dict:
        # Use this loop to take just one element to write header and
        # then stop the loop.
        mouth_dict[key].write_header(fp,sep=params.sep)
        break
    write_dict.write_dict(fp_in=fp,filename=None,dic=mouth_dict,headerkey="",sep=params.sep,lkey=0,ldebug=1,fatal=0)
    fp.close()

    # Write all information of the lakes and reservoirs to output file:
    fp = open(os.path.join(params.outputdir,"lakes_" + params.scenarioname + "_" + str(params.year) + "_P.csv"),"w")
    for key in lakeid_dict:
        # Use this loop to take just one element to write header and
        # then stop the loop.
        lakeid_dict[key].write_header(fp,sep=params.sep)
        break
    write_dict.write_dict(fp_in=fp,filename=None,dic=lakeid_dict,headerkey="",sep=params.sep,lkey=0,ldebug=1,fatal=0)
    fp.close()

    log.write_and_print(s.total("Total run"))    
    del log
    
if __name__ == "__main__":
        
    try:
        #profile.run('run_P_model(sys.argv)')
        run_P_model(sys.argv)
    except MyError, val:
        val.write()
        #print sys.exc_type
        #print traceback.print_exc()
    except (optparse.OptionValueError,
            ascraster.ASCIIGridError,
            Exception), val:
        print str(val)
        print sys.exc_type
        print traceback.print_exc()
    except:
        print "***** ERROR ******"
        print "P_model.py failed."
        print sys.exc_type
        print traceback.print_exc()
