'''Set up Environment Modules on UoW's Avon'''

#  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/>.
#
# dugan.witherick@warwick.ac.uk

iccifort_version     = '2020.4.304'
mpi_version       = '2019.9.304'
mkl_version       = '2020.4.304'
cmake_version     = '3.18.4'
python_version    = '3.8.6'
python_version_majorminor = '3.8'
scipybundle_version = '2020.11'

def init(args):

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

    # As lmod uses a hierarchical module system, these modules need to be loaded in order.
    toloads = { 'iccifort'         : 'iccifort/{}'.format(iccifort_version),
		'impi'		: 'impi/{}'.format(mpi_version),
		'imkl'		: 'imkl/{}'.format(mkl_version),
		'CMake'		: 'CMake/{}'.format(cmake_version),
		'Python'	: 'Python/{}'.format(python_version),
		'SciPy-bundle'	: 'SciPy-bundle/{}'.format(scipybundle_version),
              }

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

    # MODULESHOME
    print("\n >>> Setting up Environment Modules on UoW's Avon 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)

    # purge existing modules
    module.module('purge')

    # 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 environ, path, sep

    intel_bin_path = path.join(sep, environ['EBROOTICCIFORT'], f'compilers_and_libraries_{iccifort_version}', 'linux/bin/intel64')

    # force to build parallel binary on Avon
    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'

    args.mkl = args.blas = args.lapack = args.scalapack = path.join(sep, environ['EBROOTIMKL'], 'lib','intel64')

    args.pythonlib = path.join(sep, environ['EBROOTPYTHON'], 'lib', 'libpython%s.so'%python_version_majorminor)
    args.pythondir = path.join(sep, environ['EBROOTPYTHON'], 'include', 'python%s'%python_version_majorminor) 

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


def run(args):
    ''''''

    from os import path

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

    nprocs = max(args.nprocs, args.nworkgroups)

#    # if np=0 or np=1 run in serial; if np>1 run in parallel
#    selectcases = { 0:platfmodule.runSerial,
#                    1:platfmodule.runSerial }
#
#    # force runSerial if built in serial (mpiexec is '')
#    if not args.mpiexec:
#        nprocs = 1

    selectcases.get(nprocs, platfmodule.runParallel)(args)