#!/usr/bin/env python3

# Copyright 2021 Stanford University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from __future__ import print_function
import argparse, datetime, glob, json, multiprocessing, os, platform, shlex, shutil, subprocess, sys, traceback, tempfile
import signal

make_exe = os.environ.get('MAKE', 'make')

# Find physical core count of the machine.
if platform.system() == 'Linux':
    lines = subprocess.check_output(['lscpu', '--parse=core']).decode('utf-8')
    physical_cores = len(set(line for line in lines.strip().split('\n')
                             if not line.startswith('#')))
elif platform.system() == 'Darwin':
    physical_cores = int(
        subprocess.check_output(['sysctl', '-n', 'hw.physicalcpu']).decode('utf-8'))
elif platform.system() == 'FreeBSD':
    physical_cores = int(
        subprocess.check_output(['sysctl', '-n', 'hw.ncpu']).decode('utf-8'))
    make_exe = os.environ.get('MAKE', 'gmake')  # default needs to be GNU make
else:
    raise Exception('Unknown platform: %s' % platform.system())

# Choose a reasonable number of application cores given the
# available physical cores.
app_cores = max(physical_cores - 2, 1)

legion_cxx_tests = [
    # Tutorial
    ['tutorial/00_hello_world/hello_world', []],
    ['tutorial/01_tasks_and_futures/tasks_and_futures', []],
    ['tutorial/02_index_tasks/index_tasks', []],
    ['tutorial/03_global_vars/global_vars', []],
    ['tutorial/04_logical_regions/logical_regions', []],
    ['tutorial/05_physical_regions/physical_regions', []],
    ['tutorial/06_privileges/privileges', []],
    ['tutorial/07_partitioning/partitioning', []],
    ['tutorial/08_multiple_partitions/multiple_partitions', []],
    ['tutorial/09_custom_mapper/custom_mapper', []],

    # Examples
    ['examples/circuit/circuit', []],
    ['examples/dynamic_registration/dynamic_registration', []],
    ['examples/ghost/ghost', ['-ll:cpu', '4']],
    ['examples/ghost_pull/ghost_pull', ['-ll:cpu', '4']],
    ['examples/realm_saxpy/realm_saxpy', []],
    ['examples/realm_stencil/realm_stencil', ['-ll:cpu', '4']],
    ['examples/spmd_cgsolver/spmd_cgsolver', ['-ll:cpu', '4', '-perproc']],
    ['examples/virtual_map/virtual_map', []],
    ['examples/attach_2darray_c_fortran_layout/attach_2darray', []],
    ['examples/attach_array_daxpy/attach_array_daxpy', []],
    ['examples/attach_index_space/index_space_attach', []],
    ['examples/predication/predication', []],
    ['examples/layout_constraints/transpose', []],
    ['examples/inline_tasks/inline_tasks', []],
    ['examples/local_function_tasks/local_function_tasks', []],
    # Comment this test out until it works everywhere
    #['examples/implicit_top_task/implicit_top_task', []],

    # Tests
    ['test/rendering/rendering', ['-i', '2', '-n', '64', '-ll:cpu', '4']],
    ['test/legion_stl/test_stl', []],
]

legion_fortran_tests = [
    ['tutorial/fortran/00_hello_world/hello_world_fortran', []],
    ['tutorial/fortran/01_tasks_and_futures/tasks_and_futures_fortran', []],
    ['tutorial/fortran/02_index_tasks/index_tasks_fortran', []],
    ['tutorial/fortran/03_physical_regions/physical_regions_fortran', []],
    ['tutorial/fortran/04_privileges_accessor/privileges_accessor_fortran', []],
    ['tutorial/fortran/05_privileges_raw_ptr/privileges_raw_ptr_fortran', []],
    ['tutorial/fortran/06_partitioning/partitioning_fortran', []],
    ['tutorial/fortran/07_partitioning_fortran_task/partitioning_fortran_task_fortran', []],
    ['tutorial/fortran/08_multiple_partitions/multiple_partitions_fortran', []],
    ['tutorial/fortran/09_region_2d/region_2d_fortran', []],
    ['tutorial/fortran/10_attach_array/attach_array_fortran', []],
]

