#  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/>.

import chemsh
from ..dl_py2f import py2f, utils

# load Fortran/C Parallel Module
libparallel = utils.modutils.getLinkedLib('parallel', chemsh.CHEMSH_LIB_DIR)


class Parallel(object):
    ''''''

#    def createPyObject(self, objectRef):
    def createPyObject(self, objectRef, onedimsize, twodimsizem, twodimsizen):
        '''Pass Python object to Fortran'''

        from ctypes import byref, c_int
    
        setattr(self, 'objectPtr', py2f(objectRef))
        ierror = libparallel.create_python_object(byref(self.objectPtr),
                                                  byref(c_int(onedimsize)),
                                                  byref(c_int(twodimsizem)),
                                                  byref(c_int(twodimsizen)))
    
        if ierror:
            print(" >>> ERROR in createPyResult!")


    def cleanPyObject(self):
        '''Clean up'''

        ierror = libparallel.clean_python_object()


    def syncPyObject(self, PyObj, entry, attach=False):
        '''Synchronise across workgroups'''

        from ctypes import byref, c_char_p, c_int

#        selectcases = { int  : c_int,
#                        float: c_double }

        # sum up all results (e.g., gradients)
        if not attach:
            ierror = libparallel.synchronise_python_object(byref(py2f(PyObj)),
                                                           c_char_p(entry.strip().encode()))

        # concatenate results (e.g., conn)
        else:
            ierror = libparallel.synchronise_python_object(byref(py2f(PyObj)),
                                                           c_char_p(entry.strip().encode()),
                                                           byref(c_int(attach)))


    def syncPyResult(self, entry='gradients'):
        '''Synchronise across workgroups'''

        from ctypes import byref, c_char_p
        from numpy  import ctypeslib

#        entry = selectcases[entry.lower()].encode()

        # Fortran: a character dummy argument with length other than 1 is not interoperable
#        ierror = libparallel.synchronise_python_object(byref(c_char_p(entry)))
        ierror = libparallel.synchronise_python_object()


    @property
    def workgroup(self):
        '''Workgroup ID'''
    
        return libparallel.chemsh_parallel_c_get_wg()


    @property
    def nworkgroups(self):
        '''Number of workgroups'''
    
        return libparallel.chemsh_parallel_c_get_nwgs()


    @property
    def nprocs_per_workgroup(self):
        '''Number of processes per workgroup'''

        return libparallel.chemsh_parallel_c_get_nnodes_wg()


    @property
    def nprocs_per_wg(self):
        '''Number of processes per workgroup'''

        return self.nprocs_per_workgroup


    @property
    def nprocs_mm(self):
        '''Number of processes for MM (subset in a workgroup)'''

        return libparallel.chemsh_parallel_c_get_nnodes_subset()


    @property
    def wg_dirname(self):
        '''Workgroup directory name: for example, _wg_0123'''

        from ctypes import byref, c_char_p
    
        return c_char_p(libparallel.chemsh_parallel_c_get_wg_dirname()).value.decode()


    @property
    def world_hostnames(self):
        '''Return a list containing all hostnames of each MPI process (within the current workgroup)'''

        from ctypes import c_char_p

        mangled_names = [ 'chemshellparallelmodule_mp_c_world_hostnames_',   # intel
                          'chemshellparallelmodule_mp_c_world_hostnames_',   # gnu
                          '__chemshellparallelmodule_MOD_c_world_hostnames', # gnu 10?
                        ]

        for mangled_name in mangled_names:
            try:
                cbuff_ptr = c_char_p.in_dll(libparallel, mangled_name)
            except:
                continue
        
        # get a string containing all hostnames of each MPI process (within the current workgroup), separated by ";"
        try:
            sbuff = cbuff_ptr.value.decode().split(';')
        except:
            print('\n >>> ERROR: parallel.world_hostnames failed')
            exit(999)

        return [ h for h in sbuff if h ]


    @property
    def workgroup_hostnames(self):
        '''Return a list containing all hostnames of each MPI process (within the current workgroup)'''

        from ctypes import c_char_p

        mangled_names = [ 'chemshellparallelmodule_mp_c_hostnames_',   # intel
                          '__chemshellparallelmodule_mp_c_hostnames',  # gnu
                          '__chemshellparallelmodule_MOD_c_hostnames', # gnu 10?
                        ]

        for mangled_name in mangled_names:
            try:
                cbuff_ptr = c_char_p.in_dll(libparallel, mangled_name)
            except:
                continue
        
        # get a string containing all hostnames of each MPI process (within the current workgroup), separated by ";"
        try:
            sbuff = cbuff_ptr.value.decode().split(';')
        except:
            print('\n >>> ERROR: parallel.workgroup_hostnames failed')
            exit(999)

        return [ h for h in sbuff if h ]


    @property
    def world_ranks(self):
        ''''''

        from ctypes import c_int, POINTER
        from numpy  import array

        mangled_names = [ 'chemshellparallelmodule_mp_c_world_ranks_',   # intel
                          '__chemshellparallelmodule_mp_c_world_ranks',  # gnu
                          '__chemshellparallelmodule_MOD_c_world_ranks', # gnu 10?
                        ]

        for mangled_name in mangled_names:
            try:
                ibuff_ptr = POINTER(c_int*self.nprocs_per_wg).in_dll(libparallel, mangled_name)
            except:
                continue

        return array(ibuff_ptr.contents[:])


    @property
    def ncpus(self):
        '''Number of CPU cores on the current node (where the current Python interpreter instance is started)'''

        from multiprocessing import cpu_count

        return cpu_count()


    @property
    def prefix(self):
        '''Return .. if the current job is task-farmed and the current working directory (CWD) is under one such as _wg_XXX'''

        from os import getcwd, path
        from pathlib import Path

        if "_wg_" in path.basename(getcwd()) and self.nworkgroups > 1:
            return '..'
        else:
            return ''


# create an instance
parallel = Parallel()

