#  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

# YL 13/10/2019: there are different situations we need to consider:
#                                                                             title
#                a) element symbols, e.g. "O", "Mg"                           True
#                b) symbols with region label, e.g. "O1", "Mg1"               True
#                c) forcefield type, e.g. "H2", "H21"                         True
#                d) forcefield type, e.g. "CA"                                False
#           TODO e) lower case, e.g. "mg"                                     False
# YL 19/02/2020: "FE" and "NPH", "FE" and "NA"
def names2znums(names, exceptions={}):
    '''Translate force field atom names to atomic numbers'''

    from string import digits, punctuation
    from numpy  import argsort, array, char, isin, ndarray, searchsorted, zeros
    from ..data import atomicdata


    # force to convert to numpy.ndarray
    if type(names) is not ndarray:
        names = array(names)
    # YL 19/02/2021: it could still be a 0D array
    if names.shape == ():
        names = names[None]

    # use str below
    try:
        names = char.strip(char.decode(names))
    except:
        pass

    # YL 13/02/2021: there are names starting with digital characters
    names = char.lstrip(names, '0123456789')

    # initialise a feeder array
    znums = zeros((names.size,), dtype=int)

    # truncate to the first two characters only
    twochars = char.strip(char.ljust(names, 2))
    onechars = char.strip(char.ljust(names, 1))
    # titled names
    istitle = char.istitle(twochars)
    exact_matches = isin(twochars, atomicdata.symbols_array)

    # algorithm: get positions where names are inserted to preserve the sorted order (searchsorted is pretty fast) and
    # get the original indices according to the positions
    try:
        # find exact matches of the 2-char truncation
        znums[istitle*exact_matches] = atomicdata.symbols_sort_indices[searchsorted(atomicdata.symbols_sorted_array, twochars[istitle*exact_matches])]
        # and then of 1-char truncation for those don't match the 2-char truncation
        znums[istitle*~exact_matches] = atomicdata.symbols_sort_indices[searchsorted(atomicdata.symbols_sorted_array, onechars[istitle*~exact_matches])]
    except:
        pass

    # non-titled names (forcefield types): take only the first letter
    try:
        znums[~istitle] = atomicdata.symbols_sort_indices[searchsorted(atomicdata.symbols_sorted_array, char.ljust(char.upper(names[~istitle]), 1))]
    except:
        pass

    # apply special rules
    for name, znum in exceptions.items():
        masks = isin(char.strip(char.upper(names)), [name.upper().strip()])
        znums[masks] = znum

    return znums


def znums2names(znums):
    '''Translate atomic numbers to atom names'''

    from ..data import atomicdata

    return atomicdata.symbols_bytearray.take(znums)


# YL 16/02/2021
def removeRepeats(a1, a2):
    '''Remove entries from FF array1 that also defined in FF array2'''

    from ..io  import format
    from numpy import absolute, arange, copy, delete, isin

    indices_a1 = arange(a1.size)

    fields = a1.dtype.fields

    # dihedrals
    # YL 23/03/2021: we must check 'n' as well for multiple periodic dihedral entries may share same I-J-K-L
    if all([i in fields for i in ['I','J','K','L','n']]):
        section = 'dihedral'
        # YL 23/03/2021: because DL_FIELD uses minus values to mark special dihedrals we must compare the absolute values
        ns_a1, ns_a2 = copy(a1['n']), copy(a2['n'])
        a1['n'], a2['n'] = absolute(ns_a1), absolute(ns_a2)
        masks_a1_in_a2  = isin(a1[['I','J','K','L','n']],a2[['I','J','K','L','n']])
        # we must swap the field values otherwise a1[['I','J','K','L','n']]==a2[['I','J','K','L','n']]
        # and at the same time a1[['I','J','K','L','n']]==a2[['L','K','J','I','n']] because it's just a view
        a2[['I','J','K','L','n']] = a2[['L','K','J','I','n']]
        masks_a1_in_a2 += isin(a1[['I','J','K','L','n']],a2[['I','J','K','L','n']])
        # restore
        a2[['I','J','K','L','n']] = a2[['L','K','J','I','n']]
        a1['n'], a2['n'] = ns_a1, ns_a2

    # impropers
    elif all([i in fields for i in ['I','J','K','L']]) and 'n' not in fields:
        section = 'improper'
        masks_a1_in_a2  = isin(a1[['I','J','K','L']],a2[['I','J','K','L']])
        # we must swap the field values otherwise a1[['I','J','K','L']]==a2[['I','J','K','L']]
        # and at the same time a1[['I','J','K','L']]==a2[['L','K','J','I']] because it's just a view
        a2[['I','J','K','L']] = a2[['L','K','J','I']]
        masks_a1_in_a2 += isin(a1[['I','J','K','L']],a2[['I','J','K','L']])
        # restore
        a2[['I','J','K','L']] = a2[['L','K','J','I']]

    # angles
    elif all([i in fields for i in ['I','J','K']]):
        section = 'angle'
        masks_a1_in_a2  = isin(a1[['I','J','K']],a2[['I','J','K']])
        # we must swap the field values otherwise a1[['I','J','K']]==a2[['I','J','K']]
        # and at the same time a1[['I','J','K']]==a2[['K','J','I']] because it's just a view
        a2[['I','J','K']] = a2[['K','J','I']]
        masks_a1_in_a2 += isin(a1[['I','J','K']],a2[['I','J','K']])
        # restore
        a2[['I','J','K']] = a2[['K','J','I']]

    # bonds
    elif all([i in fields for i in ['I','J']]):
        section = 'bond'
        masks_a1_in_a2  = isin(a1[['I','J']],a2[['I','J']])
        # we must swap the field values otherwise a1[['I','J']]==a2[['I','J']]
        # and at the same time a1[['I','J']]==a2[['J','I']] because it's just a view
        a2[['I','J']] = a2[['J','I']]
        masks_a1_in_a2 += isin(a1[['I','J']],a2[['I','J']])
        # restore
        a2[['I','J']] = a2[['J','I']]

    # nonbonded
    elif all([i in fields for i in ['I']]):
        section = 'nonbonded'
        masks_a1_in_a2 = isin(a1['I'],a2['I'])

    # nothing to delete
    else:
        return a1

    # construct array of indices to remove
    indices_to_remove = indices_a1[masks_a1_in_a2]
    nindices_to_remove = indices_to_remove.size

    if nindices_to_remove:

        if chemsh.CHEMSH_DEBUG:
            print('\n'+format.wrap('Deleting existing {} forcefield {}:'.format(section, {1:'entry'}.get(nindices_to_remove, 'entries'))))

            # print dihedrals and impropers
            if all([i in fields for i in ['I','J','K','L']]):
                fmt = ' {} {} {} {}'
                for i in indices_to_remove:
                    print(fmt.format(a1[i]['I'],a1[i]['J'],a1[i]['K'],a1[i]['L']))
            # print angles
            elif all([i in fields for i in ['I','J','K']]):
                fmt = ' {} {} {}'
                for i in indices_to_remove:
                    print(fmt.format(a1[i]['I'],a1[i]['J'],a1[i]['K']))
            # print bonds
            elif all([i in fields for i in ['I','J']]):
                fmt = ' {} {}'
                for i in indices_to_remove:
                    print(fmt.format(a1[i]['I'],a1[i]['J']))
            # print nonbonded
            elif all([i in fields for i in ['I']]):
                fmt = ' {}'
                for i in indices_to_remove:
                    print(fmt.format(a1[i]['I']))

        return delete(a1, indices_to_remove)

    else:

        return a1

























