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

#  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

intel_version     = '2019.5.281'
gcc_version       = '8.3.0'
binutils_version  = '2.32'
mpi_version       = '2018.5.288'
mkl_version       = '2019.5.281'
cmake_version     = '3.15.3'
python_version    = '3.7.4'
python_version_majorminor = '3.7'
scipybundle_version = '2019.10-Python-3.7.4'

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 = { 'intel'         : 'intel/{}-GCC-{}-{}'.format(intel_version,gcc_version,binutils_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 Orac for Py-ChemShell...\n")
    module = platfmodule.getEnvModules(path.join(sep, 'usr', 'share', 'lmod', 'lmod'))
    #module = platfmodule.getEnvModules(path.join(sep, 'csc', 'orac', 'modules', 'core'))

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

    intel_bin_path = '/csc/orac/software/Core/iccifort/{}/compilers_and_libraries_{}/linux/bin/intel64'.format(intel_version, intel_version)

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

    # /csc/orac/software/MPI/intel/2019.5.281/impi/2018.5.288/imkl/2019.5.281/lib/intel64/
    args.mkl = args.blas = args.lapack = args.scalapack = path.join(sep, 'csc', 'orac', 'software', 'MPI', 'intel', '2019.5.281', 'impi',mpi_version,'imkl',mkl_version,'lib','intel64')

    # /csc/orac/software/Compiler/GCCcore/8.3.0/Python/3.7.4
    # /mnt/storage/software/languages/anaconda/anaconda.3.7-2019.07-biopython/include/python3.7m
    args.pythonlib = path.join(sep, 'csc', 'orac', 'software', 'Compiler', 'GCCcore', gcc_version, 'Python', python_version, 'lib', 'libpython%sm.so'%python_version_majorminor) 
    args.pythondir = path.join(sep, 'csc', 'orac', 'software', 'Compiler', 'GCCcore', gcc_version, 'Python', python_version, 'include', 'python%sm'%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)
