'''Interface to cluster'''

#  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 numpy           import int64, float64, full, zeros, set_printoptions
from ctypes          import CDLL, byref, Structure, c_int
from os.path         import join as os_path_join
from sys             import modules, stdout
from ..base          import constants
from ..dl_py2f       import py2f
from ..objects       import field, options, fragment, bqs
from ..io            import keyword, file
from ..dl_py2f.utils import objutils, miscutils

#set_printoptions(precision=6)

# get a reference to the current module (so we can access the module-level entries)
clusterModule = modules[__name__]

# load Fortran cluster library: use cross-platform function os.path.join() 
libpath = os_path_join(chemsh.CHEMSH_LIB_DIR, 'libcluster.so')
try:
    libcluster = CDLL(libpath)
except OSError as e:
    print(" >>> ChemShell ERROR:", e)
    print("\n     CHEMSH_ARCH is \"{}\", but perhaps you were trying to run a different build?\n".format(chemsh.CHEMSH_ARCH)) 

# create an instance of <class Options>
optionsObj = options.Options()
fieldObj   = field.Field()

# YL: we keep it at the module level so won't lose it when a subprocedure is done
clusterRef = None
fragRef    = None
fieldRef   = None

debug = byref(c_int(chemsh.CHEMSH_DEBUG))


_kwargs = {
            'cluster'       :'cluster.coords',
            'crystal_type'  :'ionic',
            'nneighbours'   : zeros((1,)                    , dtype=int64),
            'neighbours'    : full((1,constants.MAXCONN), -1, dtype=int64),
            'neighb_vectors': zeros((1,3)                   , dtype=float64),
            'origin'        : zeros((3,)                    , dtype=float64), # only an array is accepted by PY2F but can be converted from int by user input
            'frac_origin'   : True,
          }

@miscutils.time
def construct(fragObj, *args, **kwargs):
    '''Constructs the cluster'''

    from sys         import stdout
    from ..objects   import fragment
    from ..io.banner import printModuleBanner

    printModuleBanner('ChemShell Cluster Cutting Workflow')

    _kwargs = {}

    # YL: we create the cluster object here and pass it as an argument in all methods
    clusterObj = fragment.Fragment(connect_mode=fragObj.connect_mode)

    # YL NB 04/02/2021: a finite cluster should not be periodic
    #                we explicitly switch PBC off because of the automatic PBC detection introduced recently
    clusterObj.cell.pbc = 'no'

    # 1st step: cut the cluster from the periodic fragment
    cutCluster(fragObj, clusterObj, *args, **kwargs)
    clusterObj._reconnect = True

    if optionsObj.crystal_type in [ 'molecular' ]:
        return clusterObj

    # 2nd step: adjust the charge on the cluster
    adjustCharge(fragObj, clusterObj, *args, **kwargs)

    # 3rd step: calculate the periodic electrostatics
    calcElectrostatics(fragObj, clusterObj, *args,**kwargs)

    # 4th step: fit the charges on bqs
    fitBQs(fragObj, clusterObj, *args, **kwargs)

    # YL 07/01/2020: reset clusterObj._reconnect
    clusterObj._reconnect = True

    print("\n >>> Leaving ChemShell CLUSTER Module...\n")
    stdout.flush()

    return clusterObj


def __finalise(proc):
    '''Standard procedures applied to all cluster functions'''

    # wrapper of the procedure
    def wrapper(*args, **kwargs):

        result = proc(*args, **kwargs)

        return result

    libcluster.cluster_finalise(debug)

    return wrapper


@miscutils.time
def cutCluster(fragObj, clusterObj,  *args, **kwargs):
    '''Cut the cluster from a periodic system'''

    from sys    import stdout
    from numpy  import arange
    from ..base import errors

    print("\n Cutting out cluster from supercell...\n")
    stdout.flush()

    _kwargs = {
                # YL 08/05/2018 FIXME: bizarre that GNU 7.2 segfaults when 'coordinated' is '' while Intel works fine so that we apply this hotfix of 'none' (in fact any non-blank string works) for now
                'coordinated'   :'none',
                'radius_cluster': 20.0,
              }

    # for internal use only
    _internals = {
                 }

