#!/usr/bin/env python

"""
Script to run CIME tests.

Runs single tests or test suites based on either the input list or the testname or based
on an xml testlist if the xml suboption is provided.

If this tool is missing any feature that you need, please notify jgfouca@sandia.gov.
"""
from Tools.standard_script_setup import *

import update_acme_tests
from CIME.test_scheduler import TestScheduler, RUN_PHASE
from CIME.utils import expect, convert_to_seconds, compute_total_time, convert_to_babylonian_time, run_cmd_no_fail
from CIME.XML.machines import Machines
from CIME.case import Case

import argparse, math, glob

logger = logging.getLogger(__name__)

###############################################################################
def parse_command_line(args, description):
###############################################################################
    model = CIME.utils.get_model()

    if model == "cesm":
        help_str = \
"""
%s --xml-category [CATEGORY] [--xml-machine ...] [--xml-compiler ...] [ --xml-testlist ...]
OR
%s --help
OR
%s --test

\033[1mEXAMPLES:\033[0m
    \033[1;32m# Run all tests in the xml prealpha category and yellowstone machine \033[0m
    > %s --xml-machine yellowstone --xml-category prealpha
""" % ((os.path.basename(args[0]), ) * 4)

    else:
        help_str = \
"""
%s <TEST|SUITE> [<TEST|SUITE> ...] [--verbose]
OR
%s --help
OR
%s --test

\033[1mEXAMPLES:\033[0m
    \033[1;32m# Run single test \033[0m
    > %s <TESTNAME>

    \033[1;32m# Run test suite \033[0m
    > %s <SUITE>

    \033[1;32m# Run two tests \033[0m
    > %s <TESTNAME1> <TESTNAME2>

    \033[1;32m# Run two suites \033[0m
    > %s <SUITE1> <SUITE2>

    \033[1;32m# Run all tests in a suite except for one \033[0m
    > %s <SUITE> ^<TESTNAME>

    \033[1;32m# Run all tests in a suite except for tests that are in another suite \033[0m
    > %s <SUITE1> ^<SUITE2>
""" % ((os.path.basename(args[0]), ) * 9)

    parser = argparse.ArgumentParser(usage=help_str,
                                     description=description,
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    CIME.utils.setup_standard_logging_options(parser)

    parser.add_argument("--no-run", action="store_true",
                        help="Do not run generated tests")

    parser.add_argument("--no-build", action="store_true",
                        help="Do not build generated tests, implies --no-run")

    parser.add_argument("--no-setup", action="store_true",
                        help="Do not setup generated tests, implies --no-build and --no-run")

    parser.add_argument("-u", "--use-existing", action="store_true",
                        help="Use pre-existing case directories they will pick up at the latest PEND state or re-run the first failed state. Requires test-id")

    parser.add_argument("--save-timing", action="store_true",
                        help="Enable archiving of performance data.")

    parser.add_argument("--no-batch", action="store_true",
                        help="Do not submit jobs to batch system, run locally."
                        " If false, will default to machine setting.")

    parser.add_argument("--single-submit", action="store_true",
                        help="Use a single interactive allocation to run all the tests. "
                        "Can drastically reduce queue waiting. Only makes sense on batch machines.")

    parser.add_argument("-r", "--test-root",
                        help="Where test cases will be created."
                        " Will default to output root as defined in the config_machines file")

    parser.add_argument("--output-root",
                        help="Where the case output is written.")

    parser.add_argument("--baseline-root",
                        help="Specifies an root directory for baseline"
                        "datasets used for Bit-for-bit generate/compare"
                        "testing.")

    parser.add_argument("--clean", action="store_true",
                        help="Specifies if tests should be cleaned after run. If set, "
                        "all object executables, and data files will"
                        " be removed after tests are run")

    parser.add_argument("-m", "--machine",
                        help="The machine for which to build tests, this machine must be defined"
                        " in the config_machines.xml file for the given model. "
                        "Default is to match the name of the machine in the test name or "
                        "the name of the machine this script is run on to the "
                        "NODENAME_REGEX field in config_machines.xml. This option is highly "
                        "unsafe and should only be used if you know what you're doing.")

    parser.add_argument("--mpilib",
                        help="Specify the mpilib. "
                        "To see list of supported mpilibs for each machine, use the utility query_config in this directory. "
                        "The default is the first listing in MPILIBS in config_machines.xml")

    if model == "cesm":
        parser.add_argument("-c", "--compare",
                            help="While testing, compare baselines"
                            "  against the given compare directory ")

        parser.add_argument("-g", "--generate",
                            help="While testing, generate baselines"
                            " to the given generate directory; "
                            "this can also be done after the fact with bless_test_results")

        parser.add_argument("--xml-machine",
                            help="Use this machine key in the lookup in testlist.xml, default is all if any --xml- argument is used")

        parser.add_argument("--xml-compiler",
                            help="Use this compiler key in the lookup in testlist.xml, default is all if any --xml- argument is used")

        parser.add_argument("--xml-category",
                            help="Use this category key in the lookup in testlist.xml, default is all if any --xml- argument is used")

        parser.add_argument("--xml-testlist",
                            help="Use this testlist to lookup tests, default specified in config_files.xml")

        parser.add_argument("testargs", nargs="*",
                            help="Tests or test suites to run."
                            " Testname form is TEST.GRID.COMPSET[.MACHINE_COMPILER]")

    else:

        parser.add_argument("testargs", nargs="+",
                            help="Tests or test suites to run."
                            " Testname form is TEST.GRID.COMPSET[.MACHINE_COMPILER]")

        parser.add_argument("-b", "--baseline-name",
                            help="If comparing or generating baselines,"
                            " use this directory under baseline root. "
                            "Default will be current branch name.")

        parser.add_argument("-c", "--compare", action="store_true",
                            help="While testing, compare baselines")

        parser.add_argument("-g", "--generate", action="store_true",
                            help="While testing, generate baselines; "
                            "this can also be done after the fact with bless_test_results")

    parser.add_argument("--compiler",
                        help="Compiler to use to build cime.  Default will be the name in"
                        " the Testnames or the default defined for the machine.")

    parser.add_argument("-n", "--namelists-only", action="store_true",
                        help="Only perform namelist actions for tests")

    parser.add_argument("-p", "--project",
                        help="Specify a project id for the case (optional)."
                        "Used for accounting when on a batch system."
                        "The default is user-specified environment variable PROJECT")

    parser.add_argument("-t", "--test-id",
                        help="Specify an 'id' for the test. This is simply a"
                        "string that is appended to the end of a test name."
                        "If no testid is specified, then a time stamp plus random number"
                        "will be used.")

    parser.add_argument("-j", "--parallel-jobs", type=int, default=None,
                        help="Number of tasks create_test should perform simultaneously. Default "
                        "will be min(num_cores, num_tests).")

    parser.add_argument("--proc-pool", type=int, default=None,
                        help="The size of the processor pool that create_test can use. Default "
                        "is PES_PER_NODE + 25 percent.")

    parser.add_argument("--walltime", default=os.getenv("CIME_GLOBAL_WALLTIME"),
                        help="Set the wallclock limit for all tests in the suite. "
                        "Can use env var CIME_GLOBAL_WALLTIME to set this for all test.")

    parser.add_argument("-q", "--queue", default=None,
                        help="Force batch system to use a certain queue")

    parser.add_argument("-f", "--testfile",
                        help="A file containing an ascii list of tests to run")

    parser.add_argument("-o", "--allow-baseline-overwrite", action="store_true",
                        help="If the --generate option is given, then by default "
                        "an attempt to overwrite an existing baseline directory "
                        "will raise an error. Specifying this option allows "
                        "existing baseline directories to be silently overwritten.")

    parser.add_argument("--wait", action="store_true",
                        help="On batch systems, wait for submitted jobs to complete")

    parser.add_argument("--force-procs", type=int, default=None,
                        help="For all tests to run with this number of processors")

    parser.add_argument("--force-threads", type=int, default=None,
                        help="For all tests to run with this number of threads")

    parser.add_argument("-i", "--input-dir",
                        help="Use a non-default location for input files")

    args = CIME.utils.parse_args_and_handle_standard_logging_options(args, parser)

    # generate and compare flags may not point to the same directory
    if model == "cesm":
        if args.generate is not None:
            expect(not (args.generate == args.compare),
                   "Cannot generate and compare baselines at the same time")

        if args.xml_testlist is not None:
            expect(not (args.xml_machine is None and args.xml_compiler
                        is  None and args.xml_category is None),
                   "If an xml-testlist is present at least one of --xml-machine, "
                   "--xml-compiler, --xml-category must also be present")

    else:
        expect(not (args.baseline_name is not None and (not args.compare and not args.generate)),
               "Provided baseline name but did not specify compare or generate")
        expect(not (args.compare and args.generate),
               "Tried to compare and generate at same time")

    expect(not (args.namelists_only and not (args.generate or args.compare)),
           "Must provide either --compare or --generate with --namelists-only")

    if args.parallel_jobs is not None:
        expect(args.parallel_jobs > 0,
               "Invalid value for parallel_jobs: %d" % args.parallel_jobs)

    if args.use_existing:
        expect(args.test_id is not None, "Must provide test-id of pre-existing cases")

    if args.no_setup:
        args.no_build = True

    if args.no_build:
        args.no_run = True

    # Namelist-only forces some other options:
    if args.namelists_only:
        expect(not args.no_setup, "Cannot compare namelists without setup")
        args.no_build = True
        args.no_run   = True
        args.no_batch = True

    if args.single_submit:
        expect(not args.no_run, "Doesn't make sense to request single-submit if no-run is on")
        args.no_build = True
        args.no_run   = True
        args.no_batch = True

    if args.test_id is None:
        args.test_id = "%s_%s"%(CIME.utils.get_timestamp(), CIME.utils.id_generator())
    else:
        expect(CIME.utils.check_name(args.test_id, additional_chars="."),
               "invalid test-id argument provided")

    if args.testfile is not None:
        with open(args.testfile, "r") as fd:
            args.testargs.extend( [line.strip() for line in fd.read().splitlines() if line.strip()] )

    # Compute list of fully-resolved test_names
    test_extra_data = {}
    if model == "cesm":
        machine_name = args.xml_machine if args.machine is None else args.machine

        # If it's still unclear what machine to use, look at test names
        if machine_name is None:
            for test in args.testargs:
                testsplit = CIME.utils.parse_test_name(test)
                if testsplit[4] is not None:
                    if machine_name is None:
                        machine_name = testsplit[4]
                    else:
                        expect(machine_name == testsplit[4],
                               "ambiguity in machine, please use the --machine option")

        mach_obj = Machines(machine=machine_name)
        if args.testargs:
            args.compiler = mach_obj.get_default_compiler() if args.compiler is None else args.compiler
            test_names = update_acme_tests.get_full_test_names(args.testargs,
                                                               mach_obj.get_machine_name(), args.compiler)
        else:
            expect(not (args.xml_machine is None and args.xml_compiler
                        is  None and args.xml_category is None and args.xml_testlist is None),
                   "At least one of --xml-machine, --xml-testlist, "
                   "--xml-compiler, --xml-category or a valid test name must be provided.")

            test_data = CIME.test_utils.get_tests_from_xml(args.xml_machine, args.xml_category,
                                                           args.xml_compiler, args.xml_testlist,
                                                           machine_name, args.compiler)
            test_names = [item["name"] for item in test_data]
            for test_datum in test_data:
                test_extra_data[test_datum["name"]] = test_datum

        logger.info("Testnames: %s" % test_names)
    else:
        if args.machine is None:
            args.machine = update_acme_tests.infer_machine_name_from_tests(args.testargs)

        mach_obj = Machines(machine=args.machine)
        args.compiler = mach_obj.get_default_compiler() if args.compiler is None else args.compiler

        test_names = update_acme_tests.get_full_test_names(args.testargs, mach_obj.get_machine_name(), args.compiler)

    expect(mach_obj.is_valid_compiler(args.compiler),
           "Compiler %s not valid for machine %s" % (args.compiler, mach_obj.get_machine_name()))

    # Normalize compare/generate between the models
    baseline_cmp_name = None
    baseline_gen_name = None
    if args.compare or args.generate:
        if model == "cesm":
            if args.compare is not None:
                baseline_cmp_name = args.compare
            if args.generate is not None:
                baseline_gen_name = args.generate
        else:
            baseline_name = args.baseline_name if args.baseline_name else CIME.utils.get_current_branch(repo=CIME.utils.get_cime_root())
            expect(baseline_name is not None,
                   "Could not determine baseline name from branch, please use -b option")
            baseline_name = os.path.join(args.compiler, baseline_name)
            if args.compare:
                baseline_cmp_name = baseline_name
            elif args.generate:
                baseline_gen_name = baseline_name

    if args.input_dir is not None:
        args.input_dir = os.path.abspath(args.input_dir)

    return test_names, test_extra_data, args.compiler, mach_obj.get_machine_name(), args.no_run, args.no_build, args.no_setup, args.no_batch,\
        args.test_root, args.baseline_root, args.clean, baseline_cmp_name, baseline_gen_name, \
        args.namelists_only, args.project, args.test_id, args.parallel_jobs, args.walltime, \
        args.single_submit, args.proc_pool, args.use_existing, args.save_timing, args.queue, \
        args.allow_baseline_overwrite, args.output_root, args.wait, args.force_procs, args.force_threads, args.mpilib, args.input_dir

###############################################################################
def single_submit_impl(machine_name, test_id, proc_pool, project, args, job_cost_map, wall_time, test_root):
###############################################################################
    mach = Machines(machine=machine_name)
    expect(mach.has_batch_system(), "Single submit does not make sense on non-batch machine '%s'" % mach.get_machine_name())

    machine_name = mach.get_machine_name()

    if project is None:
        project = CIME.utils.get_project(mach)
        if project is None:
            project = mach.get_value("PROJECT")

    #
    # Compute arg list for second call to create_test
    #
    new_args = list(args)
    new_args.remove("--single-submit")
    new_args.append("--no-batch")
    new_args.append("--use-existing")
    no_arg_is_a_test_id_arg = True
    no_arg_is_a_proc_pool_arg = True
    no_arg_is_a_machine_arg = True
    for arg in new_args:
        if arg == "-t" or arg.startswith("--test-id"):
            no_arg_is_a_test_id_arg = False
        elif arg.startswith("--proc-pool"):
            no_arg_is_a_proc_pool_arg = False
        elif arg == "-m" or arg.startswith("--machine"):
            no_arg_is_a_machine_arg = True

    if no_arg_is_a_test_id_arg:
        new_args.append("-t %s" % test_id)
    if no_arg_is_a_proc_pool_arg:
        new_args.append("--proc-pool %d" % proc_pool)
    if no_arg_is_a_machine_arg:
        new_args.append("-m %s" % machine_name)

    #
    # Resolve batch directives manually. There is currently no other way
    # to do this without making a Case object. Make a throwaway case object
    # to help us here.
    #
    testcase_dirs = glob.glob("%s/*%s*/TestStatus" % (test_root, test_id))
    expect(testcase_dirs, "No test case dirs found!?")
    first_case = os.path.abspath(os.path.dirname(testcase_dirs[0]))
    with Case(first_case, read_only=False) as case:
        env_batch = case.get_env("batch")

        directives = env_batch.get_batch_directives(case, "case.run", raw=True)
        submit_cmd  = env_batch.get_value("batch_submit", subgroup=None)
        submit_args = env_batch.get_submit_args(case, "case.run")

    tasks_per_node = mach.get_value("PES_PER_NODE")
    num_nodes = int(math.ceil(float(proc_pool) / tasks_per_node))
    if wall_time is None:
        wall_time = compute_total_time(job_cost_map, proc_pool)
        wall_time_bab = convert_to_babylonian_time(int(wall_time))
    else:
        wall_time_bab = wall_time

    queue = env_batch.select_best_queue(proc_pool, wall_time_bab)
    wall_time_max_bab = env_batch.get_queue_specs(queue)[3]
    if wall_time_max_bab is not None:
        wall_time_max = convert_to_seconds(wall_time_max_bab)
        if wall_time_max < wall_time:
            wall_time = wall_time_max
            wall_time_bab = convert_to_babylonian_time(wall_time)

    job_id = "create_test_single_submit_%s" % test_id
    directives = directives.replace("{{ job_id }}", job_id)
    directives = directives.replace("{{ num_nodes }}", str(num_nodes))
    directives = directives.replace("{{ tasks_per_node }}", str(tasks_per_node))
    directives = directives.replace("{{ ptile }}", str(tasks_per_node))
    directives = directives.replace("{{ totaltasks }}", str(tasks_per_node * num_nodes))

    directives = directives.replace("{{ output_error_path }}", "create_test_single_submit_%s.err" % test_id)
    directives = directives.replace("{{ job_wallclock_time }}", wall_time_bab)
    directives = directives.replace("{{ job_queue }}", queue)
    if project is not None:
        directives = directives.replace("{{ project }}", project)

    expect("{{" not in directives, "Could not resolve all items in directives:\n%s" % directives)

    #
    # Make simple submit script and submit
    #

    script = "#! /bin/bash\n"
    script += "\n%s" % directives
    script += "\n"
    script += "cd %s\n"%os.getcwd()
    script += "%s %s\n" % (__file__, " ".join(new_args))

    submit_cmd = "%s %s" % (submit_cmd, submit_args)
    logger.info("Script:\n%s" % script)

    run_cmd_no_fail(submit_cmd, input_str=script, arg_stdout=None, arg_stderr=None, verbose=True)

###############################################################################
# pragma pylint: disable=protected-access
def create_test(test_names, test_data, compiler, machine_name, no_run, no_build, no_setup, no_batch, test_root,
                baseline_root, clean, baseline_cmp_name, baseline_gen_name, namelists_only, project, test_id, parallel_jobs,
                walltime, single_submit, proc_pool, use_existing, save_timing, queue, allow_baseline_overwrite, output_root, wait,
                force_procs, force_threads, mpilib, input_dir):
###############################################################################
    impl = TestScheduler(test_names, test_data=test_data,
                         no_run=no_run, no_build=no_build, no_setup=no_setup, no_batch=no_batch,
                         test_root=test_root, test_id=test_id,
                         baseline_root=baseline_root, baseline_cmp_name=baseline_cmp_name,
                         baseline_gen_name=baseline_gen_name,
                         clean=clean, machine_name=machine_name, compiler=compiler,
                         namelists_only=namelists_only,
                         project=project, parallel_jobs=parallel_jobs, walltime=walltime,
                         proc_pool=proc_pool, use_existing=use_existing, save_timing=save_timing,
                         queue=queue, allow_baseline_overwrite=allow_baseline_overwrite,
                         output_root=output_root, force_procs=force_procs, force_threads=force_threads,
                         mpilib=mpilib, input_dir=input_dir)

    success = impl.run_tests(wait=wait)

    if success and single_submit:
        # Get real test root
        test_root = impl._test_root

        job_cost_map = {}
        largest_case = 0
        for test in impl._tests:
            test_dir = impl._get_test_dir(test)
            procs_needed = impl._get_procs_needed(test, RUN_PHASE)
            time_needed = convert_to_seconds(run_cmd_no_fail("./xmlquery JOB_WALLCLOCK_TIME -value -subgroup case.test", from_dir=test_dir))
            job_cost_map[test] = (procs_needed, time_needed)
            if procs_needed > largest_case:
                largest_case = procs_needed

        if proc_pool is None:
            # Based on size of created jobs, choose a reasonable proc_pool. May need to put
            # more thought into this.
            proc_pool = 2 * largest_case

        # Create submit script
        single_submit_impl(machine_name, test_id, proc_pool, project, sys.argv[1:], job_cost_map, walltime, test_root)

    return 0 if success else CIME.utils.TESTS_FAILED_ERR_CODE

###############################################################################
def _main_func(description):
###############################################################################
    if "--test" in sys.argv:
        libroot = CIME.utils.get_python_libs_root()

        CIME.utils.run_cmd_no_fail("PYTHONPATH=%s python -m doctest %s/CIME/test_scheduler.py -v" %
                                   (libroot, libroot), arg_stdout=None, arg_stderr=None)
        return

    test_names, test_data, compiler, machine_name, no_run, no_build, no_setup, no_batch, \
    test_root, baseline_root, clean, baseline_cmp_name, baseline_gen_name, namelists_only, \
    project, test_id, parallel_jobs, walltime, single_submit, proc_pool, use_existing, \
    save_timing, queue, allow_baseline_overwrite, output_root, wait, force_procs, force_threads, mpilib, input_dir \
        = parse_command_line(sys.argv, description)

    sys.exit(create_test(test_names, test_data, compiler, machine_name, no_run, no_build, no_setup, no_batch, test_root,
                         baseline_root, clean, baseline_cmp_name, baseline_gen_name, namelists_only,
                         project, test_id, parallel_jobs, walltime, single_submit, proc_pool, use_existing, save_timing,
                         queue, allow_baseline_overwrite, output_root, wait, force_procs, force_threads, mpilib, input_dir))

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

if __name__ == "__main__":
    _main_func(__doc__)
