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

'''Some setup tools'''

def compileNWChem(args):
    '''Generate an NWChem compile script'''

    from os       import chmod, sep
    from stat     import S_IRWXU, S_IRWXG, S_IROTH, S_IXOTH
    from os.path  import abspath, dirname, isdir, join, realpath
    from datetime import datetime

    patchmodule = args.__import('patch', join(args.__CHEMSH_ROOT, 'chemsh', 'interfaces', 'nwchem', 'patch.py'))

    # must use absolute path
    args.nwchem = args.__getPathOrURL(args.nwchem)

    args.__nwchem_script = ''
    # not always available
    try:
        args.__swap_modules
    except:
        args.__swap_modules = []
    try:
        args.__restore_modules
    except:
        args.__restore_modules = []
        
    # convert to realpath (args.mkl is the "lib" dir, typically "intel64_lin" in realpath)
    args.mkl = realpath(args.mkl)
    # find "include" dir from the lib dir
    lbuff = [sep] + args.mkl.split(sep)
    for i in range(len(lbuff),2,-1):
        truncated = join(*lbuff[:i])
        mkl_incl_dir = join(truncated,'include')
        if isdir(mkl_incl_dir):
            break
    
    if args.nwchem:
        if isdir(args.nwchem):
            nwchem_location = args.nwchem
            nwchem_dest     = args.nwchem
            dest_type       = 'dir'
        # YL 17/11/2021: because the current script is run before anything and we can't know yet
        #                where NWChem will be extracted to if it's downloaded from the repository
        #                we have to defer this script for a later CMake configure()
        elif args.nwchem.lower().strip().startswith('https://'):
            nwchem_location = '@NWCHEM@'
            nwchem_dest     = join(args.__CHEMSH_ROOT, 'chemsh', 'interfaces', 'nwchem')
            dest_type       = 'url'

        # prepare the modules
        restore_modules = load_modules = unload_modules = swap_modules = ''
        for mod in args.__restore_modules:
            restore_modules += 'module restore {}\n'.format(mod)
        for mod in args.load_modules:
            load_modules += 'module load {}\n'.format(mod)
        for mod in args.unload_modules:
            unload_modules += 'module unload {}\n'.format(mod)
        for mod1, mod2 in args.__swap_modules:
            swap_modules += 'module swap {} {}\n'.format(mod1, mod2)

        # import external module `nwchem_compile.py`
        nwchem = args.__import('nwchem', join(args.__CHEMSH_ROOT, 'chemsh', 'interfaces', 'nwchem', 'compile_script.py'))

        # patches switches
        args._nwchem_patches = []
        # hack the hardcoded Cray flags that don't work on ARCHER2
        if args.platform.lower() == 'archer2':
            args._nwchem_patches += ['makefile_h','mm_bond_F']
            # defer patching if dest_type=='url'
            if dest_type != 'url':
                patchmodule.patch_makefile_h(nwchem_dest)
                patchmodule.patch_mm_bond_F(nwchem_dest)

        # patch src/GNUmakefile to insert escapes: only for certain versions of Intel
        if args.__FC_VERSION.split('.')[0] in ['2018'] and args.__fc == 'ifort' and args.mpi:
            args._nwchem_patches += ['GNUmakefile']
            # defer patching if dest_type=='url'
            if dest_type != 'url':
                patchmodule.patch_GNUmakefile(nwchem_dest)

        # this is a hack to fix the conflict between Intel compilers and system libs (see platforms/ubuntu.py)
        # we could do it using `/bin/sed` but it would be system-dependent
        _m_pi = ''
        if args.f128:
            args._nwchem_patches += ['rpspsolve_c']
            # defer patching if dest_type=='url'
            if dest_type != 'url':
                patchmodule.patch_rpspsolve_c(nwchem_dest)
            _m_pi = '''# however c99 causes M_PI undefined\necho NB: ChemShell has patched src/nwpw/nwpwlib/pseudopotential/rpspsolve.c\necho to hard engrave an M_PI value because the C99 flag causes it undefined'''

        # special flags
        try:
            extra_fflags = args._extra_fflags
            extra_cflags = args._extra_cflags
        except:
            extra_fflags = ''
            extra_cflags = ''

        # on ARCHER `ftn` needs `-dynamic`
        try:
            extra_ldflags = args._extra_ldflags
        except:
            extra_ldflags = ''

        mkl_libs_dir = 'MKL_LIBS_DIR={}'.format(args.mkl)
        mkl_incl_dir = 'MKL_INCL_DIR={}'.format(mkl_incl_dir)

        # YL 16/11/2021: changed args.__fc and args.__cc, the underlying compilers, to just args.fc and args.cc, because:
        #                to MPI compiler wrappers and the underlying compiler both can compile, however in the Makefile (GNUmakefile)
        #                -DCOMPILATION_DATE="'`date +%a_%b_%d_%H:%M:%S_%Y`'" needs to be espaced ("\\'") to work with wrappers;
        #                from now on we patch the GNUmakefile file to allow using of compiler wrappers so that ScaLAPACK can be
        #                detected when standadlone NWChem configures GA; and this only affects the standadlone NWChem binary because
        #                ChemShell only needs the static libs which are not compiled with -DCOMPILATION_DATE= or so on

        # on ARCHER the bash version is too low to allow `export FC="ftn -fPIC"`!
        # YL 10/02/2020: added "-integer-size 64"
        # YL 01/07/2020: double quotation isn't accepted with $FC on ARCHER, so moved to $FFLAGS
        if args.platform.lower() == 'archer':
            fc = 'export FC={}'.format(args.fc)
            cc = 'export CC={}'.format(args.cc)
            if extra_fflags:
                extra_fflags = '"-fPIC -integer-size 64 '+extra_fflags+'"'
                extra_cflags = '"-fPIC '+extra_fflags+'"'
            else:
                extra_fflags = '"-fPIC -integer-size 64"'
                extra_cflags = '-fPIC'
            foptions = 'export FOPTIONS={}'.format(extra_fflags)
            coptions = 'export COPTIONS={}'.format(extra_cflags)
            # YL 01/07/2020: since a unknown point ARCHER's ftn strangely needs BLACS for libmkl_core
            mkl_opt    = '"-I$MKL_INCL_DIR -L$MKL_LIBS_DIR -lmkl_intel_lp64 -lmkl_core -lmkl_intel_thread -lmkl_scalapack_lp64 -lmkl_blacs_intelmpi_lp64 -mkl=parallel"'
            blas_opt   = 'export BLAS_LIB="-I$MKL_INCL_DIR -L$MKL_LIBS_DIR -lmkl_intel_lp64 -lmkl_core -lmkl_intel_thread -lmkl_blacs_intelmpi_lp64 -lpthread -lm -mkl=parallel"\nexport BLASOPT="$BLAS_LIB"'
            lapack_opt = '"-I$MKL_INCL_DIR -L$MKL_LIBS_DIR -lmkl_intel_lp64 -lmkl_core -lmkl_intel_thread -lpthread -lm -mkl=parallel"'
        if args.platform.lower() == 'archer2' and args.arch.strip().lower() == 'cpe-gnu':
            restore_modules = 'module restore\n'
            fc = 'export FC={}'.format(args.fc)
            cc = 'export CC={}'.format(args.cc)
            # YL 03/03/2021: we copy a makefile.h to overwrite the original one because for gfortran
            #                all the flags are hardcoded and we can't use anything of our own; see:
            #                interfaces/nwchem/CMakeLists.txt
            extra_fflags =f'"-fPIC -fopenmp -fdefault-integer-8 {extra_fflags}"'
            extra_cflags =f'"-fPIC -fopenmp -ffast-math {extra_cflags}"'
            extra_fflags ='"-fPIC -fdefault-integer-8"'
            extra_cflags ='-fPIC'
            foptions   = 'export FOPTIONS={}'.format(extra_fflags)
            coptions   = 'export COPTIONS={}'.format(extra_cflags)
