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

__author__ = 'J. Adam Stephens'
__copyright__ = 'Copyright 2014-2022 National Technology & Engineering Solutions of Sandia, LLC (NTESS)'
__license__ = 'GNU Lesser General Public License'

def _create_primary_parser():
    # 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)\n" +
            "by including\nthem in COMMANDs. Insert '--' to stop option parsing\n" +
            "(to provide options to the COMMAND that collide with mpitile options).\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" +
            "\n\nExample 3: Use -- to halt parsing of the command line to prevent mpitile\n" +
            "from consuming the '-m foo' option to my_sim1.\n\n" +
            "      %(prog)s -np 16 --bind-to none -- my_sim1 -m foo"
            
            )

    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("-t","--tile-size",type=int, dest="tile_size",
            help="Provide an explicit tile size (when not equal to number of tasks)",
            default=0)
    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)")
    return parser

def _create_secondary_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("-n","-np","-c","--n",type=int, dest="applic_tasks", required=True)
    return parser


def _halt_parsing_split(tokens):
    num_splits = tokens.count("--")
    if num_splits == 0:
        return tokens, []
    elif num_splits > 1:
        print("mpitile: error: Include '--' only once per command", file=sys.stderr)
        sys.exit(-1)
    else:
        stop_parse_pos = tokens.index("--")
        return tokens[:stop_parse_pos], tokens[stop_parse_pos+1:]

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 each subset
    # are parsed separately. The first set must include any options to mpitile, and the
    # "primary" parser is used on it. Subsequent sets use the secondary parser to extract
    # the -n (and aliases) option to specify the number of tasks. All other options are
    # passed down to mpirun and the user command.
    # In addition to splitting on :, each command is split on --, which signals that parsing
    # by mpitile should cease. This allows users to provide options to their commands that
    # otherwise would collide with mpitile options (e.g. -m is an option to both mpitile
    # and the command, so the user separates the command using -- to prevent mpitile from
    # trying to parse it both times.) It is permissible to put options that mpitile
    # doesn't understand prior to --; these are assumed to be options to mpirun, and are
    # prepended to the command.
    
    # Break up the command line arguments on ":"
    all_cl = sys.argv[1:]
    command_lines = []  # will be a list of lists, one outer element per command
    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.
    # First split on --.
    to_parse, not_to_parse = _halt_parsing_split(command_lines[0])
    primary_parser = _create_primary_parser()
    tile_args, command = primary_parser.parse_known_args(to_parse)
    command += not_to_parse
    # Commands is a list of tuples that will eventually be executed.
	# Each tuple is (applic_tasks, command).
    commands = [(tile_args.applic_tasks, command)]
    if len(command_lines) > 1:  # MIMD
        secondary_parser = _create_secondary_parser()
		# Iterate the additional commands and parse them
        for line in command_lines[1:]:
            to_parse, not_to_parse = _halt_parsing_split(line)
            args, command = secondary_parser.parse_known_args(to_parse)
        commands.append( (args.applic_tasks, command + not_to_parse) )
    # 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,
                tile_size=tile_args.tile_size, parameters_file=tile_args.params_file)
    else:   
        ret = parallel.tile_run_dynamic(commands=commands,
                dedicated_master=tile_args.dedicated_master, tile_size=tile_args.tile_size,
                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)



