#  Copyright (C) 2020 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>'

# YL 15/01/2021: receiver dict with FF scheme (potential) name as top-level key
def parse(filename, _regex=None, receiver=None, prefix='', **kwargs):
    ''''''

    from numpy    import all, any, arange, array, argmax, append, asarray, char, concatenate, delete, dtype, equal, genfromtxt, isin, searchsorted, setdiff1d, vstack, unique
    from os       import path
    from copy     import copy

    selectcases = { 'bond'       : [('I','U8'),('J','U8'),('Kb','f8'),('b0','f8')],
                    'angle'      : [('I','U8'),('J','U8'),('K','U8'),('Ktheta','f8'),('Theta0','f8'),('Kub','f8'),('S0','f8')],
                    'dihedral'   : [('I','U8'),('J','U8'),('K','U8'),('L','U8'),('Kchi','f8'),('n','i8'),('delta','f8')],
                    'improper'   : [('I','U8'),('J','U8'),('K','U8'),('L','U8'),('Kpsi','f8'),('psi0','f8')],
                    'vdw'        : [('I','U8'),('epsilon','f8'),('Rmin/2','f8'),('eps,1-4','f8'),('Rmin/2,1-4','f8')],
                    # inversion
                    # shell
                    'equivalence': [('text','U36')],
                  }

    # to make it consistent with other FF entries in Py-ChemShell
    entries = { 'bond'    : 'bonds',
                'angle'   : 'angles',
                'dihedral': 'dihedrals',
                'improper': 'improper',
                'vdw'     : 'nonbonded',
              }

    # YL 13/08/2022: path prefix
    filename = path.join(prefix, filename)

    with open(filename, 'r') as fp:

        lines    = char.strip(array(fp.readlines()))
        nlines   = lines.size
        linenums = arange(nlines)

        plines = lines[char.startswith(char.lower(lines), 'potential')]
        potentials = vstack(char.split(char.strip(plines)))[:,1]

        # YL 15/01/2021: support multiple schemes (potentials)
        for potential in potentials:

            dbuff = receiver.get(potential, {})

            # loop over sections
            for section, dt in selectcases.items():
    
                line0 = argmax(char.startswith(lines, '{} '.format(section.upper())))+1
                line1 = argmax(char.startswith(lines, 'END {}'.format(section.upper())))
                text = lines[line0:line1]
                if not len(text):
                    continue
    
                dt = dtype(dt)
                fields_dt = copy(dt.names)
                # initialise array of all residue names
                if section not in dbuff:
                    dbuff.update({section:array([], dtype=dt)})
                abuff = genfromtxt(text,
                                   dtype=dt,
                                   comments='#',
                                   usecols=tuple(range(len(dt.names))))

                # YL 05/10/2021: the genfromtxt() function removes symbols from the field names, for example 'epsilon', 'Rmin/2', 'eps,1-4', 'Rmin/2,1-4'
                #                to 'I', 'epsilon', 'Rmin2', 'eps14', 'Rmin214' and we need to restore them (dt is a reference to abuff.dtype)
                abuff.dtype.names = fields_dt
    
                entryname = entries.get(section, section)
                dbuff.update({entryname:append(dbuff.get(entryname, \
                                                               array([], dtype=abuff.dtype)), \
                                                  abuff)})
    
                # remove repeated
                remove_repeated = True
                entrydata = dbuff[entryname]
                fields = entrydata.dtype.names
                # not possible to cast to 2D str arrays due to the "unsafe" rule
                # dihedrals
                # YL 23/03/2021: we must consider the "n" term for periodic dihedrals!
                if 'I' in fields and 'J' in fields and 'K' in fields and 'L' in fields and 'n' in fields:
                    definitions = entrydata[list(fields[:4])+['n']]
                    inversions = entrydata[list(fields[:4][::-1])+['n']]
                # impropers
                elif 'I' in fields and 'J' in fields and 'K' in fields and 'L' in fields and 'n' not in fields:
                    definitions = entrydata[list(fields[:4])]
                    inversions = entrydata[list(fields[:4][::-1])]
                # angles
                elif 'I' in fields and 'J' in fields and 'K' in fields and 'L' not in fields:
                    definitions = entrydata[list(fields[:3])]
                    inversions = entrydata[list(fields[:3][::-1])]
                # bonds
                elif 'I' in fields and 'J' in fields and 'K' not in fields and 'L' not in fields:
                    definitions = entrydata[list(fields[:2])]
                    inversions = entrydata[list(fields[:2][::-1])]
                else:
                    if len(fields) > 1:
                        field = list(fields[0])
                        definitions = entrydata[field]
                        # not in use but needed below
                        inversions = entrydata[field]
                    # for example: 'equivalence' which does not appear in every scheme
                    else:
                        remove_repeated = False
    
                # remove inversively repeated excluding symmetric ones
                if remove_repeated:
                    # first remove simply repeated
                    # char.strip() is not necessary if the raw data are parsed properly (especially in CHARMM .prm files)
                    # YL NB: axis not available utill NumPy v1.13!
                    u, indices_unique = unique(char.strip(array(definitions.tolist())), axis=0, return_index=True)
                    dbuff[entryname] = entrydata[indices_unique]
                    definitions = definitions[indices_unique]
                    inversions  = inversions[indices_unique]
# YL 02/07/2022: this no longer works in NumPy 1.23 and we must convert them to arrays of strings for matching
#                    masks1 = in1d(definitions.tolist(), inversions.tolist(), assume_unique=True)
                    definitions_str_array = array(list(map(str, char.strip(array(definitions.tolist())))))
                    inversions_str_array  = array(list(map(str, char.strip(array(inversions .tolist())))))
                    masks1 = isin(definitions_str_array, inversions_str_array, assume_unique=True)
                    masks2 = all(char.strip(array(definitions.tolist()))==char.strip(array(inversions.tolist())), axis=1)
                    masks = masks1^masks2
                    indices_to_remove = arange(definitions.shape[0])[masks]
                    redundant = []
                    for i, left in enumerate(indices_to_remove):
                        for j, right in enumerate(indices_to_remove[i+1:], start=i+1):
                            if str(definitions[left]) == str(inversions[right]):
                                redundant.append(right)
                    dbuff[entryname] = delete(entrydata[indices_unique], redundant)


            receiver.update({potential:dbuff})
