'''Set up Environment Modules on ARCHER'''

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

__author__ = 'You Lu <you.lu@stfc.ac.uk>'

CHEMSH_ARCH = 'gnu'

prgenv              =f'PrgEnv-{CHEMSH_ARCH}'
#prgenv_version      = '5.2.82'
prgenv_arches       = {'PrgEnv-cray':'CRAY','PrgEnv-aocc':'AOCC','PrgEnv-gnu':'GNU'}

aocc_version        = '17'
aocc_revision       = '3.191'
aocc_revision       = '0.098'
aocc_subversion     = '{}.0.{}'.format(aocc_version, aocc_revision)

cce_version         = '15.0.0'

libsci              = 'cray-libsci'
libsci_wrap_versions = {'PrgEnv-cray':'9.0' ,'PrgEnv-aocc':'2.0' ,'PrgEnv-gnu':'9.1'}
libsci_version      = '22.12.1.1'

mpich               = 'cray-mpich'
mpich_wrap_versions = {'PrgEnv-cray':'9.1' ,'PrgEnv-aocc':'2.1' ,'PrgEnv-gnu':'9.1'}
mpich_version       = '8.1.23'
cpu                 = 'x86-rome'
network             = 'ofi'


python_module       = 'cray-python'
# 3.8.5.0
python_version      = '3.9'
python_subversion   = '13.1'

cmake_version       = '3.21.3'



# YL 22/11/2021: enabled relative importing
from pathlib import Path
from sys     import path
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'



fftw_version        = '3.3.10.3'

tcl_version         = '8.6.13'

def __checkCWD(platfmodule):
    '''On ARCHER, the current working directory must be /work'''
    
    from os import getcwd, sep

    if getcwd().startswith(sep+'home'):
        print(" >>> WARNING: ChemShell built in or calculations run from /home filesystem will not be able to run on ARCHER2 compute nodes\n")
        platfmodule.promptYes()


def init(args):

    from os  import environ, path, sep
    from sys import exit, stdout

    args.arch = 'cpe-gnu'

    # default modules to load
    toloads = { 
                'cmake'  :f'cmake/{cmake_version}',
                'python3':  python_module,
                'sci'    :  libsci,
                'fabric' : 'libfabric',
#                'tcl'    : 'tcl',                   # needed by PSFGEN (NAMD) # YL 20/12/2021: somehow Tcl was deleted on ARCHER2
                'tk'     : 'tk',                    # needed by PSFGEN (NAMD)
              }

    # YL 02/03/2021: to avoid the OOM issue we do NOT load this module, see: https://docs.archer2.ac.uk/known-issues
    tounloads = { 'epcc': 'epcc-job-env',
                }

    print("\n >>> Setting up Environment Modules on ARCHER2 for ChemShell...\n")
    stdout.flush()
    platfmodule = args.__import('platforms', path.join(args.__CHEMSH_ROOT, 'chemsh', 'utils', 'platforms', '__init__.py'))
    __checkCWD(platfmodule)

    # MODULESHOME on the platform--this is basically the only variable you need to edit when create this file for a new platform
    # /usr/share/lmod/lmod
    module = platfmodule.getEnvModules(path.join(sep, 'usr', 'share', 'lmod', 'lmod'))

    # this should be loaded prior to any other modules

    # save --load-modules values for later use
    args.__extra_modules = ' '.join([ mod for mod in args.load_modules if mod not in args.filenames ])

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

    # unload first
    module.module('--quiet', 'unload', *args.unload_modules)

    # load AMD application development environment
#    module.module('savelist')
#    module.module('--quiet', 'restore', prgenv)
#    args.__restore_modules = [ prgenv ]

    if args.debug:
        module.module('--quiet', 'load', 'atp', 'valgrind4hpc')

    # YL 02/03/2021: see https://docs.archer2.ac.uk/known-issues
    toswaps = {
                'PrgEnv-cray':'PrgEnv-gnu',
              }

    module.module('--quiet', 'restore')

    # must keep this order:
    args.__swap_modules = []
    for mod in toswaps.keys():
        module.module('--quiet', 'swap', mod, toswaps[mod])
        args.__swap_modules += [(mod,toswaps[mod])]

    environ.update(LD_LIBRARY_PATH=environ['CRAY_LD_LIBRARY_PATH']+':'+environ['LD_LIBRARY_PATH'])

    # load necessary modules
    module.module('--quiet', 'load', *args.load_modules)

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

    # ARCHER2 does not allow running on the logon nodes
    args.submit = True
    args.mpiexec = 'srun'


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

    from os import environ, path, sep

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

    # hard code location of shared objects
