#  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.utils import dictutils, fileutils, objutils

def parse(filename, _regex=None, receiver=None, prefix='', **kwargs):
    '''Parse a chemical json file'''

    from numpy   import asarray, count_nonzero
    from copy    import deepcopy
    from json    import dumps, loads
    from os      import path
    from ...data import atomicdata, unitconvert


    def __getEntry(_dictbuff, key):
        '''A wrapper around dictutils.getDictEntry()'''

        # a series of keys: only a tuple is unpacked using *
        if type(key) is tuple:
            result = dictutils.getDictEntry(_dictbuff, *key)

        # a single key
        else:
            result = dictutils.getDictEntry(_dictbuff, key)

        return result


    def __setNAtoms(_XYZObj, _dictbuff):
        '''Get natoms'''

        # npts first of all
        natomss = {}
        # strictly check if the lengths of all entries providing natoms info are consistent
        for key in [ 'cores', 'atomic numbers', 'atomic types', ('elements','number'), ('elements','type') ]:
            entrybuff = __getEntry(_dictbuff, key)
            if type(entrybuff) is list:
                natomss.update({key:len(entrybuff)})

        setbuff = set(natomss.values())
        # set natoms if consistent
        if len(setbuff) == 1 and list(setbuff)[0]:
            _XYZObj.natoms = list(setbuff)[0]
        # if inconsistent
        if len(setbuff) > 1:
            raise IOError("\n >>> ChemShell ERROR: inconsistent length of arrays in JSON file!\n", natomss)


    def __assignXYZ(_XYZObj, _dictbuff, _key, mandatory=False):
        '''A more generic tool to assign values to XYZ-type objects'''

        from operator import itemgetter

        # YL 16/02/2019: keep this order as since Python 3.6+ dictionaries are orderd
        #                `znums` has to be defined before `names`
        #                see below for a patch for 3.5 and earlier Python
        #                https://stackoverflow.com/questions/39980323/are-dictionaries-ordered-in-python-3-6
        selectcases = {
                        'a'                 :'a',
                        'b'                 :'b',
                        'c'                 :'c',
                        'alpha'             :'alpha',
                        'beta'              :'beta',
                        'gamma'             :'gamma',
                        'cores'             :'coreatoms',
                        '3d'                :'coords',
                        '3d fractional'     :'fractional',
                        'charges'           :'charges',
                        'convention'        :'convention',
                       ('elements','number'):'znums',
                        # TD 09/07/2024: labels is the "canonical" cjson entry for atom names, but ChemShell
                        #                will try a number of non-standard options too
                        'labels'            :'names',
                       ('elements','type'  ):'names',
                        'atomic types'      :'names',
                        'names'             :'names',
                        'vectors'           :'vectors'
                      }

        _dictbuff = getJSONEntry(_dictbuff, _key, mandatory=mandatory)

        if not _dictbuff:
            return

        # set natoms
        __setNAtoms(_XYZObj, _dictbuff)

        # unit conversion factor (from user unit to a.u.)
        try:
            factor = unitconvert.unit2au[dictutils.getDictEntry(_dictbuff, 'unit').lower()]
        except:
            factor = 1.0

        # YL 16/02/2019: this is a solution to 3.5 and previous Python which doesn't sort dictionaries --
        #                make a list from the above dict and keep the 'names' items in the end
        lbuff = list(selectcases.items())
        # YL 06/03/2019: needs a bit more sorting for Python 3.5
        lbuff.sort(key=itemgetter(1))
        # YL 13/03/2019: further fix as there was still a bug that during the loop `lbuff` was already changed
        lsorted = [ i for i in lbuff ]
        for key, val in lbuff:
            # place the `names` entries in the end because we don't want `names` to be rewritten
            if val == 'names':
                lsorted.append(lsorted.pop(lsorted.index((key,val))))

        # parse entries
        for key, val in lsorted:

            # only convert the unit for certain entries
            if val in [ 'coords', 'a', 'b', 'c' ]:
                convertunit = True
            else:
                convertunit = False

            # skip non-existing entries
            if not hasattr(_XYZObj, val):
                continue

            try:
                buff = __getEntry(_dictbuff, key)
                # avoid None
                if not buff:
                    raise ValueError

                # unit conversion
                if convertunit:
                    buff = asarray(buff)*factor

                setattr(_XYZObj, val, buff)

            except:
                if chemsh.CHEMSH_DEBUG > 5:
                    print("\n >>> skipping JSON entry %s...\n"%str(key))

    def __checkCJSON(filename, prefix, json_dict):
        ''' Check that a file being loaded is a chemical JSON rather than a generic JSON file'''
        from sys import exit, stdout

        # Can't assume prefix has been set
        # So work out file extension
        full_filename = path.join(prefix, filename)
        extension = path.splitext(full_filename)[1]
        # If the file has been explicitly called cjson
        # it is safe to assume it is a chemical JSON
        if extension == '.cjson':
            return

        # Otherwise check if chemical json variable has been set
        cjson_ver = None

        # Check for CJSON v0
        try:
            cjson_ver = json_dict["chemical json"]
        except KeyError:
            # Check for CJSON v1 if no v0 definition
            try:
                cjson_ver = json_dict["chemicalJson"]
            except KeyError:
                pass

        # Deal with the separate cases
        # We are happy with either v0 or v1
        # Any other number is non-standard
        if cjson_ver in [0,1]:
            return
        # Not valid CJSON
        elif cjson_ver is None:
            print('\n ERROR: Chemical JSON files should either have the .cjson extension or contain the line "chemical json": 0')
            stdout.flush()
            exit(999)

    # receiver is the fragment object in this case
    frag = receiver

    # save file contents in a string buffer
    # TWK: don't alter input case
    strbuff = fileutils.file2text(filename, prefix=prefix)

    # eval string as dict: eval(expression[, globals[, locals]])
