#  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>'
    

from ...dl_py2f.utils import fileutils

# 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 arange, array, argmax, append, asarray, char, delete, dtype, genfromtxt, searchsorted, vstack
    from os       import path

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

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

        lines = char.strip(array(fp.readlines()))

        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, {})
    
            # ATOM_TYPE
            line0 = argmax(char.startswith(lines, 'ATOM_TYPE'))+1
            line1 = argmax(char.startswith(lines, 'END ATOM_TYPE'))
            # remove the leading and trailing '-'
            abuff = genfromtxt(lines[line0:line1],
                               dtype=dtype([('type','U36'),('key','U12'),('element','U2'),('mass','f8')]),
                               comments='#',
                               usecols=range(4))
            dt = abuff.dtype
            dbuff.update({'atomtypes':append(dbuff.get('atomtypes', array([], dtype=dt)), abuff)})
    
            # MOLECULE_TYPE
            # initialise array of all residue names
            dt = dtype([('type','U36'),('key','U36'),('mass','f8')])
            if 'residues' not in dbuff:
                dbuff.update({'residues':array([], dtype=dt)})
            line0 = argmax(char.startswith(lines, 'MOLECULE_TYPE')) + 1
            line1 = argmax(char.startswith(lines, 'END MOLECULE_TYPE'))
            # remove the leading and trailing '-'
            abuff = genfromtxt(lines[line0:line1],
                               dtype=dt,
                               comments='#',
                               usecols=range(3))
            # rename keys already in the databank: DL_FIELD SF should not override user-defined CHARMM RTF
            for i, residue in enumerate(abuff):
                if residue['key'] in dbuff['residues']['key']:
                    print("\n DL_FIELD .sf parser: renaming exsiting molecule key \'{}\' -> \'{}\'...\n".format(residue['key'], residue['key']+'_DLF'))
                    abuff[i]['key'] += '_DLF'
            # append residue names to possibly exsiting array
            dbuff['residues'] = append(dbuff['residues'], abuff)
    
            # chop into molecules
            masks, linenums_mol = fileutils.getChoppedMasks(filename, 'MOLECULE ', return_start_lines=True)
    
            # loop over moltypes
            moltypes = vstack(char.split(lines[linenums_mol], maxsplit=2).tolist())[:,1]
            for imol, moltype in enumerate(moltypes):
    
    #            # skip repeat in databank
    #            if moltype not in residues['type']:
    #                print("### not in types:", moltype)
    #                continue
    
                txt = lines[masks[imol]]
                dbuff.update({moltype:dict()})
                natoms = int(txt[0].split()[2])
                mol_data = dbuff[moltype]
    
                # atoms
                mol_data.update({'atom':genfromtxt(txt[1:natoms+1],
                                                   dtype=dtype([('key','U12'),('type','U36'),('charge','f8')]),
                                                   comments='#',
                                                   usecols=range(3))
                                })
    
                # bonds
                masks_bond = char.startswith(txt, 'CONNECT ')
                nbonds = masks_bond.sum()
                if nbonds:
                    dbuff[moltype].update({'bond':vstack(char.split(txt[masks_bond], maxsplit=1).tolist())[:,1]})
                else:
                    dbuff[moltype].update({'bond':array([])})
    
                # impropers
                masks_improper = char.startswith(txt, 'IMPROPER ')
                nimpropers = masks_improper.sum()
                if nimpropers:
                    abuff = genfromtxt(txt[masks_improper],
                                       dtype=str,
                                       comments='#',
                                       usecols=[1,2,3,4])
    #                dbuff[moltype].update({'impr':vstack(char.split(txt[masks_improper], maxsplit=1).tolist())[:,1]})
                    dbuff[moltype].update({'impr':abuff})
                else:
                    dbuff[moltype].update({'impr':array([])})
    
                # rigid
                masks_rigid = char.startswith(txt, 'RIGID ')
                nrigids = masks_rigid.sum()
                if nrigids:
                    dbuff[moltype].update({'rigid':vstack(char.split(txt[masks_rigid], maxsplit=1).tolist())[:,1]})
    #            else:
    #                dbuff[moltype].update({'rigid':array([])})
    
            receiver.update({potential:dbuff})
