# to replace the Python 2 version of /opt/modules/default/init/python.py on RAVEN
__author__ = 'you.lu@stfc.ac.uk'

'''Set up Environment Modules on RAVEN'''

#  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, getcwd, path, sep
    from sys            import exit

    mpi_version      = '2017.2'
    mpi_release      = '%s.174'%mpi_version

    # YL 25/07/2018: temporarily rewinded back to 2017.1 because 2017.2 results in an error:
    # OSError: /software/compilers/intel/2017.2/parallel_studio_xe_2017.2.050/compilers_and_libraries_2017/linux/mkl/lib/intel64/libmkl_def.so: undefined symbol: mkl_pdepl_s_ft_3d_z_pp
    mpi_version      = '2017.1'
    mpi_release      = '%s.132'%mpi_version

    compiler_version = 'intel-%s'%mpi_version
    cmake_version    = '3.9.1'
    anaconda_version = '3'
    python_main_version = '3.5'

    # default modules to load
    toloads = { 'cmake'   :'cmake/%s'%cmake_version,
                'anaconda':'anaconda/%s'%anaconda_version,
                'compiler':'compiler/%s'%compiler_version,
#                'mpi'     :'mpi/intel.%s'%mpi_version,
              }

    print("\n >>> Setting up Environment Modules on RAVEN for building ChemShell...\n")

    # MODULESHOME on the platform--this is basically the only variable you need to edit when create this file for a new platform
    # /opt/modules/default/Modules/bin/modulecmd
    module = utils.getEnvModules(path.join(sep, 'opt', 'modules', 'default', 'Modules'))

    # modules to load
    load = utils.getModulesToLoad(toloads, load)

    # conflicting modules to unload
    unload += [ 'intel/dummy', 'oscar-modules' ] # cray-libsci conflicts with Intel MKL and MPI

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

    # unload conflicting modules
    module.module('unload', *unload)
    
    # FIXME: sometimes the 'mpi' modules is not properly loaded, so we do it explicitly here
    module.module('load', 'mpi/intel.%s'%mpi_version)

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

    args.mpi              = True           # build parallel version on RAVEN
    args.mpi_include_path = '/software/compilers/intel/%s/include64'%mpi_release
    args.i8               = False          # build 32-bit version on RAVEN
    args.fc               = 'mpiifort'
    args.cc               = 'mpiicc'
    args.mpiexec          = 'mpirun'        # hard-coded in raven.sh.in anyway
    # so far only required by DL_POLY
    # default version is args.mpi_include_path = '/opt/cray/mpt/default/gni/mpich-INTEL/160/include'
    args.fflags           = [ '-lcilkrts', '-xAVX', '-O2', '-diag-disable=10121' ]
    args.ld_path          = [ path.join(sep, 'software', 'compilers', 'intel', mpi_version, 'lib', 'intel64') ]
    args.mkl              = path.join(sep, 'software', 'compilers', 'intel', 'mkl', 'lib', 'intel64')
    # on RAVEN, NumPy (anaconda/python3) crashes if linked against 64-bit Intel MKL library
    args.mkl_flags        = '-lmkl_core -lmkl_intel_lp64 -lmkl_intel_thread'
    # only consider the case `--mkl` rather than, for example, `--mkl /opt/intel/mkl/lib64` (with argument)
    if 'scalapack' in args and args.scalapack is True:
        args.scalapack_flags = '-lmkl_scalapack_lp64 -lmkl_blacs_intelmpi_lp64'
    else:
        args.scalapack_flags = ''
    args.pythonlib = path.join(sep, 'software', 'languages', 'anaconda%s'%anaconda_version, 'lib', 'libpython%sm.so'%python_main_version)
    args.pythondir = path.join(sep, 'software', 'languages', 'anaconda%s'%anaconda_version, 'include', 'python%sm'%python_main_version)

    utils.challengeSerial(args)


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

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

    nnodes = ceil(args.nprocs/16)
    ncpus_per_node = ceil(args.nprocs/args.nnodes)

    bigmem = ''
    
    # the default partition is "workq"
    partition = 'workq'
    if args.nodename:
        partition = args.nodename

    job_id = check_output([ 'qsub',
                            '-V',
                            '-N',                            args.jobname,
                            '-P',                            args.account,
                            '-p',                            partition,
                            '-l select=%d:ncpus=%d:mpiprocs=%d'%(args.nnodes, args.ncpus, args.ncpus),
#                            '-l select='+str(nnodes)+bigmem,
                            '-l walltime='       +           args.walltime,
                            '-v' + 'CHEMSH_INPUT='      +    args.filenames      + ','
                                 + 'CHEMSH_OUTPUT='     +    args.jobname+'.log' + ','
                                 + 'CHEMSH_NPROCS='     +str(args.nprocs)        + ','
                                 + 'CHEMSH_NPROCS_MM='  +str(args.nprocs_mm)     + ','
                                 + 'RAVEN_NPROCS_PER_NODE='+str(min(args.nprocs,24))+','
                                 + 'CHEMSH_NWORKGROUPS='+str(args.nworkgroups)   + ','
                                 + 'CHEMSH_PYTHONPATH='+path.join(path.dirname(path.realpath(__file__)), '..', '..', '..') + ','
                                 + 'CHEMSH_EXEC='       +chemsh,
                            path.join(path.dirname(path.realpath(__file__)), 'raven.sh')
                            ])

    print("\n >>>", job_id.decode().rstrip(), "submitted to RAVEN", args.nodename.title(), "compute nodes\n")




