@CHEMSH_BASH_HEADER@#!@PYTHON_EXECUTABLE@
'''Control script for ChemShell executable:
   Sets up Python 3 environment and calls the chemsh.x binary
'''

__author__  = ('You Lu <you.lu@stfc.ac.uk>')
__version__ = '23.0'

# October 2016


def parseArguments():
    ''''''

    from argparse import ArgumentParser, REMAINDER, RawTextHelpFormatter, SUPPRESS
    from os       import path
    from sys      import argv

    # file containing saved command line arguments starts with '@'
    parser = ArgumentParser(prog='chemsh',
                            description='ChemShell script',
                            fromfile_prefix_chars='@',
                            formatter_class=RawTextHelpFormatter)

    # platform options
    platforms = parser.add_argument_group('platform options')

    # plotting options
    plotting = parser.add_argument_group('plotting options')

    parser.add_argument('-v', '--version',
                         action='version',
                         default=__version__,
                         version='%(prog)s: ChemShell version {}'.format(__version__))

    parser.add_argument('--submit',
                         action='store_true',
                         default=False,
                         help='submit the job on to the compute nodes [default: %(default)s]')

    parser.add_argument('--exclusive',
                         action='store_true',
                         default=True,
                         help='exclusive node use, to avoid sharing with other jobs [default: %(default)s]')

    # TWK 17/7/23 Path to test driver
    # For internal use only (used by ChemShell test driver in job submission mode)
    # We should not assume anything here about where the test suite is,
    # so we send its path explicitly from the test driver itself.
    parser.add_argument('--test',
                         action='store',
                         default='',
                         help=SUPPRESS)

    # TWK 17/7/23 Test set specification.
    # For internal use only (used by ChemShell test driver in job submission mode)
    parser.add_argument('--sets',
                         action='store',
                         nargs='*',
                         default='',
                         help=SUPPRESS)

    parser.add_argument('-np', '--nprocs',
                         metavar='',
                         action='store',
                         type=int,
                         default=1,
                         help='total number of MPI processes [ default: %(default)d ]')

    parser.add_argument('-nppn', '--nprocs-per-node',
                         metavar='',
                         action='store',
                         type=int,
                         default=0,
                         help='number of MPI processes per node [ default: %(default)d ]')

    parser.add_argument('-c', '--cpus-per-task',
                         metavar='',
                         action='store',
                         type=int,
                         default=0,
                         help='number of CPU cores allocated per MPI process [ default: %(default)d ]')

    parser.add_argument('-nt', '--nthreads',
                         metavar='',
                         action='store',
                         type=int,
                         default=1,
                         help='total number of OpenMP threads [ default: %(default)d ]')

    parser.add_argument('-nn', '--nnodes',
                         metavar='',
                         action='store',
                         type=int,
                         default=1,
                         help='number of compute nodes [ default: %(default)d ]')

    parser.add_argument('-nc', '--ncpus',
                         metavar='',
                         action='store',
                         type=int,
                         default=1,
                         help='number of CPU cores per node [ default: %(default)d ]')

    parser.add_argument('-nwg', '--nworkgroups',
                         metavar='',
                         action='store',
                         type=int,
                         default=1,
                         help='number of workgroups [ default: %(default)d ]')

    parser.add_argument('-npmm', '--nprocs-mm',
                         metavar='',
                         action='store',
                         type=int,
                         default=0,
                         help='number of processes used for MM [ default: %(default)d ]')

    parser.add_argument('-gm', '--ga-memory',
                         metavar='',
                         action='store',
                         type=int,
                         default=1024,
                         help='maximum size of global memory per single process (MPI task) NWChem can allocate, total memory = heap memory + stack memory + global memory (heap = stack = global/2) [ default: %(default)s MB ]')

    parser.add_argument('-dbg', '--debug', '--print-level',
                         metavar='',
                         dest='debug',
                         action='store',
                         type=int,
                         default=0,
                         choices=range(6),
                         help='level of debugging printing [ default: %(default)d, allowed: %(choices)s ]')

    parser.add_argument('--prepend-pythonpath',
                         metavar='',
                         action='store',
                         default='',
                         help='path(s) to additional Python packages to prepend to sys.path at runtime (NB: not the environment variable PYTHONPATH); seperate with : when there are more than one; package paths prepended will override the compiletime --prepend-pythonpath and default system ones if there are duplicates [ default: %(default)s ]')

    parser.add_argument('-pa', '--previous-args',
                         action='store_true',
                         default=False,
                         help='take saved command line arguments from previous job (overridable by new definition) [ default: %(default)r ]')

    parser.add_argument('-w', '--wrapper',
                         metavar='',
                         action='store',
                         default='',
                         help='wrapper command for advanced analysis or debugging (e.g., --wrapper=\'map --profile\' for Allinea MAP) [ default: %(default)s ]')


    # plotting options
    plotting.add_argument('--plot',
                          metavar='',
                          action='store',
                          nargs='?',
                          default=False,
                          help='plot graphs from a data file saved for ChemShell calculation (typically in the compressed NumPy format .npz) [ default: %(default)s ]')

    plotting.add_argument('--units',
                          metavar='',
                          action='store',
                          nargs='+',
                          default=['a.u.','fs'],
                          help='units used for plotting [ default: %(default)s ]')


    # input scripts
    parser.add_argument('filenames',
                         nargs=REMAINDER)


    # platform options
    platforms.add_argument('-pl', '--platform',
                            metavar='',
                            action='store',
                            default='@PLATFORM@',
                            help='choose the platform type [ default: %(default)s ]')

    platforms.add_argument('-A', '--account',
                            metavar='',
                            action='store',
                            default='',
                            help='your budget code (PBS argument -A) [ default: %(default)s ]')

    platforms.add_argument('-J', '--jobname',
                            metavar='',
                            action='store',
                            default='jobname',
                            help='a name for your job to submit (PBS argument -N) [ default: %(default)s ]')

    platforms.add_argument('-N', '--nodename',
                            metavar='',
                            action='store',
                            default='',
                            help='specific node [ default: %(default)s ]')

    platforms.add_argument('-P', '--project',
                            metavar='',
                            action='store',
                            default='',
                            help='project name (used for SLURM argument --reservation) [ default: %(default)s ]')

    # YL 17/02/2021: we start to use the SLURM jargons since more and more HPC machines use SLURM
    platforms.add_argument('-p', '--partition',
                            metavar='',
                            action='store',
                            default='',
                            help='which partition to run on [ default: %(default)s ]')

    # YL 28/11/2021: not pragmatic due to limitation of Python - --arg '-A a -B b' won't work even quoted
    platforms.add_argument('-E', '--extra-scheduler-arg', '--extra-scheduler-argument',
                            metavar='',
                            action='store',
                            default='',
                            help='extra scheduler arguments, for example, --extra-scheduler-arg=--constraint=scarf20:--mail-user=you.lu@stfc.ac.uk where = and : have to be used [ default: %(default)s ]')

    platforms.add_argument('-q', '--qos',
                            metavar='',
                            action='store',
                            default='',
                            help='request a specific quality of service (QoS) (typically used for SLURM) [ default: %(default)s ]')
    
    platforms.add_argument('-cf', '--cpu-freq',
                            metavar='',
                            action='store',
                            type=str,
                            default='High', # YL 19/05/2024: typical values for sbatch are Low, High, HighM1, Medium,...
                            help='request a particular cpu frequency (SLURM argument) [ default: %(default)s (numbers could be accepted) ]')

    platforms.add_argument('-wt', '--walltime',
                            metavar='',
                            action='store',
                            default='00:05:00',
# TODO: convert 1234 m or 1234 min or 123 h or 123 hours to hh:mm:ss
#                            type=
                            help='maximum length of time (PBS argument -l walltime= or SLURM\'s --time=), format: d-hh:mm:ss [ default: %(default)s ]')

    platforms.add_argument('-m', '--mem', '--memory-per-cpu',
                            dest='mem',
                            metavar='',
                            action='store',
                            type=int,
                            default=2,
                            help='memory (RAM) to request per CPU [ default: %(default)s GB ]')

    # YL 04/07/2022: const doesn't work when only --mem-ratio is specified
    platforms.add_argument('-mr', '--mem-ratio',
                            metavar='',
                            nargs='?',
                            action='store',
                            type=float,
                            default=1.0,
                            const=0.25,
                            help='Ratio of memory reserved for ChemShell: typically used for launching subjobs (e.g., srun --oversubscribe --mem=XXXM namd2) [ default: %(default)s ]')

    platforms.add_argument('-s', '--space',
                            metavar='',
                            action='store',
                            default=8,
                            help='space to request for writing out files [ default: %(default)s GB ]')

    platforms.add_argument('--load-modules',
                            metavar='',
                            action='store',
                            default=[],
                            nargs='+',
                            help='modules to load [ default: %(default)s ]')

    platforms.add_argument('--unload-modules',
                            metavar='',
                            action='store',
                            default=[],
                            nargs='+',
                            help='modules to unload [ default: %(default)s ]')

    # TWK 16/11/2023 Allow user-defined environment variables to be provided to SLURM
    # in addition to those set by ChemShell
    platforms.add_argument('--export',
                           metavar='',
                           action='store',
                           default='',
                           help='environment variables to export, separated by commas (e.g. ENV1=val1,ENV2=val2)')

    # parse
    args = parser.parse_args()

    # take command line arguments from previous job (overridable by new definition)
    if args.previous_args:
        # saved command line arguments of previous job in ~/.chemsh
        previous = parser.parse_args(['@'+path.join(path.expanduser('~'), '.chemsh')])
        args = parser.parse_args(namespace=previous)