# YL 23/06/2020: since it's a general issue, this problem has been fixed at the io.keywords level
    # Check if radius_cluster is a float; if not, resolve. AJL 31/05/2020
#    _typecast_float_radius_cluster(kwargs)

    # initialise
    if chemsh.CHEMSH_DEBUG > 4:
        print("\n >>> ChemShell/CLUSTER: calling __init()...\n")
        stdout.flush()
    __init(fragObj, _kwargs, _internals, **kwargs)
    if chemsh.CHEMSH_DEBUG > 4:
        print("\n >>> ChemShell/CLUSTER: __init() called.\n")
        stdout.flush()

    natoms  = c_int(0)
    nshells = c_int(0)

    if chemsh.CHEMSH_DEBUG > 4:
        print("\n >>> ChemShell/CLUSTER: calling __py2f()...\n")
        stdout.flush()

    # do PY2F before running Fortran
    __py2f(fragObj, clusterObj)

    if chemsh.CHEMSH_DEBUG > 4:
        print("\n >>> ChemShell/CLUSTER: calling libcluster.cluster_get_size()...\n")
        stdout.flush()

    stdout.flush()
    ierror = libcluster.cluster_get_size(byref(natoms), byref(nshells), debug)

    if not natoms:
        raise errors.ChemShNoAtomError(clusterObj)

    if chemsh.CHEMSH_DEBUG > 2:
        print("\n >>> New cluster's number of atoms: {}\n".format(natoms))
        stdout.flush()

    # YL 21/10/2017: resizing fragment has nothing to do with the special treatment of no connectivity
    clusterObj.natoms = natoms.value

    if nshells:
        clusterObj.shells.natoms = nshells.value

    # TWK: suppress connect until after coords set
    # YL 21/10/2017: py2f accesses all attributes which implicitly include `conn`, and the
    #                `connect()` method failed because effectcharges() required `shells`.
    #                added `nshells` check to prevent errors when `shells` is unavailable
    #                (see Fragment.effectcharges())
    # YL 21/10/2017: another trouble is that the process here is a special case because `conn` is accessed
    #                by py2f() before `coords` are set by Fortran (different from the Python setter)
    # manually switched conn off
    clusterObj._reconnect = False

    # YL: when a fragment array is resized (via resetting natoms in this case), a py2f procedure must be done
    #clusterModule.clusterRef = py2f(clusterObj)
    __py2f(fragObj, clusterObj)

    # YL TODO: getclustercoords should be renamed setClusterCoords;
    #          or a function getClusterCoords to return the coords array then the above manual _reconnect problem should be gone
    ierror = libcluster.cluster_get_coords(byref(clusterModule.clusterRef), debug)
    if chemsh.CHEMSH_DEBUG > 4:
        print("\n >>> New cluster's coordinates: {}\n".format(clusterObj.coords))
    stdout.flush()

    # copy names from periodic frag
    clusterObj.names[:] = fragObj.names[clusterObj.map[:]-1]

    # Set shellatoms array (delete when shellatoms/coreatoms refactored)
    # YL 17/04/2018: Fragment.shellatoms has a new definition now that is the list of indices of associated shell centres
    if nshells:
        clusterObj.shellatoms[clusterObj._shells.coreatoms] = arange(nshells.value)
    stdout.flush()

    # because connectivity is omitted above, we manually do it here
    if chemsh.CHEMSH_DEBUG > 2:
        print("\n >>> Calculating connectivity table of the new cluster...\n")
        stdout.flush()
    clusterObj.connect(force_reconnect=True)

    libcluster.cluster_finalise(debug)
    stdout.flush()


