#!/usr/bin/python
import sys
import os
import re
import random
import time

if os.path.exists( "dot_master.py" ):
    sys.path += [ os.getcwd() ]
    import dot_master
else:
    dot_master = None

dot_master_fmt   = "%-39s = %s\n"
error_file       = "RMS_20XX_ERRORS_"
multi_seed_file  = "random.seeds"


##################################################################
#
# The purpose of this script is to run a per-instance RMS workflow in
# batch from eg ERT. The script expects the following input as arguments
# on the commandline
# 
#   <SIMULATION_PATH>: The path to where the simulation should run. Directory 
#                      must exist.
#
#   <REALIZATION_NR>: An integer saying which realiziton number this is. 
#
#   <RMS_VERSION> : This is a string which indicates which RMS version
#                   should be run. The allowed versions are the keys in
#                   the dictionary rms_version_table below, there should
#                   be no problems adding more versions to this table.
# 
#   <RMS_PROJECT> : This is the (full) path to the source RMS project.
# 
# 
#   <REPLACE_PATH>: The absolute path (set in the RMS project) part which 
#                   should be replaced with the current directory.     
#
#   <RMS_WORKFLOW1>: This is the name of the workflow we want to run.
#
#   <RMS_WORKFLOW2>: This is the name of the workflow we want to run.
#
#   <RMS_WORKFLOW3>: This is the name of the workflow we want to run.
#
#   
#   Example:
# 
#   bash%  run_RMS_20xx.py  0  2009.3  /d/MyField/models/rms/geomodel.pro  /OLDPATH  WorkFlow1  WorkFlow2 ...
# 
# The script works as follows:
#
#  1. The first commandline argument is the runpath of the script, and 
#     it will start by chdir() there.
#
#
#  2. The script recursively walks through the original project
#     directory, and does the following:
# 
#         a) All directories found in the original project are created
#            in the copy project, thus creating a empty project
#            skeleton.
#            If a non-directory entry with the same name is found in the
#            copy project the script will exit (this is maybe a bit over
#            the top strict??)
# 
#         b) For files in the source directory the script will do the
#            following:
#             
#            o If the name of the file is ".master" the script will
#              _copy_ the file from the source project to the copy
#              project, and on the way it will perform the string
#              substitutions asked for.
# 
#              or   
#  
#            o If a file with this name already exists in the target
#              project nothing will be done.
# 
#              else 
# 
#            o A symlink will be created from the copy project to the
#              original source project.
# 
# 
#  3. The lockfiles "project-lock-file" and ".lock" are removed.
# 
# 
#  4. The program loops as follows:
#     o Updating the global seed (incrementing by one with each subsequent iteration).
#     o Running the rms binary for each workflow.
#
#
# Observe that this script has no explicit built in support for
# per-instance files (i.e. IPL scripts).  The way that is meant to be
# solved is that an external program writes those files to the target
# project (creating necessary directories on the way) prior to calling
# this script.
#
################################################################


rms_license_file  = "/prog/roxar/licensing/geomaticLM.lic"


rms_version_table = {"2010.1.1"   : "/PATH/RMS/2010.1.1/linux-amd64-gcc_3_4-release/bin/rms",
                     "2010"       : "/PATH/RMS/2010/linux-amd64-gcc_3_2-release/bin/rms",
                     "2009"       : "/PATH/RMS/2009/linux-amd64-gcc_3_2-release/bin/rms",
                     "2009.2"     : "/PATH/RMS/2009.2/linux-amd64-gcc_3_2-release/bin/rms",
                     "2009.3"     : "/PATH/RMS/2009.3/linux-amd64-gcc_3_2-release/bin/rms",
                     "9.0.7.4"    : "/PATH/RMS/9.0.7.4/linux-amd64-gcc_3_2-release/bin/rms"}



# This function will write the error message on both  "RMS_20XX_ERRORS_" and
# stderr, and then exit.
def fatal_error( msg ):
    fileH = open(error_file , "w")
    fileH.write( msg )
    fileH.close()

    sys.exit( msg )