#    environ.update(CRAY_ADD_RPATH='yes')
#    environ.update(CRAYPE_LINK_TYPE='dynamic')


    # on ARCHER, NumPy (anaconda/python3) crashes if linked against 64-bit Intel MKL library
    args.mpi              = True           # build parallel version on ARCHER
    args.i8               = False          # build 32-bit version on ARCHER
    args.fc               = 'ftn'
    args.cc               = 'cc'
    args.cpc              = 'CC'
    args.mpiexec          = 'srun'

    # default version /opt/cray/pe/mpich/8.x.y/ofi/GNU/9.1/lib
    args.mpi_include_path = path.join(sep, 'opt', 'cray', 'pe', 'mpich', mpich_version, network, prgenv_arches.get(prgenv,'GNU'), mpich_wrap_versions.get(prgenv, '9.1'), 'include')
    args.mpi_lib_path     = path.join(sep, 'opt', 'cray', 'pe', 'mpich', mpich_version, network, prgenv_arches.get(prgenv,'GNU'), mpich_wrap_versions.get(prgenv, '9.1'), 'lib')
#    args.mpi_libraries    = path.join(sep, 'opt', 'cray', 'pe', 'mpich', mpich_version, network, prgenv_arches.get(prgenv,'GNU'), mpich_wrap_versions.get(prgenv, '9.1'), 'lib', 'libmpi.so')
    args.mpi_libraries    = '-lmpi'

    # has to be FC=ftn as on ARCHER FC=ifort doesn't compile
# already made default
#    args.__fc = args.fc
#    args.__cc = args.cc
    # use UNSETFLAGS_FORGA to smuggle `-dynamic` needed by `ftn` for linking in src/tools/GNUmakefile to configure GA
    args._extra_ldflags  = 'export UNSETFLAGS_FORGA=LDFLAGS=-dynamic'

    # on ARCHER multiple make jobs can cause problems
#    args.jobs = 1

    # we need this for cray: /opt/cray/pe/cce/10.0.3/cce-clang/x86_64/lib/libunwind.so
    if CHEMSH_ARCH == 'cray':
        cray_root = path.join(sep, 'opt', 'cray', 'pe', 'cce')
        ld_path = path.join(cray_root, cce_version, 'cce-clang', 'x86_64', 'lib')
        args.ld_path = [ ld_path ]

    # /opt/cray/pe/libsci/20.10.1.2/GNU/9.1/x86_64/lib
    libsci_root = path.join(sep, 'opt', 'cray', 'pe', 'libsci')
    libsci_lib_dir = path.join(libsci_root, libsci_version, prgenv_arches.get(prgenv, 'GNU'), libsci_wrap_versions.get(prgenv, '9.1'), 'x86_64', 'lib')
    args.mkl = args.blas = args.lapack = args.scalapack = libsci_lib_dir
    # TWK 23/08/2022 The libsci linking flag appears to cause a segfault on program exit (is it the wrong library?).
    # It should not be necessary to link in libsci explicitly anyway; removing the flag stops the segfault and hopefully has no side effects.
    #args.mkl_flags       = f'-lsci_{CHEMSH_ARCH}_mpi'
    #args.scalapack_flags = f'-lsci_{CHEMSH_ARCH}_mpi'
    args.mkl_flags       = f''
    args.scalapack_flags = f''
    args.fftw              = f'/opt/cray/pe/fftw/{fftw_version}/x86_rome/lib'
    args.fftw_include_dir  = f'/opt/cray/pe/fftw/{fftw_version}/x86_rome/include'

    # Tcl is needed by PSFGEN (NAMD): /work/y07/shared/utils/tcl/8.6.10-gcc10/include
    args.tcl_include_path = f'/work/y07/shared/utils/tcl/{tcl_version}-gcc10/include'
    # YL 20/12/2021: the Tcl module was suddenly deleted on ARCHER2 and we have to use other tcl.h for now
    # /utils/core/vmd/tcl/8.6.10-gcc10/include/tcl.h
    args.tcl_include_path = '/work/y07/shared/utils/core/vmd/tcl/8.6.10-gcc10/include'
    args.tcl_include_path = '/work/y07/shared/utils/core/tcl/8.6.13/include'