if platform.system() != 'Darwin':
    legion_cxx_tests += [
        # FIXME: Fails non-deterministically on Mac OS: https://github.com/StanfordLegion/legion/issues/213
        ['test/attach_file_mini/attach_file_mini', []],
    ]

legion_network_cxx_tests = [
    # Examples
    ['examples/mpi_interop/mpi_interop', []],

    # Tests
    ['test/bug954/bug954', ['-ll:rsize', '1024']],
]

legion_openmp_cxx_tests = [
    # Examples
    ['examples/omp_saxpy/omp_saxpy', ['-ll:ocpu', '1']],
]

legion_kokkos_cxx_tests = [
    # Examples
    ['examples/kokkos_saxpy/kokkos_saxpy', []],
]

legion_python_cxx_tests = [
    # Bindings
    ['bindings/python/legion_python', ['examples/domain.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/domain_point.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/domain_transform.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/future.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/hello.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/import.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/index_launch.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/ispace.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/method.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/must_epoch.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/partition.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/partition_by_field.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/partition_by_image.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/partition_by_image_range.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/partition_by_preimage.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/partition_by_preimage_range.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/partition_by_restriction.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/reduction.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/region.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/region_fields.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/return_region.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/single_launch.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/struct.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/trace.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/tunable.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['examples/types.py', '-ll:py', '1', '-ll:cpu', '0']],

    ['bindings/python/legion_python', ['tests/fail/privileges.py', '-ll:py', '1', '-ll:cpu', '0']],

    ['bindings/python/legion_python', ['tests/pass/copy.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['tests/pass/empty_region.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['tests/pass/print_once.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['tests/pass/privileges.py', '-ll:py', '1', '-ll:cpu', '0']],
    ['bindings/python/legion_python', ['tests/pass/no_access.py', '-ll:py', '1', '-ll:cpu', '0']],

    # Examples
    ['examples/python_interop/python_interop', ['-ll:py', '1']],

    # Tests
    ['test/python_bindings/python_bindings', ['-ll:py', '1', '-ll:cpu', '0']],
]

legion_hdf_cxx_tests = [
    # Tests
    ['test/hdf_attach_subregion_parallel/hdf_attach_subregion_parallel', ['-ll:cpu', '4']],
]

if platform.system() != 'Darwin':
    legion_hdf_cxx_tests += [
        # FIXME: Fails non-deterministically on Mac OS: https://github.com/StanfordLegion/legion/issues/213
        ['examples/attach_file/attach_file', ['-h', 'data.h5', '-d', '/path/to/data']],
    ]

def get_legion_cxx_perf_tests(nodes, cores_per_node):
    return [
        # Circuit: Heavy Compute
        ['examples/circuit/circuit',
         ['-l', '10', '-p', str(cores_per_node * nodes), '-npp', '2500', '-wpp', '10000',
          '-ll:csize', '8192', '-ll:cpu', str(cores_per_node)]],

        # Circuit: Light Compute
        ['examples/circuit/circuit',
         ['-l', '10', '-p', '100', '-npp', '2', '-wpp', '4', '-ll:cpu', '2']],
    ]

def get_regent_perf_tests(nodes, cores_per_node):
    return [
        # Circuit: Heavy Compute
        ['language/examples/circuit_sparse.rg',
         ['-l', '10', '-p', str(nodes * cores_per_node), '-npp', '2500', '-wpp', '10000',
          '-ll:csize', '8192', '-ll:cpu', str(cores_per_node), '-fflow-spmd-shardsize', str(cores_per_node)]],

        # Circuit: Light Compute
        ['language/examples/circuit_sparse.rg',
         ['-l', '10', '-p', '100', '-npp', '2', '-wpp', '4', '-ll:cpu', '2',
          '-fflow-spmd-shardsize', '2']],

        # PENNANT: Heavy Compute
        ['language/examples/pennant_fast.rg',
         ['pennant.tests/sedovbig3x30/sedovbig.pnt',
          '-seq_init', '0', '-par_init', '1', '-print_ts', '1', '-prune', '5',
          '-npieces', str(nodes * cores_per_node), '-numpcx', '1', '-numpcy', str(nodes * cores_per_node),
          '-ll:csize', '8192', '-ll:cpu', str(cores_per_node), '-fflow-spmd-shardsize', str(cores_per_node),
          '-fvectorize-unsafe', '1']],
    ]

class TestTimeoutException(Exception):
    pass

def sigalrm_handler(signum, frame):
    raise TestTimeoutException

def cmd(command, env=None, cwd=None, timelimit=None):
    print(' '.join(command))
    sys.stdout.flush()  # python 2 doesn't have flush option in print
    if timelimit:
        child = subprocess.Popen(command, env=env, cwd=cwd)
        signal.signal(signal.SIGALRM, sigalrm_handler)
        signal.alarm(timelimit)
        try:
            ret = child.wait()
            signal.alarm(0)  # disable alarm
            signal.signal(signal.SIGALRM, signal.SIG_DFL)
            if ret:
                raise subprocess.CalledProcessError(ret, command)
            return ret
        except TestTimeoutException:
            child.kill()
            raise  # re-raise
    else:
        return subprocess.check_call(command, env=env, cwd=cwd)

def run_test_regent(launcher, root_dir, tmp_dir, bin_dir, env, thread_count):
    cmd([sys.executable, os.path.join(root_dir, 'language/travis.py')], env=env)

def run_cxx(tests, flags, launcher, root_dir, bin_dir, env, thread_count, timelimit):
    prev_built_dir = None
    for test_file, test_flags in tests:
        test_dir = os.path.dirname(os.path.join(root_dir, test_file))
        if bin_dir:
            test_path = os.path.join(bin_dir, os.path.basename(test_file))
        else:
            test_path = os.path.join(root_dir, test_file)
            # build if this is in a new directory
            if test_dir != prev_built_dir:
                # and clean up the previous directory to keep disk usage down
                if prev_built_dir:
                    cmd(['find', prev_built_dir , '-type', 'f', '(', '-name', '*.a', '-o', '-name', os.path.basename(test_file), ')', '-exec', 'rm', '-v', '{}', ';'])
                cmd([make_exe, '-C', test_dir, '-j', str(thread_count)], env=env)
                prev_built_dir = test_dir
        cmd(launcher + [test_path] + flags + test_flags, env=env, cwd=test_dir, timelimit=timelimit)
    if prev_built_dir:
        cmd(['find', prev_built_dir , '-type', 'f', '(', '-name', '*.a', '-o', '-name', os.path.basename(test_file), ')', '-exec', 'rm', '-v', '{}', ';'])

def run_regent(tests, flags, launcher, root_dir, env, thread_count, timelimit):
    for test_file, test_flags in tests:
        test_dir = os.path.dirname(os.path.join(root_dir, test_file))
        test_path = os.path.join(root_dir, test_file)
        cmd(launcher + [test_path] + flags + test_flags, env=env, cwd=test_dir, timelimit=timelimit)

def precompile_regent(tests, flags, launcher, root_dir, env, thread_count):
    exe_tests = []
    for test_file, test_flags in tests:
        test_dir = os.path.dirname(os.path.join(root_dir, test_file))
        test_path = os.path.join(root_dir, test_file)

        exe = os.path.splitext(test_path)[0] + '.exe'
        env = dict(list(env.items()) + [('OBJNAME', exe)])

        cmd(launcher + [test_path] + flags + test_flags, env=env, cwd=test_dir)

        exe_tests.append([exe, test_flags])
    return exe_tests

def run_test_legion_cxx(launcher, root_dir, tmp_dir, bin_dir, env, thread_count, timelimit):
    flags = ['-logfile', 'out_%.log']
    if env['USE_CUDA'] == '1':
        flags.extend(['-ll:gpu', '1'])
    if (env['USE_KOKKOS'] == '1') and (env['USE_OPENMP'] == '1'):
        flags.extend(['-ll:ocpu', '1', '-ll:onuma', '0' ])
    run_cxx(legion_cxx_tests, flags, launcher, root_dir, bin_dir, env, thread_count, timelimit)

def run_test_legion_network_cxx(launcher, root_dir, tmp_dir, bin_dir, env, thread_count, timelimit):
    flags = ['-logfile', 'out_%.log']
    if env['USE_CUDA'] == '1':
        flags.extend(['-ll:gpu', '1'])
    if (env['USE_KOKKOS'] == '1') and (env['USE_OPENMP'] == '1'):
        flags.extend(['-ll:ocpu', '1', '-ll:onuma', '0' ])
    run_cxx(legion_network_cxx_tests, flags, launcher, root_dir, bin_dir, env, thread_count, timelimit)

def run_test_legion_openmp_cxx(launcher, root_dir, tmp_dir, bin_dir, env, thread_count, timelimit):
    flags = ['-logfile', 'out_%.log']
    if env['USE_CUDA'] == '1':
        flags.extend(['-ll:gpu', '1'])
    if (env['USE_KOKKOS'] == '1') and (env['USE_OPENMP'] == '1'):
        flags.extend(['-ll:ocpu', '1', '-ll:onuma', '0' ])
    run_cxx(legion_openmp_cxx_tests, flags, launcher, root_dir, bin_dir, env, thread_count, timelimit)

def run_test_legion_kokkos_cxx(launcher, root_dir, tmp_dir, bin_dir, env, thread_count, timelimit):
    flags = ['-logfile', 'out_%.log']
    if env['USE_CUDA'] == '1':
        flags.extend(['-ll:gpu', '1'])
    if (env['USE_KOKKOS'] == '1') and (env['USE_OPENMP'] == '1'):
        flags.extend(['-ll:ocpu', '1', '-ll:onuma', '0' ])
    run_cxx(legion_kokkos_cxx_tests, flags, launcher, root_dir, bin_dir, env, thread_count, timelimit)

def run_test_legion_python_cxx(launcher, root_dir, tmp_dir, bin_dir, env, thread_count, timelimit):
    # Hack: legion_python currently requires the module name to come first
    flags = [] # ['-logfile', 'out_%.log']
    python_dir = os.path.join(root_dir, 'bindings', 'python')
    # Hack: Fix up the environment so that Python can find all the examples.
    env = dict(list(env.items()) + [
        ('PYTHONPATH', ':'.join([python_dir])),
        ('LD_LIBRARY_PATH', ':'.join([python_dir])),
    ])
    # Clean up around python because we are going to make shared objects
    # which is not something that anyone else does
    cmd([make_exe, '-C', python_dir, 'clean'], env=env)
    run_cxx(legion_python_cxx_tests, flags, launcher, root_dir, bin_dir, env, thread_count, timelimit)
    cmd([make_exe, '-C', python_dir, 'clean'], env=env)

def run_test_legion_hdf_cxx(launcher, root_dir, tmp_dir, bin_dir, env, thread_count, timelimit):
    flags = ['-logfile', 'out_%.log']
    if env['USE_CUDA'] == '1':
        flags.extend(['-ll:gpu', '1'])
    if (env['USE_KOKKOS'] == '1') and (env['USE_OPENMP'] == '1'):
        flags.extend(['-ll:ocpu', '1', '-ll:onuma', '0' ])
    run_cxx(legion_hdf_cxx_tests, flags, launcher, root_dir, bin_dir, env, thread_count, timelimit)

def run_test_legion_fortran(launcher, root_dir, tmp_dir, bin_dir, env, thread_count, timelimit):
    flags = ['-logfile', 'out_%.log']
    if env['USE_CUDA'] == '1':
        flags.extend(['-ll:gpu', '1'])
    run_cxx(legion_fortran_tests, flags, launcher, root_dir, bin_dir, env, thread_count, timelimit)

def run_test_fuzzer(launcher, root_dir, tmp_dir, bin_dir, env, thread_count):
    env = dict(list(env.items()) + [('WARN_AS_ERROR', '0')])
    fuzz_dir = os.path.join(tmp_dir, 'fuzz-tester')
    cmd(['git', 'clone', 'https://github.com/StanfordLegion/fuzz-tester', fuzz_dir])
    # TODO; Merge deppart branch into master after this makes it to stable Legion branch
    cmd(['git', 'checkout', 'deppart'], cwd=fuzz_dir)
    cmd(['python3', 'main.py'], env=env, cwd=fuzz_dir)

def run_test_realm(launcher, root_dir, tmp_dir, bin_dir, env, thread_count, timelimit):
    test_dir = os.path.join(root_dir, 'test/realm')
    cmd([make_exe, '-C', test_dir, 'DEBUG=0', 'clean'], env=env)
    cmd([make_exe, '-C', test_dir, 'DEBUG=0', '-j', str(thread_count), 'build'], env=env)
    cmd([make_exe, '-C', test_dir, 'DEBUG=0', 'run_all'], env=env, timelimit=timelimit)

def run_test_external1(launcher, root_dir, tmp_dir, bin_dir, env, thread_count, timelimit):
    flags = ['-logfile', 'out_%.log']

    # Realm perf test (move back to perf test when integrated with perf.py)
    perf_dir = os.path.join(root_dir, 'test/performance/realm')
    cmd([make_exe, '-C', perf_dir, 'DEBUG=0', 'clean_all'], env=env)
    cmd([make_exe, '-C', perf_dir, 'DEBUG=0', 'build_all'], env=env)
    cmd([make_exe, '-C', perf_dir, 'DEBUG=0', 'RUNMODE=short', 'run_all'], env=env, timelimit=timelimit)

    # Fast Direct Solver
    # Contact: Chao Chen <cchen10@stanford.edu>
    solver_dir = os.path.join(tmp_dir, 'fastSolver2')
    cmd(['git', 'clone', 'https://github.com/Charles-Chao-Chen/fastSolver2.git', solver_dir])
    # cmd(['git', 'checkout', '4c7a59de63dd46a0abcc7f296fa3b0f511e5e6d2', ], cwd=solver_dir)
    solver = [[os.path.join(solver_dir, 'spmd_driver/solver'),
        ['-machine', '1', '-core', '8', '-mtxlvl', '6', '-ll:cpu', '8', '-ll:csize', '1024']]]
    run_cxx(solver, flags, launcher, root_dir, None, env, thread_count, timelimit)

    # Parallel Research Kernels: Stencil
    # Contact: Wonchan Lee <wonchan@cs.stanford.edu>
    prk_dir = os.path.join(tmp_dir, 'prk')
    cmd(['git', 'clone', 'https://github.com/magnatelee/PRK.git', prk_dir])
    # This uses a custom Makefile that requires additional
    # configuration. Rather than go to that trouble it's easier to
    # just use a copy of the standard Makefile template.
    stencil_dir = os.path.join(prk_dir, 'LEGION', 'Stencil')
    stencil_env = dict(list(env.items()) + [
        ('OUTFILE', 'stencil'),
        ('GEN_SRC', 'stencil.cc'),
        ('CXXFLAGS', (env['CXXFLAGS'] if 'CXXFLAGS' in env else '') +
         ' -DRADIUS=2 -DRESTRICT_KEYWORD -DDISABLE_BARRIER_MIGRATION'),
    ])
    makefile = os.path.join(root_dir, 'apps/Makefile.template')
    cmd([make_exe, '-f', makefile, '-C', stencil_dir, '-j', str(thread_count)], env=stencil_env)
    stencil = os.path.join(stencil_dir, 'stencil')
    # HACK: work around stencil mapper issue with -ll:ext_sysmem 0
    cmd([stencil, '4', '10', '1000', '-ll:ext_sysmem', '0'], timelimit=timelimit)

    # SNAP
    # Contact: Mike Bauer <mbauer@nvidia.com>
    snap_dir = os.path.join(tmp_dir, 'snap')
    cmd(['git', 'clone', 'https://github.com/StanfordLegion/Legion-SNAP.git', snap_dir])
    # This can't handle flags before application arguments, so place
    # them after.
    snap = [[os.path.join(snap_dir, 'src/snap'),
             [os.path.join(snap_dir, 'input/mms.in')] + flags]]
    run_cxx(snap, [], launcher, root_dir, None, env, thread_count, timelimit)

    # Soleil-X
    # Contact: Manolis Papadakis <mpapadak@stanford.edu>
    soleil_dir = os.path.join(tmp_dir, 'soleil-x')
    cmd(['git', 'clone', 'https://github.com/stanfordhpccenter/soleil-x.git', soleil_dir])
    soleil_env = dict(list(env.items()) + [
        ('LEGION_DIR', root_dir),
        ('SOLEIL_DIR', soleil_dir),
        ('CC', 'gcc'),