# you.lu@stfc.ac.uk
'''Set up environment modules on STFC Scientific Computing's SCARF cluster.'''

#  Copyright (C) 2017 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

python_module       = 'Python-bundle-PyPI/2023.10-GCCcore-13.2.0'
intel_version       = '2023.2'
intel_subversion    = '2023.2.1'
intel_year          = intel_version.split('.')[0]
python_version      = '3.11'
python_subversion   = '3.11.5'
fftw_version        = '3.3.10'
gcc_version         = '13.2.0'

# YL 17/02/2021: enabled relative importing
from pathlib import Path
from sys     import path, stdout
file = Path(__file__).resolve()
parent, top = file.parent, file.parents[2]
# add top and remove parent
path.append(str(top))
try:
     path.remove(str(parent))
except ValueError:
     pass
import platforms.scarf
__package__ = 'utils.platforms'


def init(args):

    from platform import node
    from os       import environ, path, sep

    toloads = { # YL 21/02/2024: New modules, requires consitent GCC versions 
                'AMD':'AMDmodules', # TD 01/03/24 Most modules hidden behind this 
                'python3': python_module,
                'intel'  :f'intel-compilers/{intel_subversion}',
                'CMake'  :f'CMake/3.27.6-GCCcore-{gcc_version}',
                'numpy'  :f'SciPy-bundle/2023.11-gfbf-2023b',
                'mpi'    :f'iimpi/2023b',
                'mkl'    :f'imkl/2023.2.0',
                'fftw'   :f'imkl-FFTW/2023.2.0-iimpi-2023b',
              }

    #toswaps = { f'intel/{intel_version_py}': f'intel/{intel_version}',
    #          }

    platfmodule = args.__import('platforms', path.join(args.__CHEMSH_ROOT, 'chemsh', 'utils', 'platforms', '__init__.py'))

    # MODULESHOME on the platform--this is basically the only variable you need to edit when create this file for a new platform
    print("\n >>> Setting up environment modules on STFC Scientific Computing's SCARF cluster for ChemShell...\n")
    module = platfmodule.getEnvModules(path.join(sep, 'usr', 'share', 'lmod', 'lmod'))

    environ.setdefault('LOADEDMODULES', '')
    # Gobal Arrays
    environ.setdefault('ARMCI_DEFAULT_SHMMAX_UBOUND', '131072')
   
# TODO make a function
 
    # modules to load
    # YL 09/11/2021: we no longer do the pre-screening because of complicated dependencies on SCARF
#    args.load_modules = platfmodule.getModulesToLoad(toloads, args.load_modules)

    # conflicting modules to unload
    args.unload_modules += [ ]

    # unload conflicting modules
    module.module('unload', *args.unload_modules)

    # load necessary modules
    module.module('load', *toloads.values())

    # third, load user modules
    module.module('load', *args.load_modules)

    # print a summary
    module.module('list')


# END TODO make a function

def setup(args):
    ''''''

    from os import environ, path, sep

    # force to build parallel binary on SCARF
    args.mpi     =  True
    args.i8      =  False
    args.fc      = 'mpiifort'
    args.cc      = 'mpiicc'
    args.cpc     = 'mpiicpc'
    # /apps20/sw/amd/impi/2021.10.0-intel-compilers-2023.2.1/mpi/2021.10.0/bin/mpirun
    args.mpiexec = path.join(sep, 'apps20', 'sw', 'amd', 'impi', '2021.10.0-intel-compilers-2023.2.1', 'mpi', '2021.10.0', 'bin', 'mpirun')

    if args.force_serial:
        args.mpi     =  False
        args.fc      = 'ifort'
        args.cc      = 'icc'
        args.cpc     = 'icpc'
        args.mpiexec = ''

    # /apps20/sw/amd/imkl/2023.2.0/mkl/2023.2.0/
    MKLROOT = path.join(sep, 'apps20', 'sw', 'amd', 'imkl', '2023.2.0', 'mkl', '2023.2.0')
    args.mkl = args.blas = args.lapack = args.scalapack = path.join(MKLROOT, 'lib', 'intel64')

    # this is used to fix the potential conflict between Intel compilers' math.h and the system math.h
    #args.math_header_macros = '-D__linux__ -D__USE_MISC -D__NetBSD__ -D__FreeBSD__ -U__PURE_INTEL_C99_HEADERS__'

    # since CMake 3.12 $ENV{MPI_ROOT} is needed by find_package(MPI)
    # /apps20/sw/amd/impi/2021.10.0-intel-compilers-2023.2.1/mpi/
    # /opt/ibm/platform_mpi
    environ['MPI_ROOT'] = path.join(sep, 'opt', 'ibm', 'platform_mpi')

    # /apps20/sw/amd/Python/3.11.5-GCCcore-13.2.0/lib/libpython3.11.so
    args.pythonlib = path.join(sep, 'apps20', 'sw', 'amd', 'Python', '3.11.5-GCCcore-13.2.0', 'lib', 'libpython3.11.so')
    # /apps20/sw/amd/Python/3.11.5-GCCcore-13.2.0/include/python3.11/
    args.pythondir = path.join(sep, 'apps20', 'sw', 'amd', 'Python', '3.11.5-GCCcore-13.2.0', 'include', 'python3.11')

    mkl_suffix = 'lp64'
    # YL TODO 05/03/2021: to be renamed as mkl_link_flags