#        writeLog(args, '_chemsh_run.log')

    # FIXME: currently doesn't save when -pa does not appear
    else:
        # save arguments to file .chemsh under home directory (works for all platforms)
        # should not use bin/intel or bin/gnu because a build on supercomputer is shared by many users
        # YL 03/09/2021: on some supercomputers /home isn't available for compute nodes
        try:
            with open(path.join(path.expanduser('~'), '.chemsh'), 'w') as fp:
                fp.write('\n'.join(argv[1:]))
        except:
            pass

    return args


def setExec(args):
    '''Get the path of the executable binary chemsh.x'''

    from os.path import join, dirname, realpath

# YL: this won't work if we use hybrid script as __file__ will be stdin then
#    args.__CHEMSH_EXEC = realpath(join(dirname(__file__), 'chemsh.x'))
    args.__CHEMSH_EXEC   = '@CHEMSH_EXECUTABLE@'
    args.__CHEMSH_SCRIPT = realpath(join(args.__CHEMSH_EXEC, '..', 'chemsh'))


def importModule(modname, filepath, submodules=None, **kwargs):
    '''Construct a module through its file location'''

    from importlib.util import spec_from_file_location, module_from_spec
    from os.path        import abspath, dirname, join
    from sys            import path

    # so that we can do \`from . import *\` in the package (e.g., in tools/__init__.py)
    pkgdir = abspath(join(dirname(filepath), '..'))
    if pkgdir not in path:
        path.append(pkgdir)

    spec   = spec_from_file_location(modname, filepath, submodule_search_locations=submodules)
    module = module_from_spec(spec)

    # YL 26/02/2021: allowed additional attributes via keyword arguments
    #                it has to be before exec_module() which loads everything
    for k, v in kwargs.items():
        setattr(module, k, v)

    spec.loader.exec_module(module)

    return module