#            foptions   = ''
#            coptions   = ''
            cc         = ''
            mkl_opt    = '"-L{} {}"'.format(args.mkl, args.mkl_flags)
            blas_opt   = '"-L{} {}"'.format(args.mkl, args.mkl_flags)
            blas_opt   = 'export BLASOPT=" "'
            lapack_opt = '"-L{} {}"'.format(args.mkl, args.mkl_flags)
            lapack_opt = '" "'
            extra_ldflags = 'export LDOPTIONS=-Wl,--unresolved-symbols=ignore-all'
            extra_ldflags = ''
            mkl_libs_dir = ''
            mkl_incl_dir = ''
            # to avoid luout=6 so that we can choose a different unit number
            _m_pi = '# to bypass the default unit=6\nexport USE_SUBGROUPS=1'
        elif args.platform.lower() == 'archer2':
            fc = 'export FC={}'.format(args.fc)
            cc = 'export CC={}'.format(args.cc)
            # -dC  : disable some compiletime errors (especially for bugs in NWChem such as unmatched numbers of arguments)
            # -eZ  : enable preprocessing
            # -N132: line width 132
            # -homp: enable OpenMP
            # -sinteger64: default 64-bit integer
            # -M112,241: disable specific warnings
            if extra_fflags:
                extra_fflags = '"-fPIC -dC -eZ -N132 -homp -sinteger64 -DCRAYFORTRAN -DUSE_POSIXF -M112,241'+extra_fflags+'"'
                extra_cflags = '"-fPIC '+extra_fflags+'"'
            else:
                extra_fflags = '"-fPIC -dC -eZ -N132 -homp -sinteger64 -DCRAYFORTRAN -DUSE_POSIXF -M112,241"'
                extra_cflags = '-fPIC'
            foptions   = 'export FOPTIONS={}'.format(extra_fflags)
            coptions   = 'export COPTIONS={}'.format(extra_cflags)
            mkl_opt    = '"-L{} {}"'.format(args.mkl, args.mkl_flags)
            blas_opt   = 'export BLAS_LIB="-L{} {}"'.format(args.mkl, args.mkl_flags)
            lapack_opt = '"-L{} {}"'.format(args.mkl, args.mkl_flags)
            extra_ldflags = 'export LDOPTIONS=-Wl,--unresolved-symbols=ignore-all'
        # intel
        else:
            if args.mpi:
                intel_omp_flags = '-mkl=parallel'
            else:
                intel_omp_flags = '-mkl=sequential'
            if int(args.__FC_VERSION.split('.')[0]) >= 2021:
                # YL 16/06/2023: in GA 5.8.2 (NWChem 7.2.0) linking against OpenMP is required for detecting ScaLAPACK
                if args.mpi:
                    intel_omp_flags = '-qmkl=parallel -liomp5'
                else:
                    intel_omp_flags = '-qmkl=sequential'
            fc = 'export FC="{} -fPIC {} -integer-size 64"'.format(args.fc, extra_fflags)
            cc = 'export CC="{} -fPIC {}"'.format(args.cc, extra_cflags)
            foptions = ''
            coptions = ''
            mkl_opt    = f'"-I$MKL_INCL_DIR -L$MKL_LIBS_DIR -lmkl_intel_lp64 -lmkl_core -lmkl_intel_thread -lmkl_scalapack_lp64 -lmkl_blacs_intelmpi_lp64 {intel_omp_flags}"'
            blas_opt   = f'export BLAS_LIB="-I$MKL_INCL_DIR -L$MKL_LIBS_DIR -lmkl_intel_lp64 -lmkl_core -lmkl_intel_thread -lpthread -lm {intel_omp_flags}"\nexport BLASOPT="$BLAS_LIB"'
            lapack_opt = f'"-I$MKL_INCL_DIR -L$MKL_LIBS_DIR -lmkl_intel_lp64 -lmkl_core -lmkl_intel_thread -lpthread -lm {intel_omp_flags}"'

        # progress monitor only for installing from URL
        if dest_type == 'url':
            _progress = 'python3 {} $NWCHEM_TOP/src/config.log $NWCHEM_TOP/src/64_to_32.log $NWCHEM_TOP/src/make.log &'.format(join(args.__CHEMSH_ROOT, 'chemsh', 'interfaces', 'nwchem', 'progress.py'))
        else:
            _progress = ''

        # dynamically generate the platform-specific information
        # do NOT use compiler wrappers
        strbuff = nwchem.getScriptTemplate(version=7).format(_swap_modules=swap_modules,
                                                             _restore_modules=restore_modules,
                                                             _load_modules=load_modules,
                                                             _unload_modules=unload_modules,
                                                             _njobs=args.jobs,
                                                             _nwchem=nwchem_location,
                                                             _nwchem_version='@NWCHEM_VERSION@',
                                                             _ga_name='@GA_NAME@',
                                                             _now=datetime.now().strftime("%H:%M:%S %a %d %b %Y"),
                                                             _fc=fc,
                                                             _cc=cc,
                                                             _foptions=foptions,
                                                             _coptions=coptions,
#                                                             _fflags=extra_fflags,
#                                                             _cflags=extra_cflags,
                                                             _mkl_opt=mkl_opt,
                                                             _blas_opt=blas_opt,
                                                             _lapack_opt=lapack_opt,
                                                             _mkl_libs_dir=mkl_libs_dir,
                                                             _mkl_incl_dir=mkl_incl_dir,
                                                             _mpi_lib_path='@MPI_LIB_PATH@',
                                                             _mpi_include_path='@MPI_INCLUDE_PATH@',
                                                             _ldflags=extra_ldflags,
                                                             _m_pi=_m_pi,
                                                             _progress=_progress,
                                                            )
        # write out the script
        if dest_type == 'dir':
            filename = join(args.nwchem, 'src', 'compile_nwchem_7_for_chemsh_on_{}.sh'.format(args.platform))
        elif dest_type == 'url':
            filename = join(nwchem_dest, 'compile_nwchem_7_for_chemsh_on_{}.sh'.format(args.platform))
        args.__nwchem_script = filename
        with open(filename, 'w') as fp:
            fp.write(strbuff)

        # chmod to "-rwxrwxr-x"
        chmod(filename, S_IRWXU+S_IRWXG+S_IROTH+S_IXOTH)


