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

__author__ = 'You Lu <you.lu@stfc.ac.uk>'

from ...base          import constants
from ...dl_py2f.utils import miscutils

contents = [ 'atom', 'bond', 'angle', 'dihedral' ]


# YL IMPORTANT 04/11/2021: we assume the data dict used here is already unpacked
#                          with repeated molecules and repeated atoms of molecule!
@miscutils.time
def excludeQM(data, qm_atoms, qmmm_pairs):
    '''Exclude terms of certain QM-MM interactions from the given FIELD data'''

    from numpy import append, arange, array, char, delete, empty, isin
    from sys import stdout

    qm_atoms_from_1 = array(qm_atoms) + 1
    # YL 05/11/2021: we must convert the list of QM-JQ-MM pairs to a structured array to use in1d()
    abuff_qmmm_pairs = array(qmmm_pairs)
    qmmm_pairs_from_1 = empty((abuff_qmmm_pairs.shape[0],), dtype=[('I','i4'),('J','i4')])
    qmmm_pairs_from_1['I'] = abuff_qmmm_pairs[:,0] + 1
    qmmm_pairs_from_1['J'] = abuff_qmmm_pairs[:,1] + 1

    for key, val in data.items():

        if key in [ 'ntypes', 'vdw', 'unit' ]:
            continue

        ### bond ###

        dtype_bond = val['bond']['data'].dtype
        Is = val['bond']['data']['I']
        Js = val['bond']['data']['J']
    
        # exclude QM-QM
        # NB: indices in the data dictionary start with 1 as in DL_POLY
        masks_qm_qm = isin(Is, qm_atoms_from_1)*isin(Js, qm_atoms_from_1)
        # NB: do not index using val['bond']['data'][masks_qm_qm]['Kb'] which doesn't return a reference
        val['bond']['data']['Kb'][masks_qm_qm] = 0.0
        val['bond']['data']['b0'][masks_qm_qm] = 1.0
    
        # exclude QM-(JQ)-MM: 
        # YL 08/11/2021: this should be unnecessary assuming all I-J pairs can't be QM-(JQ)-MM
        bonds_IJ = val['bond']['data'][['I','J']]
        bonds_JI = val['bond']['data'][['J','I']]
        qjm_IJ   = qmmm_pairs_from_1[['I','J']]
        qjm_JI   = qmmm_pairs_from_1[['J','I']]
        # YL 02/07/2022: in1d(bonds_IJ, qjm_IJ) no longer works under NumPy 1.23
        # YL 07/07/2022: note that this has made the result different!
        bonds_IJ_str_array = list(map(str, val['bond']['data'][['I','J']]))
        bonds_JI_str_array = list(map(str, val['bond']['data'][['J','I']]))
        qjm_IJ_str_array = list(map(str, qjm_IJ.tolist()))
        qjm_JI_str_array = list(map(str, qjm_JI.tolist()))
#        masks_bonds_in_QJM = in1d(bonds_IJ, qjm_IJ) + in1d(bonds_IJ, qjm_JI)
        masks_bonds_in_QJM = isin(bonds_IJ_str_array, qjm_IJ_str_array) + isin(bonds_IJ_str_array, qjm_JI_str_array)
        # NB: do not index using val['bond']['data'][masks_qm_qm]['Kb'] which doesn't return a reference
        val['bond']['data']['Kb'][masks_bonds_in_QJM] = 0.0
        val['bond']['data']['b0'][masks_bonds_in_QJM] = 1.0
        # YL 07/07/2022: should be 'harm'
        val['bond']['data']['type'][masks_bonds_in_QJM] = 'harm'
        # YL 05/11/2021: QM-(J)-MM not found in existing Is/Js should be added to array
