#!/usr/bin/python2
from __future__ import print_function
import os
import sys
import argparse
import dakota.interfacing.parallel as parallel 

__author__ = 'J. Adam Stephens'
__copyright__ = 'Copyright 2014 Sandia Corporation'
__license__ = 'GNU Lesser General Public License'

def main():
    # mpitile supports SIMD and MIMD models. See the manpage for mpirun for a description.
    # To support MIMD, the command line arguments are split on colons (":"), and the
    # parser is invoked on each subset.
    
    # The first set of tokens must contain options to mpitile (e.g. --static).
    usage = "%(prog)s [%(prog)s options and global mpirun options] COMMAND1 : COMMAND2 : ..."
    parser = argparse.ArgumentParser(usage=usage,
            formatter_class=argparse.RawDescriptionHelpFormatter,
            description="Run COMMAND (or MIMD COMMANDs) in parallel under SLURM+OpenMPI on an available tile.",
            epilog="%(prog)s is a wrapper for OpenMPI's MPI launcher mpirun. The mpirun\n" +
            "command must be on the PATH.\n\n" +
            "Options and arguments may be passed to mpirun (or the user's code) " +
            "by including\nthem in COMMANDs.\n\n" +
            "Example 1: Launch my_sim using 32 tasks, assuming a dedicated master node. Pass\n" +
            "the option --bind-to none to mpirun and -verbose=OFF to my_sim.\n\n" +
            "      %(prog)s -np 32 --dedicated-master NODE --bind-to none my_sim -verbose=OFF" +
            "\n\nExample 2: Launch my_sim1 and my_sim2 in MIMD configuration using 16 " +
            "tasks\nand static scheduling.\n\n" +
            "      %(prog)s -np 16 --static --eval-num=$num my_sim1 : -np 16 my_sim2")

    parser.add_argument("-n","-np","-c","--n",type=int, dest="applic_tasks",
            help="Number of MPI processes (tasks)", required=True)
    parser.add_argument("-m","--dedicated-master",action="store",dest="dedicated_master",
            help="Reserve the first NODE or TILE for Dakota", choices=['NODE','TILE'])
    parser.add_argument("-u","--lock-id", dest="lock_id",
            help="Write lock files named <LOCK_ID>.<tile>. (Default: <SLURM_JOB_ID>.<tile>)")
    parser.add_argument("-l","--lock-dir",dest="lock_dir",
            help="Specify lock file directotry (Default: $HOME/.DakotaEvalTiling)")
    parser.add_argument("-s", "--static-scheduling",dest="static",action="store_true",
            default=False, 
            help="Use static scheduling instead of dynamic (no lock files). Use only in " +
            "concert with Dakota input keywords 'local_evaluation_scheduling static'")
    parser.add_argument("-e", "--eval-num", type=int, dest="eval_num",
            help="The evaluation number determines tile placement when using static " +
            "scheduling. (This option is ignored for the default, dynamic scheduling.)")
    parser.add_argument("-p", "--params-file", type=str, dest="params_file",
            help="Extract the evaluation number from the Dakota parameters file." + 
            "(This option is ignored for the default, dynamic scheduling, and when " +
            "--eval-num is given)")
    #parser.add_argument("command",nargs="*", help=argparse.SUPPRESS)

    # Break up the command line arguments on ":"
    all_cl = sys.argv[1:]
    command_lines = []
    while True:
        try:
            cindex = all_cl.index(":")
            command_lines.append(all_cl[:cindex])
            del all_cl[:cindex+1]
        except ValueError:
            command_lines.append(all_cl[:])
            break

    # The first (and usually the only) set of tokens contains options to 
    # mpitile. Process these separately in order to extract them.
    tile_args, command = parser.parse_known_args(command_lines[0])
    # Commands is a list of tuples. Each tuple is (applic_tasks, command)
    commands = [(tile_args.applic_tasks, command)]
    for line in command_lines[1:]:
        args, command = parser.parse_known_args(line)
        commands.append( (args.applic_tasks, command) )
        if args.dedicated_master or args.lock_dir or args.lock_id or \
                args.static or args.eval_num or args.params_file:
            print("WARNING: All mpitile options in MIMD COMMANDs except APPLIC_PROCS will be ignored.")

    # Call tile_run_static or tile_run_dynamic to launch the user's code(s)
    if tile_args.static:
        if not tile_args.eval_num and not tile_args.params_file:
            print("An evaluation number (--eval-num) or Dakota parameters file (--params-file) " +
            "must be given when using --static.", file=sys.stderr)
            sys.exit(1)

        ret = parallel.tile_run_static(commands=commands,
                dedicated_master=tile_args.dedicated_master, eval_num=tile_args.eval_num, 
                parameters_file=tile_args.params_file)
    else:   
        ret = parallel.tile_run_dynamic(commands=commands,
                dedicated_master=tile_args.dedicated_master, lock_id=tile_args.lock_id, 
            lock_dir=tile_args.lock_dir)

    sys.exit(ret)

if __name__ == "__main__":
    try:
        main()
    except (parallel.ResourceError, parallel.MgrEnvError) as e:
        print("mpitile: " + str(e),file=sys.stderr)
        sys.exit(1)