# This function will copy the file @src_entry to the file
# @target_entry, and on the way it will perform two different string
# substitutions. The .master file is organized like a key value list:
#
#    key1                   = Value1      
#    key2                   = Value2
#    ......
#      
# The input argument @value_subst is a list of tuples, where each
# occurence of the first element in the tuple is replaced with the
# second element in the tuple. 
#
# The input argument @set_list is a dictionary, if the key in the
# .master file is in the @set_list dictionary the value in the .master
# will be replaced with the value in the set_List dictionary.

def copy_dot_master( src_entry , target_entry , subst_list , set_list):
    print "Copying %s -> %s" % (src_entry , target_entry)
    if os.path.exists( target_entry ):
        os.unlink( target_entry )

    srcH    = open( src_entry    , "r")
    targetH = open( target_entry , "w")

    for line in srcH.readlines():
        for (old , new) in subst_list:
            line = line.replace( old , new )
            
        tmp = re.split( "\s*=\s*" , line)
        key = tmp[0]
        if set_list.has_key( key) :
            targetH.write( dot_master_fmt % (key , set_list[key]))
        else:
            targetH.write( line )

    srcH.close()
    targetH.close()




def create_project_directory( arg_dict , path , entries ):
    offset       = arg_dict["offset"]
    target_root  = arg_dict["target_root"]
    src_path     = arg_dict["src_path"]
    subst        = arg_dict["subst_list"]
    set_list     = arg_dict["set_list"]
    
    newpath = "%s/%s" % ( target_root , path[offset:] )

    #Creating all the directories of the project
    if not os.path.exists( newpath ):
        os.makedirs( newpath )
    else:
        if os.path.islink( newpath ):
            fatal_error("Entry:%s already exists as a symbolic link. Clean up first..." % newpath)
            
        if not os.path.isdir( newpath ):
            fatal_error("Entry:%s already exists - and it is not a directory. Clean up first..." % newpath)

    # Linking in all the files
    for entry in entries:
        src_entry    = "%s/%s" % (path , entry)
        if os.path.isfile( src_entry ):
            target_entry = "%s/%s" % (newpath , entry)
            if entry == ".master":
                copy_dot_master( src_entry , target_entry , subst , set_list)
            else:
                if not os.path.exists( target_entry ):
                    os.symlink( src_entry , target_entry )





# This function duplicate the directory structure of the project
# located in src_path. The newly created project will be created in
# the current directory. All files are symlinked.

def create_project_directories( src_path , project , target_path , subst_list , set_list):
    src_project = os.path.join( src_path , project )
    arg_dict = {"offset"     : len(src_path) ,
                "target_root": target_path  , 
                "src_path"   : src_path      ,
                "subst_list" : subst_list ,
                "set_list"   : set_list} 
    print "Creating all project directories: %s/%s" % (arg_dict["target_root"] , project )
    os.path.walk( src_project , create_project_directory , arg_dict)
    print "Setting up RMS project complete."
    

#################################################################
# This function will update the seed in the master_file. The function
# will update the following lines in the master_file:
# 
# global_seed   =    12345678
# seeds(n)      =    09999999
# 
# All the lines will be updated with the __same_seed__. This function is
# only called once, with the root level .master file as argument, if
# there are other random seeds around they are not touched.
#
# 
# What seed to use:
# -----------------
# 
# There are three different systems for choosing seed to use; the three
# methods are tried out in the following order:
# 
#     1. If the file "RMS_SEED" exists in current working directory the
#        script will use the value found in that file as seed.
# 
#     2. If the file random.seeds exists the script  will use integer
#        nr @iens as seed. The format of the @multi_seed_file is just a
#        list of integer, each on a separate line. The first integer should be the number of seeds
# 
#     3. If neither of the seed files exist the script will use
#        /dev/urandom to initialize the python rng, and then
#        random.randint() to get a seed.
#
# The seed which is actually used in the end is appended to the file
# RMS_SEED_USED. Observe that for both the files RMS_SEED and
# @multi_seed_file the format "better" be right - the script will fail
# hard if these files are not formatted correctly.


