# to replace the Python 2 version of /cm/local/apps/environment-modules/current/init/python.py on Neale
# you.lu@stfc.ac.uk
'''Set up Environment Modules on Hartree Centre Neale'''

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

def initPlatform(args, utils=None, load=[], unload=[]):

    from os             import environ, path, sep
    from re             import sub
    from importlib.util import spec_from_file_location, module_from_spec

    intel_version  = '2016.1.10'
    python_version = '3.6.0'

    toloads = { 'cmake'    :'cmake',
                'python3'  :'python3',
                'intel_cs' :'intel_cs',
                'intel_mpi':'intel_mpi',
                'allinea'  :'allinea',
                'slurm'    :'slurm'
              }

    # MODULESHOME on the platform--this is basically the only variable you need to edit when create this file for a new platform
    print("\n >>> Setting up Environment Modules on Hartree Centre Neale for building ChemShell...\n")
    module = utils.getEnvModules(path.join(sep, 'cm', 'local', 'apps', 'environment-modules', 'current'))

    environ.setdefault('LOADEDMODULES', '')
    environ.setdefault('ARMCI_DEFAULT_SHMMAX_UBOUND', '131072')
    
    # modules to load
    load = utils.getModulesToLoad(toloads, load)

    # conflicting modules to unload
    unload += [ ]

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

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

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

    # force to build parallel binary on Neale
    args.mpi       =  True
    args.fc        = 'mpiifort'
    args.cc        = 'mpiicc'
    args.mpiexec   = 'srun'
    args.mkl       = '/cm/shared/hartree/intel/intel_cs/%s/mkl'%intel_version
    args.mkl       = path.join(sep, 'cm', 'shared', 'hartree', 'intel', 'intel_cs', intel_version, 'mkl', 'lib', 'intel64')
    args.mkl_flags = '-lmkl_core -lmkl_intel_lp64 -lmkl_scalapack_lp64 -lmkl_blacs_intelmpi_lp64 -lmkl_intel_thread'
    args.scalapack = '/cm/shared/hartree/intel/intel_cs/%s/mkl'%intel_version
    args.pythonlib = path.join(sep, 'cm', 'shared', 'hartree', 'python3', python_version, 'lib',     'libpython3.6m.so')
    args.pythondir = path.join(sep, 'cm', 'shared', 'hartree', 'python3', python_version, 'include', 'python3.6m')
    args.ld_path   = [ path.join(sep, 'cm', 'shared', 'hartree', 'intel', 'intel_cs', intel_version, 'lib', 'intel64') ]


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

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

#    initPlatform(load=args.load_modules, unload=args.unload_modules)

    environ['CHEMSH_EXEC']        = chemsh
    environ['CHEMSH_NPROCS']      = str(args.nprocs)
    environ['CHEMSH_NWORKGROUPS'] = str(args.nworkgroups)
    environ['CHEMSH_INPUT']       = args.filenames
    environ['CHEMSH_OUTPUT']      = args.jobname+'.log'
    environ['CHEMSH_WRAPPER']     = args.wrapper

    nnodes = ceil(args.nprocs/16)
#    ncpus_per_node = ceil(args.nprocs/args.nnodes)
    if not args.nodename:
        args.nodename = 'defq'

# YL TODO: -o and -e
    job_id = check_output([ 'sbatch',
                            '--get-user-env',
                            '--ntasks-per-node', str(16),
                            '-J',     args.jobname,
                            '-o',     args.jobname+'.log',
                            '-e',     args.jobname+'.err',
                            '-t',     args.walltime,
                            '-p',     args.nodename,
                            '-N', str(nnodes),
                            '-n', str(args.nprocs),
                            path.join(path.dirname(path.realpath(__file__)), 'neale.sh')
                            ])
#                            '--export=' + 'input=' +    args.filenames      + ','
#                                        + 'output='+    args.jobname+'.log' + ','
#                                        + 'nprocs='+str(args.nprocs)        + ','
#                                        + 'exec='  +chemsh,

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




