'''Set up Environment Modules on COSMA (Durham)'''

#  Copyright (C) 2023 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/>.

# you.lu@stfc.ac.uk

oneapi_version    = '2022.3.0'
intel_version     = '2022.2.1'
intel_subversion  = '2018.3.222'
mpi_version       = '2021.7.1'
intel_year        = intel_subversion.split('.')[0]
python_version    = '3.10'
python_subversion = '3.10.7'
python_module     =f'python/{python_subversion}'
fftw_version      = '3.3.10'
cmake_version     = '3.25.1'


# 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.cosma
__package__ = 'utils.platforms'


def init(args):

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

    toloads = { 
                'python'  : f'python/{python_subversion}',
                'cmake'   : f'cmake/{cmake_version}',
                'intel'   : f'intel_comp/{oneapi_version}',
                'compiler': f'compiler/{intel_version}',
                'mpi'     : f'mpi/{mpi_version}',
                'mkl'     : f'mkl/{intel_version}',
                'fftw'    : f'fftw/{fftw_version}',
              }
    tounloads = {}

    toswaps = {
              }

    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 UKRI STFC's SCARF for ChemShell...\n")
    # /cosma/local/Modules/default
    module = platfmodule.getEnvModules(path.join(sep, 'cosma', 'local', 'Modules', 'default'))

    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)

    # modules to unload
    args.unload_modules = platfmodule.getModulesToLoad(tounloads, args.unload_modules)

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

    # load necessary modules
#    module.module('load', *args.load_modules)
    # YL 09/11/2021: first, load numpy which also load Intel 17 and remember $PATH and $LD_LIBRARY_PATH
    module.module('load', *toloads.values())

    # YL 08/11/2021: second, swapping Intel 17/18 will reset environment variables including NumPy and especially other user-loaded modules such as OpenMPI 
    args.__swap_modules = []
    for mod in toswaps.keys():
        module.module('swap', mod, toswaps[mod])
        args.__swap_modules += [(mod,toswaps[mod])]

    # 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'
    args.arch = 'intel'

    # since CMake 3.12 $ENV{MPI_ROOT} is needed by find_package(MPI)
    mpi_root = path.join(sep, 'cosma', 'local', 'intel', f'oneAPI_{oneapi_version}', 'mpi')
    environ['MPI_ROOT'] = mpi_root

    # /cosma/local/intel/oneAPI_2022.3.0/mpi/2021.7.1/bin/mpirun
    args.mpiexec = path.join(mpi_root, f'{mpi_version}', 'bin', 'mpirun')

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

    # /cosma/local/intel/oneAPI_2022.3.0/mkl/latest/lib/intel64
    MKLROOT = path.join(sep, 'cosma', 'local', 'intel', f'oneAPI_{oneapi_version}', 'mkl', 'latest')
    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__'

    # /cosma/local/Python/3.10.7/lib/libpython3.10.so
    args.pythonlib = path.join(sep, 'cosma', 'local', 'Python', python_subversion, 'lib', f'libpython{python_version}.so')
    # /cosma/local/Python/3.10.7/include/python3.10
    args.pythondir = path.join(sep, 'cosma', 'local', 'Python', python_subversion, 'include', f'python{python_version}')

    mkl_suffix = 'lp64'
    # YL TODO 05/03/2021: to be renamed as mkl_link_flags
    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'
    # /cosma/local/Modules/modulefiles/libraries/fftw/3.3.7
#    args.fftw              = f'/cosma/local/Modules/modulefiles/libraries/fftw/{fftw_version}'
#    args.fftw_include_dir  = f'/apps/eb/software/FFTW/{fftw_version}-intel-2018b/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 {}'.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 'cosma8'

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

    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
                            '-A',     args.account,        # --account  : budget account name
                            '-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__)), 'cosma.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 UKRI STFC SCARF')
    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('') 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()