#        masks_QJM_in_bonds = in1d(qjm_IJ, bonds_IJ) + in1d(qjm_IJ, bonds_JI)
        masks_QJM_in_bonds = isin(qjm_IJ_str_array, bonds_IJ_str_array) + isin(qjm_IJ_str_array, bonds_JI_str_array)
        nqjms = (~masks_QJM_in_bonds).sum()
        extra_qjm = empty((nqjms,), dtype=val['bond']['data'].dtype)
        extra_qjm['type'] = 'harm'
        extra_qjm['Kb'] = 0.0
        extra_qjm['b0'] = 1.0
        extra_qjm['I'] = qmmm_pairs_from_1['I'][~masks_QJM_in_bonds]
        extra_qjm['J'] = qmmm_pairs_from_1['J'][~masks_QJM_in_bonds]

        # exclude all QM-QM combinations
        # bond dtype: [('type','U8'),('I','i4'),('J','i4'),('Kb','f8'),('b0','f8')]
        extra_qm_qm = array([], dtype=val['bond']['data'].dtype)
        for i, I in enumerate(qm_atoms_from_1):
            for J in qm_atoms_from_1[i+1:]:
                abuff = empty((1,), dtype=val['bond']['data'].dtype)
                abuff[0] = 'harm', I, J, 0.0, 1.0
                extra_qm_qm = append(extra_qm_qm, abuff)

        # also append the above extra_qjm
        extra_qm_qm = append(extra_qm_qm, extra_qjm)

        # delete the existing QM-QM terms first
        # NB: boolean array not taken until NumPy v1.19!
        to_delete = arange(val['bond']['nentries'])[masks_qm_qm]
        array = delete(val['bond']['data'], to_delete)
        val['bond'].update({'data':array})

        # append to form a new array
        val['bond'].update({'data':append(val['bond']['data'], extra_qm_qm)})
        # update the number of data
        nqqs = int(qm_atoms_from_1.size*(qm_atoms_from_1.size-1)/2)
        val['bond'].update({'nqqs':nqqs,'nqjms':nqjms,'nqqs_existing':masks_qm_qm.sum()})
        val['bond'].update({'nentries':val['bond']['nentries']-masks_qm_qm.sum()+nqjms+nqqs})

        ### angle ###
    
        Is = val['angle']['data']['I']
        Js = val['angle']['data']['J']
        Ks = val['angle']['data']['K']
    
        # exclude Q-Q-Q
        masks_qm_qm_qm = isin(Is, qm_atoms_from_1)*isin(Js, qm_atoms_from_1)*isin(Ks, qm_atoms_from_1)
    
        # exclude Q-Q-Q, Q-Q-M, and M-Q-Q
        masks_qm_qm_mm = isin(Is, qm_atoms_from_1)*isin(Js, qm_atoms_from_1)*~isin(Ks, qm_atoms_from_1)
    
        # exclude Q-Q-Q, Q-Q-M, and M-Q-Q
        masks_mm_qm_qm = ~isin(Is, qm_atoms_from_1)*isin(Js, qm_atoms_from_1)*isin(Ks, qm_atoms_from_1)
    
        # NB: boolean array not taken until NumPy v1.19!
        to_delete = arange(val['angle']['nentries'])[masks_qm_qm_qm+masks_qm_qm_mm+masks_mm_qm_qm]
        array = delete(val['angle']['data'], to_delete)
        val['angle'].update({'data':array})
        ndeletions = masks_qm_qm_qm.sum() + masks_qm_qm_mm.sum() + masks_mm_qm_qm.sum()
        nentries = val['angle']['nentries'] - ndeletions
        val['angle'].update({'nentries':nentries})
        val['angle'].update({'ndeletions':ndeletions})
        
        ### dihedral ###
    
        Is = val['dihedral']['data']['I']
        Js = val['dihedral']['data']['J']
        Ks = val['dihedral']['data']['K']
        Ls = val['dihedral']['data']['L']
    
        # exclude Q-Q-Q-Q
        masks_qm_qm_qm_qm = isin(Is, qm_atoms_from_1)*isin(Js, qm_atoms_from_1)*isin(Ks, qm_atoms_from_1)*isin(Ls, qm_atoms_from_1)
    
        # exclude Q-Q-Q-M
        masks_qm_qm_qm_mm = isin(Is, qm_atoms_from_1)*isin(Js, qm_atoms_from_1)*isin(Ks, qm_atoms_from_1)*~isin(Ls, qm_atoms_from_1)
    
        # exclude M-Q-Q-Q
        masks_mm_qm_qm_qm = ~isin(Is, qm_atoms_from_1)*isin(Js, qm_atoms_from_1)*isin(Ks, qm_atoms_from_1)*isin(Ls, qm_atoms_from_1)
    
        # exclude M-Q-Q-M
        masks_mm_qm_qm_mm = ~isin(Is, qm_atoms_from_1)*isin(Js, qm_atoms_from_1)*isin(Ks, qm_atoms_from_1)*~isin(Ls, qm_atoms_from_1)
    
        # NB: boolean array not taken until NumPy v1.19!
        to_delete = arange(val['dihedral']['nentries'])[masks_qm_qm_qm_qm+masks_qm_qm_qm_mm+masks_mm_qm_qm_qm+masks_mm_qm_qm_mm]
        array = delete(val['dihedral']['data'], to_delete)
        val['dihedral'].update({'data':array})
        ndeletions = masks_qm_qm_qm_qm.sum() + masks_qm_qm_qm_mm.sum() + masks_mm_qm_qm_qm.sum() + masks_mm_qm_qm_mm.sum()
        nentries = val['dihedral']['nentries'] - ndeletions
        val['dihedral'].update({'nentries':nentries})
        val['dihedral'].update({'ndeletions':ndeletions})
   
        # charges for Coulombic
        val['atom']['data']['charge'][qm_atoms] = 0.0



