__author__ = 'you.lu@stfc.ac.uk'

'''Set up Environment Modules on THOMAS'''

#  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/>.

# YL 15/04/2024: compilers/intel/2022.2
intel_version    = 2022
intel_subversion = 2
# YL 15/04/2024: mpi/intel/2021.6.0/intel 
mpi_version    = 2021
mpi_subversion = 6
mpi_update     = 0
cmake_version    = ''               # not in use
python_version   = 3.9
python_subversion = 0
fftw_version      = '3.3.8'

def init(args):

    from os      import environ, getcwd, sep
    from os.path import join
    from sys     import exit

    module_version = 'current'

    # default modules to load
    toloads = { 'cmake'    : 'cmake/3.21.1', 
                'compilers':f'compilers/intel/{intel_version}.{intel_subversion}',
                'mpi'      :f'mpi/intel/{mpi_version}.{mpi_subversion}.{mpi_update}/intel',
                'python'   :f'python3/{python_version}',
#                'fftw'     :f'fftw/{fftw_version}-impi/intel-2018',
              }

    tounloads = { 'modules'  : 'default-modules/2018',
                  'gcc': 'gcc-libs',
#                  'compilers': 'compilers/intel/2018/update3',
#                  'mpi'      : 'mpi/intel/2018/update3/intel',
                }

    print("\n >>> Setting up Environment Modules on MMM Hub Young for ChemShell...\n")

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

    # `/shared/ucl/apps/modules/current/bin/modulecmd
    module_home = join(sep, 'shared', 'ucl', 'apps', 'modules', module_version)
    module = platfmodule.getEnvModules(module_home)

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

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

    # unload conflicting modules
    module.module('unload', '--silent', *args.unload_modules)
    
    # load necessary modules
    module.module('load', *args.load_modules)

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

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

    from os import path, sep

    args.mpi     =  True           # build parallel version
    args.i8      =  False          # THOMAS NumPy crashes with 64-bit version
    args.fc      = 'mpiifort'
    args.cc      = 'mpiicc'
    args.cpc     = 'mpiicpc'
    args.mpiexec = 'mpirun'

    args.mpi_include_path = f'/shared/ucl/apps/intel/{intel_version}.{intel_subversion}/mpi/{mpi_version}.{mpi_subversion}.{mpi_update}/include'
    args.mpi_lib_path     = f'/shared/ucl/apps/intel/{intel_version}.{intel_subversion}/mpi/{mpi_version}.{mpi_subversion}.{mpi_update}/lib'

    # /shared/ucl/apps/intel/2019.Update5/lib/intel64/
    args.ld_path = [ path.join(sep, 'shared', 'ucl', 'apps', 'intel', '{}.Update{}'.format(intel_version, intel_subversion), 'lib', 'intel64') ]
    # /shared/ucl/apps/intel/2019.Update5/mkl/lib/intel64/
    args.mkl = args.blas = args.lapack = args.scalapack = path.join(sep, 'shared', 'ucl', 'apps', 'intel', '{}.{}'.format(intel_version, intel_subversion), 'mkl', 'latest', 'lib', 'intel64')

    # YL 18/07/2021: updated
    # /shared/ucl/apps/intel/2019.Update5/mkl
    MKLROOT = path.join(sep, 'shared', 'ucl', 'apps', 'intel', '{}.{}'.format(intel_version, intel_subversion), 'mkl', 'latest')
    # 32-bit
    mkl_suffix = 'lp64'
    args.mkl_flags         = f'-L{args.mkl} -lmkl_intel_{mkl_suffix} -lmkl_intel_thread -lmkl_core -liomp5 -lpthread -lm -ldl'
    args.mkl_compile_flags = f'-I{MKLROOT}{sep}include{sep}intel64{sep}lp64 -I{MKLROOT}{sep}include'
    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'/shared/ucl/apps/fftw/{fftw_version}-impi/intel-2018/lib'
    args.fftw_include_dir  = f'/shared/ucl/apps/fftw/{fftw_version}-impi/intel-2018/include'

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

    # /shared/ucl/apps/python/3.9.0/gnu-4.9.2/lib/libpython3.9.so
    args.pythonlib = path.join(sep, 'shared', 'ucl', 'apps', 'python', '{}.{}'.format(python_version, python_subversion), 'gnu-4.9.2', 'lib', 'libpython{}.so'.format(python_version))
    # /shared/ucl/apps/python/3.9.0/gnu-4.9.2/include/python3.9
    args.pythondir = path.join(sep, 'shared', 'ucl', 'apps', 'python', '{}.{}'.format(python_version, python_subversion), 'gnu-4.9.2', 'include', 'python{}'.format(python_version))

    # make `chemsh` a bash/python "heredoc" for the users' convenience
    args.__bash_header = [ '#!/bin/bash --login', 'module load python3/{}'.format(python_version), '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, getcwd, path
    from subprocess import check_output
    from shutil     import copy

    ncores_per_node = 40

    # there are 24 core per node on THOMAS 
    if args.nprocs == 1:
        args.nprocs = ncores_per_node

    if args.nprocs%ncores_per_node:
        print("\n WARNING: the requested number of processes (-np) should be a multiple of 40\n")

    # there are 192 GB RAM on each node
    args.mem = 4.8

    nnodes = ceil(args.nprocs/ncores_per_node)

    username = environ['USER']

    # Scratch space on YOUNG
    scratch = '/home/%s/Scratch'%username

    # use 'Free' as the default project code
    args.project = args.project or 'Free'
    projects = [ 'Gold', 'Free', 'Test' ]
    if args.project not in projects:
        print(" >>> ChemShell ERROR: On Young, a project name (the -P argument) has to be one of: %s"%' '.join(projects))
        exit(999) 

    # YOUNG nodes are diskless and there is no $TMPDIR available, so -l tmpfs=xxG should not be requested
    job_id = check_output([ 'qsub',
                            '-V',
                            '-S', '/bin/bash',
                            '-l', 'h_rt='               +    args.walltime,
                            '-l', 'mem=%dG'             %int(args.mem),
                            '-N',                            args.jobname,
                            '-P',                            args.project,
                            '-A',                            args.account,
                            '-pe', 'mpi',                str(args.nprocs),
                            '-wd', getcwd(),
                            '-v',  'CHEMSH_INPUT='      +' '.join(args.filenames)+ ','
                                 + 'CHEMSH_OUTPUT='     +    args.jobname+'.log' + ','
                                 + 'CHEMSH_WRAPPER='    +    args.wrapper        + ','
                                 + 'CHEMSH_NPROCS='     +str(args.nprocs)        + ','
                                 + 'CHEMSH_NPROCS_MM='  +str(args.nprocs_mm)     + ','
                                 + 'CHEMSH_NWORKGROUPS='+str(args.nworkgroups)   + ','
                                 + 'CHEMSH_GA_MEMORY='  +str(args.ga_memory)     + ','
                                 + 'CHEMSH_DEBUG='      +str(args.debug)         + ','
                                 + 'CHEMSH_PYTHONPATH='+path.join(path.dirname(path.realpath(__file__)), '..', '..', '..') + ','
                                 + 'CHEMSH_EXEC='       +args.__CHEMSH_EXEC      + ','
                                 + 'CHEMSH_TEST_DRIVER='+    args.test           + ','
                                 + 'CHEMSH_TESTSETS='   +' '.join(args.sets)     + ','
                                 + 'SCRATCH='           +str(scratch) + ','
                                 + args.export,
                            path.join(path.dirname(path.realpath(__file__)), 'young.sh')
                            ])

    print("\n >>>", job_id.decode().rstrip(), "to MMM Hub Young compute nodes\n")