@miscutils.time
def adjustCharge(fragObj, clusterObj, *args, **kwargs):
    '''Adjust the charge on the cluster'''

    from sys import stdout

    print("\n >>> ChemShell Cluster: adjusting Charges...\n")
    stdout.flush()

    _kwargs = {
                # YL 01/08/2018: 'adjust_method' or 'adjust_mode' will be a better name for this variable
                'adjust_charge'   :'', # 'coordination_scaled' | 'valence_model'
                'valence_matrix'  : zeros((fragObj.nspecies,fragObj.nspecies), dtype=float64),
              }

    # YL: for internal use only
    _internals = {
                 }

    # initialise
    __init(fragObj, _kwargs, _internals, **kwargs)

    # do PY2F before running Fortran
    __py2f(fragObj, clusterObj)

    # now run interface
    libcluster.cluster_adjust_charges(byref(clusterModule.clusterRef), debug)

    # return __init(fragObj, clusterObj, _kwargs, *args, **kwargs, task=2)
    libcluster.cluster_finalise(debug)


@miscutils.time
def calcElectrostatics(fragObj, clusterObj, *args, **kwargs):
    '''Calculate the electrostatic pot on the periodic pot'''

    from sys import stdout

    print("\n >>> ChemShell CLUSTER: calculating electrostatics...\n")
    stdout.flush()

    _kwargs = {
                'adjust_mode'   : 1,
#                'valence_matrix': zeros(shape=(0,3), dtype=float64),
                'radius_cluster': 20.0,
                'grid_mode'     : 0,
                'grid_points'   : 0,
              }

    _internals = {
                 }
 
# YL 23/06/2020: since it's a general issue, this problem has been fixed at the io.keywords level
    # Check if radius_cluster is a float; if not, resolve. AJL 31/05/2020
#    _typecast_float_radius_cluster(kwargs)

    # initialise
    __init(fragObj, _kwargs, _internals, **kwargs)

    fieldObj.npoints = fragObj.natoms + fragObj.nshells

    # do PY2F before running Fortran
    __py2f(fragObj, clusterObj)

    # now run interface
    libcluster.cluster_calc_electrostatics(byref(clusterModule.fieldRef), debug)
          
    print("\n >>> ChemShell CLUSTER: electrostatic field:", fieldObj.efield)
    print("\n >>> ChemShell CLUSTER: field grid:"         , fieldObj.grid)
    print("\n >>> ChemShell CLUSTER: potential:"          , fieldObj.pot)

    #return __init(fragObj, clusterObj, _kwargs, *args, **kwargs, task=5)

    #__finalise()
    libcluster.cluster_finalise(debug)


#@__finalise
@miscutils.time
def fitBQs(fragObj, clusterObj,  *args, **kwargs):
    '''Fit background point charges'''

    from sys   import stdout
    from numpy import argwhere, nonzero, pi

    print("\n >>> ChemShell CLUSTER: fitting BQs...\n")
    stdout.flush()

    _kwargs = {
                'bq_density'    : 4,
                # AJL, 13/6/18. This a bit esoteric without knowing the code - only works when prefixed with "bq_"?
                # a weight factor for the target charge
                'bq_fit_weight' : 1.0, 
                'bq_layer'      : 0.0,
                'bq_margin'     : 3.5,
                'bq_symbol'     :'F',
                'radius_active' : 20.0,
                'radius_cluster': 40.0,
                'totalcharge'   : None,
              }

    _internals = {
                 }

# YL 23/06/2020: since it's a general issue, this problem has been fixed at the io.keywords level
    # Check if radius_cluster is a float; if not, resolve. AJL 31/05/2020