# FIXME: this is a hack of JSON to allow string without quotes (specifically, atomic types),
#        however we should remove this since it involves non-standard JSON
#        extending JSONEncoder/JSONDecoder does not work!
    # to accept strings without quotes ("" or ''), let positional argument 'locals'
    #     be namespace created from atomicdata.symbols (namely, H='H', He='He',...)
    # this can only be realised after applied str.title() since some atomic symbols
    #     in lower case (In:in, As:as, etc.) are reserved by Python.
    dictbuff = eval(strbuff, None, literals2constants())

    # convert single quote to double
    # TWK don't alter input case
    strbuff = dumps(dictbuff)
    foobar = loads(strbuff)

    # TD 12/03/24 Check the file is a Chemical JSON
    __checkCJSON(filename, prefix, foobar)

    # add some trailing spaces to make sure no garbled characters enter the string passed to Fortran
# TODO: invent a generic method to do so, since also needed in, for example, <method> fragpointer
# FIXME: frag.title must be defined prior to other attributes, otherwise 
    # title (required)
#    frag.title = foobar["name"]+" "*255

    # step 1: initialise natoms (required)
    __assignXYZ(frag, foobar, 'atoms', mandatory=True)

    # step 2: read in lattice parameters/vectors
    __assignXYZ(frag.cell, foobar, 'unit cell', mandatory=False)
    # YL 15/10/2020: avogadro2 uses "unitCell"
    # TD 12/03/2024: This is CJSON v1, this is the only keyword that ChemShell knows about that changed
    # between the two versions (They changed from spaces to camel case, but this only matters for keywords
    # that are multiple words) Still worth trying both though because if a file has the cjson extension
    # but doesn't set "chemical json" we don't know which version we are reading ahead of time
    __assignXYZ(frag.cell, foobar, 'unitCell', mandatory=False)

    # step 3: read in fractional coordiantes ('3d fractional' is not part of 'unit cell' due to the chemical JSON protocol)
    #         only do this when there is a "unit cell" entry
    if count_nonzero(asarray(frag.cell.vectors)) or count_nonzero(asarray(frag.cell.consts)):
        __assignXYZ(frag.cell, foobar, 'atoms', mandatory=False)

    # step 4: block 'shells' (optional)
    # use the dummy <Cell> object to convert fractional coordinates to cartesian. NB: name mangling
    frag.shells._Shells__cell = objutils.getDeepCopy(frag._cell)
    __assignXYZ(frag.shells, foobar, 'shells')

    # step 5: coords (required)
    # have to do this one more time because the connectivity depends on shell charges!
    __assignXYZ(frag, foobar, 'atoms', mandatory=True)

    # TODO: bqs (optional)
    dictbuff = getJSONEntry(foobar, "bqs")

