#  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 init(args):
    '''Initialise the platform specifications'''

    # nothing to do on Ubuntu


def setup(args):
    '''Platform specifications for setup'''

    selectcases = {
                    # mpi_lib_path, mpi_include_path, and mpiexec
                    'openmpi':{ 'mpi_lib_path'    :'/usr/lib/x86_64-linux-gnu/openmpi/lib',
                                'mpi_include_path':'/usr/lib/x86_64-linux-gnu/openmpi/include',
                                'mpiexec'         :'mpiexec.openmpi',
                                'scalapack_flags' :'-lscalapack-openmpi' },
                    'mpich'  :{ 'mpi_lib_path'    :'/usr/lib/mpich/lib',
                                'mpi_include_path':'/usr/lib/mpich/include',
                                'mpiexec'         :'mpiexec.mpich',
                                'scalapack_flags' :'-lscalapack-mpich' },
                    'intel'  :{ 'mpi_lib_path'    :'',
                                'mpi_include_path':'',
                                'mpiexec'         :'mpiexec' },
                  }


    if args.__CHEMSH_ARCH == 'gnu':
        # correct the default value in `setup` if not overridden by the user
        args.__setarg('blas'            , '/usr/lib/x86_64-linux-gnu/blas')
        args.__setarg('lapack'          , '/usr/lib/x86_64-linux-gnu/lapack')
        args.__setarg('fftw'            , '/usr/lib/x86_64-linux-gnu')
        args.__setarg('fftw_include_dir', '/usr/include')
        if ('scalapack' in args and not args.scalapack) or 'scalapack' not in args:
            args.scalapack = '/usr/lib/x86_64-linux-gnu'

        

        # default MPI implementations: OpenMPI --mpi
        #                              MPICH   --mpi=mpich
        if args.mpi:
            if args.mpi in selectcases.keys():
                args.mpi, args.mpidist = True, args.mpi
            # assign the default values if not done by the user
            for arg in [ 'mpiexec', 'mpi_lib_path', 'mpi_include_path', 'scalapack_flags' ]:
                if not getattr(args, arg):
                    args.__setarg(arg, selectcases[getattr(args, 'mpidist', 'openmpi')][arg])

    elif args.__CHEMSH_ARCH == 'intel':

        # TD 28/05/24: Reasonable default compiler choices for intel oneAPI 2024.0
        #              Users can override these choices if they want
        # Parallel
        if args.mpi:
            args.__setarg('cc', 'mpiicx')
            args.__setarg('fc', 'mpiifort')
            args.__setarg('cpc', 'mpiicpx')
        # Serial
        else:
            args.__setarg('cc', 'icx')
            args.__setarg('fc', 'ifort')
            args.__setarg('cpc', 'icpx')

        # Print a warning about oneAPI versions
        print("============================================================================================")
        print(" Default Intel build via ubuntu platform file (chemsh/utils/platforms/ubuntu.py) ")
        print("============================================================================================")
        print(" Depending on the version of oneAPI you have installed, some extra options might be needed")
        print(" You might need to set the following options to ./setup explicitly to match your environment")
        print("")
        print(" -fc                   [Specify Fortran compiler]")
        print(" -cc                   [Specify C compiler]")
        print(" -cpc                  [Specify C++ compiler]")
        print(" -mpi_include_path     [Include path(s) for MPI header]")
        print(" -mpi_lib_path         [Specify MPI lib path]")
        print(" -mkl                  [Enable Intel MKL and specify MKL path]")
        print(" -fftw                 [Specify the path to FFTW library directory]")
        print(" -fftw_include_dir     [Specify the full path of directory containing FFTW header files]")
        print("")
        print(" ChemShell will set these options automatically for new oneAPI versions (>=2024.0)")
        print(" Older versions will probably require the old Intel C compiler, icc/mpiicc")
        print("============================================================================================")
        # Note that in older versions, older compiler versions might need to be set
        # 2024.0 is suggested as this is the first version where icc is deprecated
        # A bit messy to work out exactly which oneAPI versions work with which compilers

        # Intel's OneAPI (Intel Compilers v21+)
        if float(args.__FC_VERSION) >= 2021.0:

            # this is used to fix the potential conflict between Intel compilers' math.h and the Ubuntu system math.h
            args.math_header_macros = ''
    
            # YL 22/02/2021: allow the user to override the default (default is '/opt/intel/mkl/lib/intel64')
            args.__setarg('mpi_include_path', '/opt/intel/oneapi/mpi/latest/include')
            args.__setarg('mpi_lib_path'    , '\"/opt/intel/oneapi/mpi/latest/lib;/opt/intel/oneapi/mpi/latest/lib/release\"')
            args.__setarg('mkl'   , '/opt/intel/oneapi/mkl/latest/lib/intel64')
            args.__setarg('blas'  , args.mkl)
            args.__setarg('lapack', args.mkl)
    
            # YL 22/02/2021: allow the user to override the default (default is '')
            args.__setarg('fftw'            , '/opt/intel/oneapi/mkl/latest/share/mkl/interfaces/fftw3xf')
            args.__setarg('fftw_include_dir', '/opt/intel/oneapi/mkl/latest/include/fftw')

        else:

            # this is used to fix the potential conflict between Intel compilers' math.h and the Ubuntu system math.h
            args.math_header_macros = '-U__linux__ -U__USE_MISC -U__NetBSD__ -D__FreeBSD__ -U__PURE_INTEL_C99_HEADERS__'
    
            # YL 22/02/2021: allow the user to override the default (default is '/opt/intel/mkl/lib/intel64')
            args.__setarg('mkl'   , '/opt/intel/mkl/lib/intel64')
            args.__setarg('blas'  , args.mkl)
            args.__setarg('lapack', args.mkl)
    
            # YL 22/02/2021: allow the user to override the default (default is '')
            # TD 08/07/2024: FFTW path got changed at some point before version 2024.0
            #                Used to be '/opt/intel/oneapi/mkl/latest/interfaces/fftw3xf'
            args.__setarg('fftw'            , '/opt/intel/mkl/interfaces/fftw3xf')
            args.__setarg('fftw_include_dir', '/opt/intel/mkl/include/fftw')

        # TWK 4/3/2021 Don't override the default MKL flags
        if 'i8' in args and args.i8:
            mkl_suffix = '_ilp64'
        else:
            mkl_suffix = '_lp64'
        
        # COMPILETIME: when not specified by user