def choosePlatform(args, plot=False):
    ''''''

    from os.path        import join, dirname, realpath
    from importlib.util import spec_from_file_location, module_from_spec

    # hardcode the platform according to the compile argument
    # but allow the user at runtime to override the platform hardwired at compiletime
    args.platform = args.platform.lower() or '@PLATFORM@'.lower()
    args.__import = importModule

# YL: this won't work if we use hybrid script as __file__ will be stdin then
    args.__CHEMSH_ROOT = '@CHEMSH_DIR@'

    # load the platforms module
    platfmodule = importModule(args.platform, join(args.__CHEMSH_ROOT, 'chemsh', 'utils', 'platforms', '__init__.py'))

# FUTURE: moved to branch 'tcl2py' 
#    toolsmodule = importModule('tools', join(CHEMSH_ROOT, 'chemsh', 'utils', 'tools', '__init__.py'))
#    toolsmodule.run(args, CHEMSH_ROOT, importer=importModule)

    # initialise platform
    platfmodule.init(args)

    # a ChemShell preprocessor to get "top-level" parameters and save as environment
    preprocess(args)

    # YL 26/10/2018: insert in rather than replace the existing PYTHONPATH
    # YL 30/10/2018: this works only for the current Python session but not for the Python/C API interpreter
    #sys_path.append(getPath())

    # set args.__CHEMSH_EXEC (chemsh.x)
    setExec(args)

    platfmodule.run(args)