#    _typecast_float_radius_cluster(kwargs)

    # initialise
    __init(fragObj, _kwargs, _internals, **kwargs)

    # bq_density is related to the number of bqs added
    # If bq_density=n, the number of bqs in the surrounding shell 
    # is 4n^2 + 2 (3d case). 
    # In the 2d case number of bqs determined by bq_density and bq_layer
    # (extra rings of charges to accommodate the adsorption layer). 
    # Extra fitting sites are added above the surface to ensure a good
    # fit for the bq_layer charges.
    #
    # In addition there are 6 bqs placed far away from the cluster
    # to account for the 'floating zero' of the potential and electric
    # field

    # 3D
    if fragObj.cell.ndimensions == 3:
        clusterObj.bqs.nbqs = 4*optionsObj.bq_density**2 + 8

    # 2D
    # YL TODO: remove either the following Python or the identical replicated in Fortran
    elif fragObj.cell.ndimensions == 2:
        npolars  = 6
        step_z   = pi * (optionsObj.radius_cluster + optionsObj.bq_margin) * 2.0 / (4 * optionsObj.bq_density - 1)
        n_layer  = int((optionsObj.bq_layer + optionsObj.bq_margin) / step_z + 1)
        nlsqs    = (2 * optionsObj.bq_density * (optionsObj.bq_density + 1) + 1) + (4 * optionsObj.bq_density * n_layer)
        clusterObj.bqs.nbqs = npolars + nlsqs

    else:
        print("Crystal fragment dimensionality not supported")
        raise ValueError

    # charge constraint
    if optionsObj.totalcharge == None:
#        optionsObj.totalcharge = clusterObj.effectcharges.sum()
        # YL 20/08/2018: maybe we should unchain the total charge as we did before
        # optionsObj.bq_fit_weight = 1.0E-10
        # TWK 27/1/2020: set default behaviour to reproduce Tcl-ChemShell (i.e. no target charge in bq fit)
        # Rather than try to indirectly get lsq_constraint to zero to match TCS, instead we can set
        # bq_fit_weight to zero directly here, which has exactly the same effect in the bq fit routines
        optionsObj.bq_fit_weight = 0.0
    else:
        optionsObj.totalcharge -= clusterObj.effectcharges.sum()

    # field 
    if not clusterModule.fieldObj:
        clusterModule.fieldObj = field.Field()
    clusterModule.fieldObj.npoints = fragObj.natoms + fragObj.nshells

    # do PY2F before running Fortran
    __py2f(fragObj, clusterObj)

    # run the Fortran function
    libcluster.cluster_fit_bqs(byref(clusterModule.clusterRef),
                               byref(clusterModule.fieldRef),
                               debug)

    print("\n>>> BQ charges:\n", clusterObj.bqs.charges)

    # This hack transmutes the added BQs into "F atoms"
    # This mimics the behaviour of Tcl ChemShell (which in turn followed construct)
    # where bqs are directly set as F (eventually F5) for use in GULP
    # It would probably be better to handle this closer to the GULP interface
    # at which point this code can be deleted, but for now stay consistent with TCS
    # for testing purposes.
    # AJL, 12/06/18: The element choice should be softcoded.
    #tmpNames = ['F'] * clusterObj.bqs.nbqs
    tmpNames = [optionsObj.bq_symbol] * clusterObj.bqs.nbqs
    tmpFrag = fragment.Fragment(coords=clusterObj.bqs.coords,
                                charges=clusterObj.bqs.charges,
                                names=tmpNames)

    # YL 20/04/2018:
    clusterObj.insert(tmpFrag, clusterObj.natoms)

    if clusterObj.nshells:
        clusterObj.shellatoms = [ i+tmpFrag.natoms if i > clusterObj.natoms else i for i in clusterObj.shellatoms ]
        clusterObj._shells.coreatoms[clusterObj.shellatoms[nonzero(clusterObj.shellatoms+1)]] = argwhere(clusterObj.shellatoms+1).ravel()

    # FIXME 2017-08-03 xyz delete method is not currently working for bqs
    clusterObj.bqs = bqs.BQs()

    # connectivity must be recalculated after BQs are created
    clusterObj.connect(force_reconnect=True)

    #__finalise()
    libcluster.cluster_finalise(debug)