#-finit-local-zero
    if args.debug:
        args.cflags = [ '-O0', '-g', '-fopenmp -ffast-math -fno-finite-math-only -Wno-all' ]
        args.fflags = [ '-O0', '-g', '-fopenmp -fall-intrinsics -fallow-argument-mismatch -ffree-line-length-255 -finit-integer=0 -finit-real=zero -fcheck=all -Wno-all' ]
    else:
        # CMake also use CMake_LANGUAGE_FLAGS for linking so do not define anything here but only handle it in CMakeLists.txt
        args.cflags = [ '-O3', '-fopenmp -ffast-math -fno-finite-math-only -Wno-all' ]
        # we have to specify -O levels for the submodules individually because DL_PY2F can't be compiled by Cray with -O1+ options now
        args.fflags =  [ '-O2', '-fopenmp -fall-intrinsics -fallow-argument-mismatch -ffree-line-length-255 -finit-integer=0 -finit-real=zero -Wno-all' ]
        args.fflags += [ '-Wno-implicit-procedure -Wno-argument-mismatch' ]
        args._extra_fflags = f'-target-cpu={cpu} -target-network={network}'
        args._extra_cflags = f'-target-cpu={cpu} -target-network={network}'

    # /opt/cray/pe/python/3.8.5.0/lib/libpython3.8.so.1.0
    args.pythonlib = path.join(sep, 'opt', 'cray', 'pe', 'python', python_version+'.'+python_subversion, 'lib', 'libpython{}.so'.format(python_version))
    args.pythondir = path.join(sep, 'opt', 'cray', 'pe', 'python', python_version+'.'+python_subversion, 'include', 'python{}'.format(python_version))

    platfmodule.challengeSerial(args)

    # 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