def preprocess(args):
    '''Preprocess the input file to get some global variables'''

#    # GA memory sizes
#    preprocessors = [ 'HEAP', 'STACK', 'GLOBAL' ]
#    try:
#        with open(args.filenames) as fp:
#            for key in preprocessors:
#                for line in fp.readlines():
#                    if key in line:
#                        try:
#                            exec(line)
#                            #environ[key] = str(locals()[key])
#                            setattr(args, key, locals()[key])
#                        except:
#                            pass
#                        continue
#                    else:
#                        setattr(args, key, 0)
#                        continue
#    except:
#        return


# TODO: unify with the writeSetupLog in setup.py
def writeLog(args, filename):
    '''Write setup.py arguments to setup.log'''

    from time    import time as time_time
    from time    import asctime, localtime
    from getpass import getuser

    # special treatment if --previous-args is used

    def _getRunArgs():
        ''''''
   
        from sys import argv
   
        setExec(args)
        
        strbuff = ' '.join([args.__CHEMSH_SCRIPT]+argv[1:])
   
        return strbuff

    with open(filename, 'a') as fp:
        strbuff = "" + 24*"=" + "\n"
        fp.write(strbuff)
        strbuff = "" + asctime(localtime(time_time())) + "\n"
        fp.write(strbuff)
        strbuff = "" + 24*"=" + "\n\n"
        fp.write(strbuff)
        strbuff = "By " + getuser() + "\n\n"
        fp.write(strbuff)
        fp.write(_getRunArgs())
        fp.write('\n\n\n')


def main():
    ''''''

    from os          import environ, path
    from collections import OrderedDict
    import sys

    # insert location of ChemShell site-packages for missing Python dependencies automatically installed at compiletime
    sys.path.insert(1, '@CHEMSH_SITE_PACKAGES@')

    # prepend compiletime additional sys.path (after the 1st element which is normally the working directory)
    pythonpath = '@PYTHONPATH@'.split(':')
    sys.path = sys.path[0:1] + [ p.strip() for p in pythonpath if p ] + sys.path[1:]
    # YL 01/03/2021: we allow duplicates in the above insertion but remove duplicated elements
    #                from the right side so that the user can prioritise a path even if is's
    #                already in sys.path
    sys.path = list(OrderedDict.fromkeys(sys.path))

    # YL 30/10/2018: PYTHONPATH is not good
    environ['PYTHONPATH'] = ''

    args = parseArguments()

    # prepend runtime additional sys.path (after the 1st element which is normally the working directory)
    sys.path = sys.path[0:1] + [ p.strip() for p in args.prepend_pythonpath.split(':') if p ] + sys.path[1:]
    # YL 01/03/2021: we allow duplicates in the above insertion but remove duplicated elements
    #                from the right side so that the user can prioritise a path even if is's
    #                already in sys.path
    sys.path = list(OrderedDict.fromkeys(sys.path))

    # YL 01/10/2020: remove the input script filenames which may be taken as modulenames
    #                by args.load/unload_modules which have the attribute `nargs='+'`
    for name in args.load_modules:
        if path.isfile(name):
            args.load_modules.remove(name)
            args.filenames += [name]
    for name in args.unload_modules:
        if path.isfile(name):
            args.unload_modules.remove(name)
            args.filenames += [name]

    # YL 20/09/2019: always make args.filenames a list to enable sequetial dealing with multiple input scripts
    if not args.filenames:
        args.filenames = ['']

    args.mpiexec = '@MPIEXEC@'

    if not args.previous_args:
        writeLog(args, '_chemsh_run.log')

    # a plotting job
    if args.plot:
        return choosePlatform(args, plot=True)
    
    # select platform to run specified jobs
    choosePlatform(args)


if __name__ == '__main__':

    main()

@CHEMSH_BASH_FOOTER@