def get_seed( target_path , project , iens):
    single_seed_file = "%s/RMS_SEED" % target_path
    if os.path.exists( single_seed_file ):
        # Using existing single seed file
        fileH = open( single_seed_file , "r" )
        new_seed  = int( fileH.readline( ) )
        fileH.close()
    elif os.path.exists( multi_seed_file ):
        fileH = open( multi_seed_file , "r")
        seed_list = [ int(x) for x in fileH.readlines() ]
        fileH.close()
        if seed_list[0] <= iens:
            fatal_error("Asking for seed:%d seed_file:%s only has %d seeds\n" % (iens , multi_seed_file , seed_list[0]))
        new_seed = seed_list[iens + 1]
    else:
        # Generate a seed from /dev/urandom
        fileH = open( "/dev/urandom" , "r")
        buffer = fileH.read( 64 )
        fileH.close()
        random.seed( buffer )
        new_seed = random.randint( 0 , 21047483000 )

    fileH = open("%s/RMS_SEED_USED" % target_path , "a+")
    fileH.write("%s ... %d\n" % (time.strftime("%d-%m-%Y %H:%M:%S" , time.localtime(time.time())) , new_seed))
    fileH.close()
    return new_seed


def set_seed( seed ):
    master_file = "%s/%s/.master" % ( target_path , project )
    if not os.path.exists( master_file ):
        fatal_error("Could not find rms master file: %s - seems like a broken project" % master_file )
    
    fileH = open(master_file , "r")
    linelist = fileH.readlines()
    fileH.close()

    fileH = open( master_file , "w" )
    for line in linelist:
        line = re.sub("^global_seed\s+=\s+\d+" ,    "global_seed     = %d" % seed , line)
        line = re.sub("^seeds\((\d+)\)\s+=\s+\d+" , "seeds(\1)       = %d" % seed , line)
        fileH.write( line )
    fileH.close()


#################################################################

def init( arglist ):
    if len(arglist) < 7: 
        msg = """
The run_RMS_20xx script needs the following arguments
\n  cwd  iens  rms_version  path_to_project   replace_path workflow1 | workflow2   workflow3 ...

The available rms_versions are:%s
""" % rms_version_table.keys()
        fatal_error( msg )

    cwd         = sys.argv[1]
    iens        = int( sys.argv[2] )
    
    rms_version = sys.argv[3]
    if rms_version_table.has_key( rms_version ):
        rms_executable = rms_version_table[ rms_version ]
    else:
        fatal_error("Sorry: rms_version:%s not recognized - available:%s" % ( rms_version , rms_version_table.keys()))

        
    src_project = arglist[4]
    if os.path.exists( src_project ):
        if os.path.isdir( src_project ):
            (src_path , project) = os.path.split( src_project )
        else:
            fatal_error("Fatal error - %s is not a directory" % src_project)
    else:
        fatal_error("Fatal error - project:%s does not exist" % src_project)

    replace_path = sys.argv[5]

    rms_workflows = sys.argv[6:]

    return (cwd , iens , rms_executable , src_path , project , replace_path , rms_workflows )


#################################################################



def unlink_lockfiles( project ):
    lock_file1 = "%s/.lock" % project
    if os.path.exists( lock_file1 ):
        os.unlink( lock_file1 )

    lock_file2 = "%s/project_lock_file" % project
    if os.path.exists( lock_file2 ):
        os.unlink( lock_file2 )



#################################################################    

def run_rms( rms_executable , project , workflow ):
    os.environ["LM_LICENSE_FILE"] = rms_license_file
    cmd = "%s -nomesa -project %s -batch %s" % ( rms_executable , project , workflow )
    print "Starting rms with command: \"%s\"" % cmd
    os.system( cmd )
    print "RMS run is complete"
    


#################################################################
#################################################################
# Main program starting:

(target_path , iens , rms_executable , src_path , project , replace_path , rms_workflows) = init( sys.argv )
if os.path.exists( target_path ):
    os.chdir( target_path )
else:
    fatal_error("Directory:%s does not exist \n" % target_path )

subst_list = [(replace_path , target_path)]

# If the dot_master module is present we ignore the replace_path commandline argument
if dot_master:
    subst_list = dot_master.subst_list
    set_list = dot_master.set_list
else:
    set_list = {}
    
create_project_directories( src_path , project , target_path , subst_list , set_list)
unlink_lockfiles( project )        
global_seed = get_seed( target_path , project , iens)
for workflow in rms_workflows:
    set_seed( global_seed )
    run_rms( rms_executable , project , workflow )
    global_seed += 1