#        if 'mkl_flags' in args and not args.mkl_flags:
#            args.mkl_flags = '-lmkl_core -lmkl_intel%s'%mkl_suffix
#            if args.mpi:
#                args.mkl_flags += ' -lmkl_intel_thread'

        # YL 15/11/2021: it turned out we at least need ScaLAPACK flags
        if args.mpi:
            if ('scalapack' in args and not args.scalapack) or 'scalapack' not in args:
                args.scalapack = args.mkl
#        args.scalapack_flags = f'-lmkl_blacs_intelmpi{mkl_suffix} -lmkl_scalapack{mkl_suffix}'

            args.__setarg('mpiexec', 'mpirun')

        if args.f128:
            # it fails on some systems (glib v.??)
#            args.fflags.append('-D_Float128=__float128 -D_Float32=float -D_Float64=double -D_Float32x=_Float64 -D_Float64x=_Float128 -D__PURE_INTEL_C99_HEADERS__')
#            args.cflags.append('-D_Float128=__float128 -D_Float32=float -D_Float64=double -D_Float32x=_Float64 -D_Float64x=_Float128 -D__PURE_INTEL_C99_HEADERS__')
            args._extra_fflags = '-D_Float32=float -D_Float32x=double -D_Float64=double -D_Float64x=__float128'
            args._extra_ctypedefs = '-D_Float32=float -D_Float32x=double -D_Float64=double -D_Float64x=__float128'
            args._extra_cflags = args._extra_fflags+' -D__PURE_INTEL_C99_HEADERS__'
            args.fflags.append(str(args._extra_fflags))
            args.cflags.append(str(args._extra_cflags))
            # another bug: /opt/intel/compilers_and_libraries_2017.4.196/linux/compiler/include/math.h(1300): error: identifier "_LIB_VERSION_TYPE" is undefined
            # https://software.intel.com/en-us/forums/intel-c-compiler/topic/760979


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)