def __py2f(fragObj, clusterObj):
    ''''''

    #
    clusterModule.optionsRef = py2f(optionsObj)
    clusterModule.fragRef    = py2f(fragObj)
    clusterModule.fieldRef   = py2f(fieldObj)
    clusterModule.clusterRef = py2f(clusterObj)

    libcluster.cluster_init(byref(clusterModule.fragRef),
                            byref(clusterModule.optionsRef),
                            debug)


def __init(fragObj, __kwargs, __internals, **kwargs):
    '''A wrapper to run __setAttrs() and __initOptions()'''

    __setAttrs(__kwargs, __internals)
    __setOptions(fragObj, **kwargs)


# set attributes 
def __setAttrs(__kwargs, __internals):

    optionsObj._kwargs    = __kwargs
    optionsObj._internals = __internals


#@miscutils.time
# initialise
@objutils.init(optionsObj, inheritfrom=modules[__name__])
def __setOptions(fragObj, **kwargs):
    '''Initialise optionsObj'''

    from numpy import array, dot, set_printoptions, inf
    from .     import neighbours

    # let NumPy print complete array rather than "...."-shortened
    set_printoptions(threshold=inf)

    # read in keyword arguments: use strict=False because there is no way to check all the _kwargs in different functions at the beginning
    keyword.parseKeys(optionsObj, strict=False, **kwargs)

    # get origin
    for old_key in ['origin_atom', 'origin_point']:
        if old_key in kwargs:
            setattr(optionsObj, 'origin', array([kwargs[old_key]]))
            print("\n >>> Since v19.1 keyword '{}' has been changed to 'origin' (accepting either int or array value).".format(old_key))
            print("     Sorry for the inconvenience. The assigned value has been forwarded to 'origin', nonetheless.\n")

    # given value has to be converted to ndarray as it could be a list
    if array(optionsObj.origin).shape == (3,):
        if optionsObj.frac_origin:
            optionsObj.origin = array(optionsObj.origin, dtype=float64)
        # convert to fractional
        else:
            #optionsObj.origin = array(dot(optionsObj.origin, fragObj.cell._c2f_matrix), dtype=float64)
            optionsObj.origin = array(dot(optionsObj.origin, fragObj.cell.reciprocal.T), dtype=float64)
    # YL 22/11/2019: from now io.keyword.parseKeys() takes [val] instead of val if default type is list/array
    elif array(optionsObj.origin).shape == (1,):
        if optionsObj.origin[0]+1:
            origin, optionsObj.origin = optionsObj.origin[0], zeros((3,), dtype=float64)
            if optionsObj.frac_origin:
                optionsObj.origin[:] = fragObj.cell.fractional[origin]
            # convert to fractional
            else:
                #optionsObj.origin[:] = array(dot(fragObj.coords[origin], fragObj.cell._c2f_matrix), dtype=float64)
                optionsObj.origin[:] = array(dot(fragObj.coords[origin], fragObj.cell.reciprocal.T), dtype=float64)
        else:
            optionsObj.origin = zeros((3,), dtype=float64)

    # print("\n Fractional origin (a.u.): %20.10f %20.10f %20.10f"%(optionsObj.origin[0], optionsObj.origin[1], optionsObj.origin[2]))

    # get the neighbour list
    neighbours.setNeighbours(fragObj, optionsObj)
 

# FIXME: not in use?
def get_cluster_region(fragObj):
    '''Assigns the QM/MM regions to the cluster fragment'''

    selectcases = {
                    'region': ['qm_explicit','qm_implicit','active','atoms','all'],
                  }

# YL 23/06/2020: since it's a general issue, this problem has been fixed at the io.keywords level
#def _typecast_float_radius_cluster(kwargs):
#    '''Check if radius_cluster is a float, and if not typecast to prevent uncontrolled errors in the fortran code'''
#
#    if(not isinstance(kwargs['radius_cluster'], float)):
#        kwargs['radius_cluster'] = float(kwargs['radius_cluster'])