#    # TODO: connections (optional)
#    try:
#        nconns   = len(foobar["bonds"]["connections"]["index"])
#        frag.conn = asarray(foobar["bonds"]["connections"]["index"]).reshape(nconns/2,2)
#    except:
#        pass



def dump(filename, frag, **kwargs):
    ''''''

    from json              import dump
    from collections       import defaultdict
    from ...base.constants import MAXCONN

    # TODO If needed, this is a generic ChemShell JSON for generic objects
    #with open(filename, 'w') as fp:
    #     dump(objutils.obj2dict(obj), fp, sort_keys=False, indent=4)

    # Here we are saving a Fragment as Chemical JSON (with ChemShell-specific extensions)
    tree = lambda: defaultdict(tree)
    fj = tree()

    # System properties 
    fj["chemical json"] = 0
    fj["name"] = frag.title.strip()
    # TODO only write out if charge has been explicitly set
    fj["charge"] = frag.totalcharge

    # Cell
    if frag.cell.natoms:
        consts = frag.cell.consts
        fj["unit cell"]["unit"] = "a.u."
        fj["unit cell"]["vectors"] = frag.cell.vectors[:frag.cell.ndimensions,:].reshape(1,3*frag.cell.ndimensions).tolist()[0]

    # Atoms
    fj["atoms"]["labels"] = [x.decode() for x in frag.names.tolist()]
    fj["atoms"]["elements"]["number"] = frag.znums.tolist()
    fj["atoms"]["coords"]["unit"] = "a.u."
    fj["atoms"]["coords"]["3d"]     = frag.coords.reshape(1,3*frag.natoms).tolist()[0]
    if frag.hasCharges():
        fj["atoms"]["charges"] = frag.charges.tolist()
 
    # Shells (extension to Chemical JSON)
    # TODO Fractional shell coords
    if frag.nshells > 0:
        shells = frag.shells
        fj["shells"]["unit"] = "a.u."
        fj["shells"]["3d"] = shells.coords.reshape(1,3*frag.nshells).tolist()[0]
        # JSON does not like dtype=int64 apparently
        cores = shells.coreatoms
        fj["shells"]["cores"] =  [int(s) for s in cores]
        fj["shells"]["charges"] = shells.charges.tolist()

    # Point charges (extension to Chemical JSON)
    # TODO Fractional point charge coords
    if frag.nbqs > 0:
        bqs = frag.bqs 
        fj["bqs"]["unit"] = "a.u."
        fj["bqs"]["3d"] = bqs.coords.reshape(1,3*frag.nbqs).tolist()[0]
        fj["bqs"]["charges"] = bqs.charges.tolist()

    # Connectivity table
    nconn = 0
    for i in range(frag.natoms):
        for j in range(MAXCONN):
            if frag.conn[i][j] > i:
                nconn += 1
    if nconn > 0:
        fj["bonds"]["connections"]["index"] = []
        for i in range(frag.natoms):
            for j in range(MAXCONN):
                if frag.conn[i][j] > i:
                    fj["bonds"]["connections"]["index"].append(i)
                    # JSON does not like dtype=int64 apparently
                    fj["bonds"]["connections"]["index"].append(int(frag.conn[i][j]))


    with open(filename, 'w') as fp:
        dump(fj, fp, sort_keys=False, indent=4)


def literals2constants():
    ''''''

    from ...data  import atomicdata

    # YL 15/10/2020: in Avogadro2 there is a section "selected" with values `true` and `false`
    # which are useless in ChemShell but have to be defined
    false = False
    true  = True

    literal_constants = ','.join(atomicdata.symbols)

    literal_constants += ',' + ','.join(['true','false','none'])

    namespace = "(" + literal_constants + ")"

    list_constants = atomicdata.symbols + [ True, False, None ]

    exec(namespace + " = " + str(list_constants))

    return locals()

def getJSONEntry(dictbuff, key, key2=None, mandatory=False):
    '''Return an entry by retrieving up to two keywords'''

    # usage: does not return Error when 'mandatory' is False
    def _getOptional():
        # if mandatory raise error
        if mandatory:
            raise IOError("\nChemShell >>> ERROR: setting up '"+key+"' failed.\n")
        # if optional return None
        else:
            return None

    try:
        # if one keyword
        entry = dictbuff[key]
        # if two keywords
        if key2:
            key = key2
            entry = entry[key]
    except:
        entry = _getOptional()

    return entry