#    args.mkl_flags = f'{args.mkl}{sep}libmkl_blas95_{mkl_suffix}.a {args.mkl}{sep}libmkl_lapack95_{mkl_suffix}.a -L{args.mkl} -lmkl_scalapack_{mkl_suffix} -lmkl_intel_{mkl_suffix} -lmkl_intel_thread -lmkl_core -lmkl_blacs_intelmpi_{mkl_suffix} -liomp5 -lpthread -lm -ldl'
    args.mkl_flags         = f'-lmkl_core -lmkl_intel_{mkl_suffix} -lmkl_intel_thread -liomp5 -lpthread -lm -ldl'
    args.mkl_compile_flags = f'-I{MKLROOT}{sep}include{sep}intel64{sep}lp64 -I{MKLROOT}{sep}include -mkl=parallel'
    args.scalapack_flags   = f'-lmkl_core -lmkl_intel_{mkl_suffix} -lmkl_intel_thread -lmkl_blacs_intelmpi_{mkl_suffix} -lmkl_scalapack_{mkl_suffix} -liomp5 -lpthread -lm -ldl'
    
    args.fftw              = f'/apps20/sw/amd/imkl-FFTW/2023.2.0-iimpi-2023b/lib'
    #args.fftw_include_dir  = f'/apps20/sw/amd/FFTW/{fftw_version}-GCC-13.2.0/include' 

    # NB: --ldflags holds s list
    # use the following to link in Intel MPI statically
#    args.ldflags = '-static_mpi', '-L/apps/intel/2017/compilers_and_libraries_2017.2.174/linux/mpi/lib64 -lmpi'

    platfmodule = args.__import('platforms', path.join(args.__CHEMSH_ROOT, 'chemsh', 'utils', 'platforms', '__init__.py'))

    # make `chemsh` a bash/python "heredoc" for the users' convenience
    args.__bash_header = [ '#!/bin/bash --login', 'module load AMDmodules', 'module load {}'.format(python_module), 'read -r -d \\\"\\\" CMD << EOF' ]
    args.__bash_footer = platfmodule.__bash_footer


