'''Set up Environment Modules on UoB's BlueCrystal4'''

#  Copyright (C) 2019 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

python_module     = 'languages/anaconda3'
anaconda_version  = '2019.07-3.7.3-biopython'
anaconda_dirname  = '3.7-2019.07-biopython'
intel_version     = '2018-u3'
intel_release     = '2018.3.222'
cmake_version     = '3.13.2'
mkl_version       = '2017.1.132-GCC-5.4.0-2.26'
mpi_version       = '2017.1.132-GCC-5.4.0-2.26'

intel_subversion  = '2017.1.132'
intel_year        = intel_subversion.split('.')[0]
python_version    = '3.7'
python_subversion = '3.6.5'

def init(args):

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

    toloads = { 'tools/cmake'    :'tools/cmake/{}'.format(cmake_version),
                python_module    :'{}/{}'.format(python_module, anaconda_version),
                'languages/intel':'languages/intel/{}'.format(intel_version),
                'imkl'           :'imkl/{}'.format(mkl_version),
                'impi'           :'impi/{}'.format(mpi_version),
              }

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

    # MODULESHOME
    print("\n >>> Setting up Environment Modules on UoB's BlueCrystal4 for Py-ChemShell...\n")
    module = platfmodule.getEnvModules(path.join(sep, 'usr', 'share', 'lmod', 'lmod'))

# TODO make a function
 
    # modules to load
    args.load_modules = platfmodule.getModulesToLoad(toloads, args.load_modules)

    # conflicting modules to unload
    args.unload_modules += [ 'languages/java' ]

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

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

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

# END TODO make a function

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

    from os import path, sep

    intel_bin_path = '/mnt/storage/software/languages/intel/intel-parallel_studio_xe_{}/compilers_and_libraries_{}/linux/bin/intel64/ifort'.format(intel_version.replace('-','_'), intel_release)

    # force to build parallel binary on BC4
    args.mpi       =  True
    args.i8        =  False
    args.fc        = 'mpiifort -fc={}/ifort'.format(intel_bin_path)
    args.cc        = 'mpiicc -cc={}/icc'.format(intel_bin_path)
    args.cpc       = 'mpiicpc -cxx={}/icpc'.format(intel_bin_path)
    args.mpiexec   = 'mpirun'

    # /mnt/storage/apps/intel/impi/2017.1.132/lib64
    # /mnt/storage/apps/intel/mkl/lib/intel64
    args.mkl = args.blas = args.lapack = args.scalapack = path.join(sep, 'mnt', 'storage', 'apps', 'intel', 'mkl', 'lib', 'intel64')

    # /mnt/storage/software/languages/anaconda/anaconda.3.7-2019.07-biopython/include/python3.7m
    args.pythonlib = path.join(sep, 'mnt', 'storage', 'software', 'languages', 'anaconda', 'anaconda.{}'.format(anaconda_dirname), 'lib', 'libpython%sm.so'%python_version)
    args.pythondir = path.join(sep, 'mnt', 'storage', 'software', 'languages', 'anaconda', 'anaconda.{}'.format(anaconda_dirname), 'include', 'python%sm'%python_version)

    mkl_suffix = '_lp64'
    args.mkl_flags       = '-lmkl_core -lmkl_intel%s'%mkl_suffix
    args.scalapack_flags = '-lmkl_blacs_intelmpi%s -lmkl_scalapack%s'%(mkl_suffix, mkl_suffix)

    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: BC4 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

    # system defaults
    args.partition  = args.partition or 'cpu'
    args.nprocs = args.nprocs or 28
    nnodes = ceil(args.nprocs/28)

    if type(args.mem) is int:
        memory = str(args.mem*1024)+'M'

    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']       = 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)

    # Slurm's --export will override the environment variables
    job_id = check_output([ 'sbatch',
                            '--get-user-env',
                            '-J',     args.jobname,
                            '-o',     args.jobname+'.log',
                            '-e',     args.jobname+'.err',
                            '-p',     args.partition,
                            '-t',     args.walltime,
                            '-n', str(args.nprocs),
                            '-N', str(nnodes),
                            '--mem',  memory,
#                            '--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__)), 'bluecrystal4.sh')
                            ])

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