# because in FIELD indices restart from 1 in each molecule we have to ravel out the data
def unravelData(dict_data, reduce=False):
    '''Flatten the data with repeated molecules and molecule data with repeated atoms and return a new dict'''

    from numpy import arange, argsort, concatenate, empty, repeat, tile
    from ...io import files

    # dict of dtypes
    contents2dtypes = files.dlp_field.contents2dtypes

    # create a dict buffer as result to return
    dict_to_return = {}
    # YL 05/11/2021: for non-reduced form we have to use temporary arrays
    if not reduce:
        dict_abuffs = {}
    title = ''

    # van der Waals
    # YL 05/11/2021: we introduce a switch wo return a reduced version of data for convenient use or not
    if reduce:
        dict_to_return.update({'vdw':dict_data['vdw']['data']})
    else:
        dict_to_return.update({'vdw':dict_data['vdw']})
    for content in contents:
        if reduce:
            dict_to_return.update({content:empty(shape=(0,), dtype=contents2dtypes[content])})
        else:
            dict_abuffs.update({content:empty(shape=(0,), dtype=contents2dtypes[content])})

    natoms = 0
    for key, val in dict_data.items():
        # process only molecule(s)
        if type(key) is int:
            # concat the titles
            if key:
                title += '+'+val['title'].strip('name ')
            else:
                title += val['title'].strip('name ')
            nmolecules = val['nmolecules']
            natoms_mol = val['atom']['data']['nrepeats'].sum()
            for content in contents:
                # skip when unavailable
                if 'data' not in val[content]:
                    continue
                # tile the array by the number of molecules
                # single-record array
                if val[content]['data'].shape == ():
                    abuff = tile(val[content]['data'][None], nmolecules)
                    nentries = 1
                else:
                    abuff = tile(val[content]['data'], nmolecules)
                    nentries = val[content]['data'].shape[0]
                # repeat each molecule by the number of atoms and add the increment
                offset = repeat(arange(0,nmolecules*natoms_mol,natoms_mol), nentries) + natoms
                # add the increments to atom indices
                for i in [ 'I', 'J', 'K', 'L' ]:
                    if i not in val[content]['data'].dtype.names:
                        continue
                    # add the previous total number of atoms
                    abuff[i] += offset
                # append to the existing data
                if reduce:
                    dict_to_return.update({content:concatenate((dict_to_return[content],abuff))})
                else:
                    dict_abuffs.update({content:concatenate((dict_abuffs[content],abuff))})
            natoms += natoms_mol*nmolecules