def run(args):
    '''Submit the job'''

    from math       import ceil
    from os         import environ, path
    from sys        import stdout
    from subprocess import check_output

    help(args)

    # TWK 23/1/2024 Allow user to specify partition and qos options of their choice
    if args.partition:
        partition = args.partition
    else:
        partition = 'standard'
        print(" ARCHER2: using standard partition.")

    if args.qos:
        qos = args.qos
    else:
        qos = 'standard'
        print(" ARCHER2: using standard QoS.")
    
    # TD 13/2/2024 Testing suggests that the default clock speed lowering on archer2 is not energy efficient for ChemShell
    # Present default behaviour is to reduce the clock speed to 2.0 GHz
    # YL 19/05/2024: cpu_freq needs an initial value when the user doesn't specify the argument
    if args.cpu_freq == 'High' or args.cpu_freq not in [ '2250000', '2000000', '1500000' ]:
        args.cpu_freq = '2250000'
    print(f' ARCHER2: requesting CPU frequency {int(args.cpu_freq)/1000000:.2f} GHz')

    # TWK 4/1/2024 Assume cpus_per_task is the same as the number of OpenMP threads, unless specified by user
    ncpus_per_task = args.cpus_per_task or args.nthreads
    # NUMA region on ARCHER2 is 16 cores
    sensible_nthreads = [ 1, 2, 4, 8, 16 ]
    if ncpus_per_task not in sensible_nthreads:
        print("\n ChemShell@ARCHER2 WARNING: number of threads set to {} but recommended to be 1, 2, 4, 8 or 16\n".format(ncpus_per_task))

    # args.nprocs-per-node is defaulted to 0
    # Default MPI processes per node is obtained by dividing the 128 cores of an ARCHER2 node by number of cores allocated to each MPI task
    ntasks_per_node = args.nprocs_per_node or ( 128 // ncpus_per_task )
    nnodes = ceil(args.nprocs/ntasks_per_node)
    if args.nprocs%ntasks_per_node and args.nprocs > ntasks_per_node:
        print("\n ChemShell@ARCHER2 WARNING: number of requested MPI processes via `--np {}` not divisible by number of tasks per node\n".format(args.nprocs))

    # YL 20/12/2021: on the full ARCHER2 NAMD needs to be called with srun however --exclusive is enforced and causes
    #                the "step creation temporarily disabled, retrying (Requested nodes are busy)" error because all the
    #                memory resources are assigned to ChemShell; we need to reserve some memory for NAMD
    mem_per_node = 256000
    if args.mem_ratio == 1.0:
        mem_archer2 = 0
    else:
        mem_archer2 = int(mem_per_node*args.mem_ratio)
        mem_archer2 *= nnodes

#    environ['CHEMSH_EXEC']        =          args.__CHEMSH_EXEC
#    environ['CHEMSH_NPROCS']      =      str(args.nprocs)
#    environ['CHEMSH_NPROCS_MM']   =      str(args.nprocs_mm)
#    environ['CHEMSH_NWORKGROUPS'] =      str(args.nworkgroups)
#    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)
#                                        + 'CHEMSH_EXTRA_MODS=' +    args.__extra_modules+ ','

#                            '-w', '001314',                # --nodelist
    # Slurm's --export will override the environment variables
    job_id = check_output([ 'sbatch',
#                            '--get-user-env',
                            '-A',     args.account,        # --account  : budget account name
                            '-J',     args.jobname,        # --job-name : job name
                            '-p',     partition,           # --partition: partition
                            '-q',     qos,                 # --qos      : quality of service
                            '-e',     args.jobname+'.err', # --error    : error filename
                            '-o',     args.jobname+'.out', # --output   : output filename
                            '-t',     args.walltime,       # --time     : time
                            '-n', str(args.nprocs),        # --ntasks   : number of MPI processes
                            '-N', str(nnodes),             # --nodes    : number of nodes
                            '--ntasks-per-node={}'.format(ntasks_per_node),
                            '--cpus-per-task={}'.format(ncpus_per_task),
                            '--cpu-freq={}'.format(args.cpu_freq),
                            '--oversubscribe',             # won't work because --exclusive is enforced on ARCHER2
                            '--export=' + 'CHEMSH_INPUT=' +' '.join(args.filenames)     + ','
                                        + 'CHEMSH_OUTPUT='     +    args.jobname+'.log' + ','
                                        + 'CHEMSH_JOBNAME='    +    args.jobname        + ','
                                        + 'CHEMSH_WRAPPER='    +    args.wrapper        + ','
                                        + 'CHEMSH_NPROCS='     +str(args.nprocs)        + ','
                                        + 'CHEMSH_GA_MEMORY='  +str(args.ga_memory)     + ','
                                        + 'CHEMSH_NPROCS_MM='  +str(args.nprocs_mm)     + ','
                                        + 'CHEMSH_NTHREADS='   +str(args.nthreads)      + ','
                                        + 'CHEMSH_NWORKGROUPS='+str(args.nworkgroups)   + ','
                                        + 'CHEMSH_DEBUG='      +str(args.debug)         + ','
                                        + 'CHEMSH_EXEC='       +    args.__CHEMSH_EXEC  + ','
                                        + 'CHEMSH_ARCH='       +   'cray'               + ','
                                        + 'CHEMSH_UNLOAD_MODS='+    ' '.join([ mod for mod in args.unload_modules if mod not in args.filenames ])+ ','
                                        + 'CHEMSH_LOAD_MODS='  +    ' '.join([ mod for mod in args.load_modules if mod not in args.filenames ])+ ','
                                        + 'CHEMSH_SWAP_MODS='  +    ';'.join([ k+' '+v for k,v in args.__swap_modules ])+ ','
                                        + 'CHEMSH_TEST_DRIVER='+    args.test           + ','
                                        + 'CHEMSH_TESTSETS='   +' '.join(args.sets)     + ','
                                        + 'MEM_ARCHER2='       +str(mem_archer2) + ','
                                        + args.export,
                            path.join(path.dirname(path.realpath(__file__)), 'archer2.sh')
                            ])

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




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

    from . import _slurm, textwrap_indent

    print('\n ARCHER2 Full System')
    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, int(v/2)))

    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))
    
    print('')