# YL 26/10/2018: SCARF has migrated to SLURM since Oct 2018
def run(args):
    '''Submit the job'''

    from math       import ceil
    from os         import environ, path
    from subprocess import check_output
    from .          import _slurm

    help(args)

    exclusive = {True:'--exclusive=user',False:''}[args.exclusive]

    args.partition = args.partition or 'scarf'

    if args.partition not in _slurm.partitions:
        print('\n Requested --partition={} is not valid. Using \"scarf\" as default'.format(args.partition))
        args.partition = 'scarf'

    print('\n Requested --partition={} has {} CPU cores per node'.format(args.partition, _slurm.partitions2ncpus_per_node[args.partition]))

    args.nprocs = args.nprocs or _slurm.partitions2ncpus_per_node[args.partition]
    nnodes = ceil(args.nprocs/_slurm.partitions2ncpus_per_node[args.partition])

    environ['MPI_EXEC']           =          args.mpiexec
    environ['CHEMSH_EXEC']        =          args.__CHEMSH_EXEC
    environ['CHEMSH_NPROCS']      =      str(args.nprocs)
    environ['CHEMSH_NPROCS_MM']   =      str(args.nprocs_mm)
    environ['CHEMSH_NTHREADS']    =      str(args.nthreads)
    environ['CHEMSH_NWORKGROUPS'] =      str(args.nworkgroups)
    environ['CHEMSH_GA_MEMORY']   =      str(args.ga_memory)
    environ['CHEMSH_INPUT']       = ' '.join(args.filenames)
    environ['CHEMSH_OUTPUT']      =          args.jobname+'.log'
    environ['CHEMSH_WRAPPER']     =          args.wrapper
    environ['CHEMSH_JOBNAME']     =          args.jobname
    environ['CHEMSH_DEBUG']       =      str(args.debug)
    environ['CHEMSH_TEST_DRIVER'] =          args.test
    environ['CHEMSH_TESTSETS']    = ' '.join(args.sets)

    # TWK 17/11/2023 Parse user-defined environment variables
    if args.export:
        userenvs = dict(envvar.split("=") for envvar in args.export.split(","))
        environ.update(userenvs)

    if args.extra_scheduler_arg:
        extra_arguments = args.extra_scheduler_arg.split(':')
    else:
        extra_arguments = ['--comment=chemshell']

    # Slurm's --export will override the environment variables
    job_id = check_output([ 'sbatch',
                            '--export=ALL',
                             exclusive,
                            '--reservation={}'.format(args.project),
                            '--mem-per-cpu={}'.format(args.mem*1024), # request memory per CPU core in MB
                            '-p',     args.partition,
                            '-J',     args.jobname,
                            '-o',     args.jobname+'.out',
                            '-e',     args.jobname+'.err',
                            '-t',     args.walltime,
                            '-n', str(args.nprocs),
                            *extra_arguments,
#                            '-N', str(nnodes),
#                            '--export=' + 'CHEMSH_INPUT='      +    args.filenames      + ','
#                                        + 'CHEMSH_OUTPUT='     +    args.jobname+'.log' + ','
#                                        + 'CHEMSH_WRAPPER='    +    args.wrapper        + ','
#                                        + 'CHEMSH_NPROCS='     +str(args.nprocs)        + ','
#                                        + 'CHEMSH_NWORKGROUPS='+str(args.nworkgroups)   + ','
#                                        + 'CHEMSH_EXEC='  +chemsh,
                            path.join(path.dirname(path.realpath(__file__)), 'scarf.sh')
                            ])

    print("\n\n >>>", job_id.decode().rstrip(), "to SCARF compute nodes\n\n")


def help(args):
    '''Print help information for job submitting'''

    from numpy import char
    from .     import _slurm, textwrap_indent

    print("\n STFC Scientific Computing's SCARF cluster")
    print('\n Available SLURM quality of service (QoS) to choose using `chemsh --qos=`\n{}'.format(textwrap_indent(', '.join(_slurm.qos))))
    print('\n Available SLURM partitions to choose using `chemsh --partition=`\n{}'.format(textwrap_indent(', '.join(_slurm.partitions))))
    print('\n Available SLURM reservations to choose using `chemsh --project=`\n{}'.format(textwrap_indent(', '.join(_slurm.reservations))))

    print('\n Average number of CPU cores per node in each partition:')
    max_len = str(_slurm.max_partition_str_len)
    for k, v in _slurm.partitions2ncpus_per_node.items():
        print((' {:<'+max_len+'}: {:>4d}').format(k, v))

    print('\n Number of CPU cores and nodes in each subsection:')
    subsections = [ ', '.join([ s for s in sbuff.split(',') if s.startswith('scarf') or s.startswith('gpu') ]) for sbuff in _slurm.subsections ]
    max_len = str(char.str_len(subsections).max())
    for f, c, n in zip(subsections, _slurm.ncpus_per_subsection, _slurm.nodes_per_subsection):
        if not f:
            continue
        print((' {:<'+max_len+'}: {:>4d}  {}').format(f, c, n))

    print('\n Maximum walltime allowed in each partition:')
    max_len = str(_slurm.max_partition_str_len)
    for k, v in _slurm.partitions2walltimes.items():
        print((' {:<'+max_len+'}: {:>12s}').format(k, v))

    stdout.flush()