#    if not reduce:
#        sorter_bond = argsort(dict_abuffs['bond'], order=['type','I','J'])
#        dict_abuffs.update({'bond':dict_abuffs['bond'][sorter_bond]})

    # if the original structure of dict_data is to be kept
    if not reduce:
        for key, val in dict_data.items():
            # process only 'unit', 'ntypes', and 'vdw'
            if type(key) is not int:
                dict_to_return.update({key:dict_data[key]})
        # merge molecules
        dict_to_return.update({'ntypes':1})
        # first copy all
        dict_to_return.update({0:dict_data[0]})
        dict_to_return[0].update({'nmolecules':1})
        dict_to_return[0].update({'title':title})
        for content in contents:
            dict_to_return[0][content].update({'data':dict_abuffs[content]})
            dict_to_return[0][content].update({'nentries':dict_abuffs[content].size})
        dict_to_return[0]['atom']['data']['nrepeats'] = 1

    return dict_to_return
    

# YL TODO 27/10/2021: we also want to compare the values
@miscutils.time
def compareTwoFIELDs(f1, f2, tolerance=constants.MEDIUM):
    '''Compare two FIELD files to find the difference'''

    from numpy   import any, arange, argsort, argwhere, char, isin, set_printoptions, unique, zeros
    from sys     import maxsize, stdout
    from ...io   import files
    from ...base import forcefield


    def _compareEntryUnsorted(_a1, _a2, _content, _entry, _indices):
        '''Compare two arrays' specific entry'''

        _n1 = _a1.shape[0]
        _n2 = _a2.shape[0]
        if _n1 == _n2:
            size = int((_n1+_n2)/2)
            for i in range(size):
                r1 = _a1[i]
                r2 = _a2[i]
                _index = _indices[i] + 1
                if r1[_entry] != r2[_entry]:
                    print(f'{_content} {_index}:', r1, r2)
        stdout.flush()


    def _compareEntrySorted(_a1, _a2, _content, _order1=[], _order2=[], _entries=[]):
        '''Compare two arrays' specific entries, sorted according to _order (list of fields)'''

        _n1 = _a1.shape[0]
        _n2 = _a2.shape[0]
        if _n1 == _n2:
            size = int((_n1+_n2)/2)
            _sorter_ab1 = argsort(_a1, order=_order1)
            _sorter_ab2 = argsort(_a2, order=_order2)
            _s1 = _a1[_sorter_ab1]
            _s2 = _a2[_sorter_ab2]
            for _entry in _entries:
                print(f'\n {_content} containing different {_entry}:')
                for i in range(size):
                    _r1 = _s1[i]
                    _r2 = _s2[i]
                    if _r1[_entry] != _r2[_entry]:
                        print(_r1, _r2)
        stdout.flush()


    set_printoptions(threshold=maxsize)

    dict_both = {}
    for filename in [ f1, f2 ]:

        # read in the data
        print(f'\n Reading in data from file {filename}...')
        dbuff = files.dlp_field.parse(filename)

        dict_file = unravelData(dbuff, reduce=True)
        dict_both.update({filename:dict_file})

    # compare atoms
    print('\n Comparing atoms...')
    print(' ^^^^^^^^^^^^^^^^^^')
    a1 = dict_both[f1]['atom']
    a2 = dict_both[f2]['atom']
    # a1 in a2
    masks1 = isin(a1['name'],a2['name'])
    indices = argwhere(masks1).ravel()
    nmissings = (~masks1).sum()
    if nmissings:
        print(f'\n {nmissings} missing {f1} atoms in {f2}:')
        print(a1[~masks1])
    else:
        print(f'\n No missing {f1} atoms in {f2}:')
    # a2 in a1
    masks2 = isin(a2['name'],a1['name'])
    nmissings = (~masks2).sum()
    if nmissings:
        print(f'\n {nmissings} missing {f2} atoms in {f1}:')
        print(a2[~masks2])
    else:
        print(f'\n No missing {f2} atoms in {f1}:')
    # compare type and values
    print('\n Atoms containing different charges')
    print(f' {f1}                             {f2}')
    _compareEntryUnsorted(a1[masks1], a2[masks2], 'atom', 'charge', indices)

    # compare bonds
    print('\n Comparing bonds...')
    print(' ^^^^^^^^^^^^^^^^^^')
    a1 = dict_both[f1]['bond']
    a2 = dict_both[f2]['bond']
    # a1 in a2
    missing_a1_in_a2 = forcefield.removeRepeats(a1, a2)
    nmissings = missing_a1_in_a2.size
    if nmissings:
        print(f'\n {nmissings} missing {f1} bonds in {f2}:')
        print(missing_a1_in_a2)
    else:
        print(f'\n No missing {f1} bonds in {f2}')
    # a2 in a1
    missing_a2_in_a1 = forcefield.removeRepeats(a2, a1)
    nmissings = missing_a2_in_a1.size
    if nmissings:
        print(f'\n {nmissings} missing {f2} bonds in {f1}:')
        print(missing_a2_in_a1)
    else:
        print(f'\n No missing {f2} bonds in {f1}')
    print(f' {f1}                             {f2}')
    stdout.flush()
    # YL 23/08/2023: due to the recent change in NumPy we have to convert these arrays of tuples (generated by fancy indexing) to arrays of strings
    # a1 in a2
    masks1 = isin(char.mod('%s',a1[['I','J']]), char.mod('%s',a2[['I','J']]))
    masks2 = isin(char.mod('%s',a2[['I','J']]), char.mod('%s',a1[['I','J']]))
    _compareEntrySorted(a1[masks1], a2[masks2], 'bonds', _entries=['type','Kb','b0'], _order1=['I','J'], _order2=['I','J'])
    # a2 in a1
    masks1 = isin(char.mod('%s',a1[['I','J']]), char.mod('%s',a2[['J','I']]))
    masks2 = isin(char.mod('%s',a2[['I','J']]), char.mod('%s',a1[['J','I']]))
    _compareEntrySorted(a1[masks1], a2[masks2], 'bonds', _entries=['type','Kb','b0'], _order1=['I','J'], _order2=['J','I'])

    # compare angles
    print('\n Comparing angles...')
    print(' ^^^^^^^^^^^^^^^^^^^')
    a1 = dict_both[f1]['angle']
    a2 = dict_both[f2]['angle']
    # a1 in a2
    missing_a1_in_a2 = forcefield.removeRepeats(a1, a2)
    nmissings = missing_a1_in_a2.size
    if nmissings:
        print(f'\n {nmissings} missing {f1} angles in {f2}:')
        print(missing_a1_in_a2)
    else:
        print(f'\n No missing {f1} angles in {f2}')
    print(f' {f1}                             {f2}')
    stdout.flush()
    masks1 = isin(char.mod('%s',a1[['I','J','K']]), char.mod('%s',a2[['I','J','K']]))
    masks2 = isin(char.mod('%s',a2[['I','J','K']]), char.mod('%s',a1[['I','J','K']]))
    _compareEntrySorted(a1[masks1], a2[masks2], 'angles', _entries=['type','K','Ktheta','Theta0'], _order1=['I','J','K'], _order2=['I','J','K'])
    # a2 in a1
    missing_a2_in_a1 = forcefield.removeRepeats(a2, a1)
    nmissings = missing_a2_in_a1.size
    if nmissings:
        print(f'\n {nmissings} missing {f2} angles in {f1}:')
        print(missing_a2_in_a1)
    else:
        print(f'\n No missing {f2} angles in {f1}')
    stdout.flush()
    masks1 = isin(char.mod('%s',a1[['I','J','K']]), char.mod('%s',a2[['K','J','I']]))
    masks2 = isin(char.mod('%s',a2[['I','J','K']]), char.mod('%s',a1[['K','J','I']]))
    _compareEntrySorted(a1[masks1], a2[masks2], 'angles', _entries=['type','K','Ktheta','Theta0'], _order1=['I','J','K'], _order2=['K','J','I'])

    # compare dihedrals
    print('\n Comparing dihedrals...')
    print(' ^^^^^^^^^^^^^^^^^^^^^^')
    a1 = dict_both[f1]['dihedral']
    a2 = dict_both[f2]['dihedral']
    # a1 in a2
    missing_a1_in_a2 = forcefield.removeRepeats(a1, a2)
    nmissings = missing_a1_in_a2.size
    if nmissings:
        print(f'\n {nmissings} missing {f1} dihedrals in {f2}:')
        print(missing_a1_in_a2)
    else:
        print(f'\n No missing {f1} dihedrals in {f2}:')
    print(f' {f1}                             {f2}')
    stdout.flush()
    masks1 = isin(char.mod('%s',a1[['I','J','K','L']]), char.mod('%s',a2[['I','J','K','L']]))
    masks2 = isin(char.mod('%s',a2[['I','J','K','L']]), char.mod('%s',a1[['I','J','K','L']]))
    _compareEntrySorted(a1[masks1], a2[masks2], 'dihedrals', _entries=['type','Kchi','delta','n','14e','14vdw'], _order1=['I','J','K','L'], _order2=['I','J','K','L'])
    # a2 in a1
    missing_a2_in_a1 = forcefield.removeRepeats(a2, a1)
    nmissings = missing_a2_in_a1.size
    if nmissings:
        print(f'\n {nmissings} missing {f2} dihedrals in {f1}:')
        print(missing_a2_in_a1)
    else:
        print(f'\n No missing {f2} dihedrals in {f1}:')
    stdout.flush()
    masks1 = isin(char.mod('%s',a1[['I','J','K','L']]), char.mod('%s',a2[['L','K','J','I']]))
    masks2 = isin(char.mod('%s',a2[['I','J','K','L']]), char.mod('%s',a1[['L','K','J','I']]))
    _compareEntrySorted(a1[masks1], a2[masks2], 'dihedrals', _entries=['type','Kchi','delta','n','14e','14vdw'], _order1=['I','J','K','L'], _order2=['L','K','J','I'])

    # compare vdw
    print('\n Comparing vdWs...')
    print(' ^^^^^^^^^^^^^^^^^')
    a1 = dict_both[f1]['vdw']
    a2 = dict_both[f2]['vdw']
#    masks1 = in1d(a1['A'],a2['A'])*in1d(a1['B'],a2['B'])
#    masks2 = in1d(a1['A'],a2['B'])*in1d(a1['B'],a2['A'])
    
    # a1 in a2
    masks_ab1 = isin(a1['A'],a2['A'])*isin(a1['B'],a2['B'])
    masks_ba1 = isin(a1['A'],a2['B'])*isin(a1['B'],a2['A'])
    masks1 = masks_ab1 + masks_ba1

    nmissings = (~masks1).sum()
    if nmissings:
        print(f'\n {nmissings} missing {f1} vdWs in {f2}:')
        print(a1[~masks1])
    else:
        print(f'\n No missing {f1} vdWs in {f2}')
    # a2 in a1
    masks_ab2 = isin(a2['A'],a1['A'])*isin(a2['B'],a1['B'])
    masks_ba2 = isin(a2['A'],a1['B'])*isin(a2['B'],a1['A'])
    masks2 = masks_ab2 + masks_ba2
    nmissings = (~masks2).sum()
    if nmissings:
        print(f'\n {nmissings} missing {f2} vdWs in {f2}:')
        print(a2[~masks2])
    else:
        print(f'\n No missing {f2} vdWs in {f2}')
    # compare type and values
    _compareEntrySorted(a1[masks_ab1], a2[masks_ab2], 'vdWs', _entries=['type','a','b'], _order1=['A','B'], _order2=['A','B'])

#    if masks_ba1.size == masks_ba2.size:
#        size = int((masks_ba1.size+masks_ba2.size)/2)
#        sorter_ba1 = argsort(a1[masks_ba1], order=['A','B'])
#        sorter_ba2 = argsort(a2[masks_ba2], order=['A','B'])
#        s1 = a1[masks_ba1][sorter_ba1]
#        s2 = a2[masks_ba2][sorter_ba2]
#        for i in range(size):
#            r1 = s1[i]
#            r2 = s2[i]
#            if r1['type'] != r2['type']:
#                print(r1,r2)
#        for i in range(size):
#            r1 = s1[i]
#            r2 = s2[i]
#            if r1['a']-r2['a']>tolerance or r1['b']-r2['b']>tolerance:
#                print(r1,r2)
#
#
