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

# YL 26/10/2021: raised to module level for convenient reuse
# dtypes of sections
# YL 05/12/2023: here "names" are actually FF types
contents2dtypes = { 'atom'    : [('name','U8'),('mass','f8'),('charge','f8'),('nrepeats','i4'),('frozen','i4')],
                    'bond'    : [('type','U8'),('I','i4'),('J','i4'),('Kb','f8'),('b0','f8')],
                    'angle'   : [('type','U8'),('I','i4'),('J','i4'),('K','i4'),('Ktheta','f8'),('Theta0','f8')],
                    'dihedral': [('type','U8'),('I','i4'),('J','i4'),('K','i4'),('L','i4'),('Kchi','f8'),('delta','f8'),('n','i4'),('14e','f8'),('14vdw','f8')],
                    'vdw'     : [('A','U8'),('B','U8'),('type','U8'),('a','f8'),('b','f8')]
                  }

from sys import stdout

def parse(filename, prefix='', **kwargs):
    '''A parser of DL_POLY's FIELD file'''

    from numpy import arange, array, char, genfromtxt, searchsorted, set_printoptions, zeros
    from os    import path

    dbuff = {}

    delimiters = { 'bond'    : [5,5,5,12,12],
                   'angle'   : [5,5,5,5,12,12],
                   'dihedral': [5,5,5,5,5,12,12,12,12,12],
                 }

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

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

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

        # all line numbers of lines starting with "nummols" and "finish"
        linenums_molecule = linenums[char.startswith(lines_lower, 'nummols')]
        linenums_finish   = linenums[char.startswith(lines_lower, 'finish')]

        # unit
        linenum_unit = linenums[char.startswith(lines_lower, 'units')][0]
        dbuff.update({'unit':lines_lower[linenum_unit].split()[1]})

        # DL_POLY <=4: molecular types
        try:
            linenum_unit = linenums[char.startswith(lines_lower, 'molecular types')][0]
            dbuff.update({'ntypes':int(lines_lower[linenum_unit].split()[2])})
        # YL 06/12/2021: DL_POLY >=5: molecules
        except IndexError:
            linenum_unit = linenums[char.startswith(lines_lower, 'molecules')][0]
            dbuff.update({'ntypes':int(lines_lower[linenum_unit].split()[1])})

        if linenums_molecule.size != linenums_finish.size:
            print("\n >>> ERROR: Incomplete DL_POLY FIELD file {}\n".format(filename))

        if linenums_molecule.size != dbuff['ntypes']:
            dbuff['ntypes'] = linenums_molecule.size

        # loop over molecules
        for i, molecule in enumerate(linenums_molecule):

            # add a molecule entry
            dbuff.update({i:{}})

            l0, l1 = linenums_molecule[i], linenums_finish[i]

            # title if exists
            if lines_lower[l0-1].startswith('molecule') and not lines_lower[l0-1].startswith('molecules'):
                dbuff[i].update({'title':' '.join(lines[l0-1].split()[1:])})
            else:
                dbuff[i].update({'title':'molecule\'s title'})

            lines_molecule = lines_lower[l0:l1]
            lines_molecule_original = lines[l0:l1]
            sublinenums = arange(lines_molecule.size)

            # number of molecules
            dbuff[i].update({'nmolecules':int(lines[l0].split()[1])})

            # atoms
            dbuff[i].update({'atom':{}})
            sublinenum_atom = sublinenums[char.startswith(lines_molecule, 'atom')][0]
            natoms = int(lines_molecule[sublinenum_atom].split()[1])
            dbuff[i]['atom'].update({'nentries':natoms})
            sublinenums_atom = sublinenums[sublinenum_atom+1:sublinenum_atom+natoms+1]
            abuff = genfromtxt(lines_molecule_original[sublinenums_atom],
                               dtype=contents2dtypes['atom'],
                               comments='#',
                               invalid_raise=True)
            dbuff[i]['atom'].update({'data':abuff})

            # other entries
            for entry, dtype in contents2dtypes.items():
                if entry in [ 'atom', 'vdw' ]:
                    continue
                dbuff[i].update({entry:{}})
                try:
                    sublinenum_entry = sublinenums[char.startswith(lines_molecule, entry)][0]
                # no such entry
                except IndexError:
                    continue
                nentries = int(lines_molecule[sublinenum_entry].split()[1])
                dbuff[i][entry].update({'nentries':nentries})
                sublinenums_entry = sublinenums[sublinenum_entry+1:sublinenum_entry+nentries+1]
                set_printoptions(suppress=True)
                # free-format
                try:
                    abuff = genfromtxt(lines_molecule[sublinenums_entry],
                                       dtype=contents2dtypes[entry],
                                       filling_values=0.0,
                                       comments='#',
                                       invalid_raise=True)
                # fixed-format generated by Tcl-ChemShell
                except ValueError:
                    abuff = genfromtxt(lines_molecule[sublinenums_entry],
                                       dtype=contents2dtypes[entry],
                                       filling_values=0.0,
                                       comments='#',
                                       delimiter=delimiters[entry],
                                       invalid_raise=True)
                # YL 06/11/2021: trim characters
                for name in abuff.dtype.names:
                    if abuff[name].dtype.name.startswith('str'):
                        abuff[name] = char.strip(abuff[name])
                dbuff[i][entry].update({'data':abuff})

#                               converters={ 9 : lambda x:float(x or 0.0)},
            # vdw
            dbuff.update({'vdw':{}})
            linenum_vdw = linenums[char.startswith(lines_lower, 'vdw')][0]
            nvdws = int(lines[linenum_vdw].split()[1])
            dbuff['vdw'].update({'nentries':nvdws})
            linenums_vdw = linenums[linenum_vdw+1:linenum_vdw+nvdws+1]
            # YL 24/08/2023: there could be a line of "close" that needs to be excluded
            masks_close = char.startswith(char.lower(char.strip(lines[linenums_vdw])), 'close')
            linenums_vdw = linenums_vdw[~masks_close]
            abuff = genfromtxt(lines[linenums_vdw],
                               dtype=contents2dtypes['vdw'],
                               comments='#',
                               filling_values=0.0,
                               invalid_raise=False)
            # YL 06/11/2021: trim characters
            for name in abuff.dtype.names:
                if abuff[name].dtype.name.startswith('str'):
                    abuff[name] = char.strip(abuff[name])
            dbuff['vdw'].update({'data':abuff})

    return dbuff


def write(filename, field_data, **kwargs):
    '''A writer of field_data (<class dict>)'''

    from io    import BytesIO
    from numpy import savetxt
    from time  import asctime, localtime, time

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

        # header
        fp.write('Generated by ChemShell, {}\n'.format(asctime(localtime(time()))))

        # unit
        fp.write('units {}\n'.format(field_data['unit']))

        # number of molecular types
        fp.write('molecules {}\n'.format(field_data['ntypes']))

        # loop over molecule types
        for key, val in field_data.items():

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

            # title
            if val.get('title', ''):
                # YL 13/07/2022: the title string's length cannot exceed 40 (39+1) in DL_POLY
                fp.write('{}\n'.format(val['title'][:39]))

            # number of molecules
            fp.write('nummols {}\n'.format(val['nmolecules']))

            # atoms
            fp.write('atoms {}\n'.format(val['atom']['nentries']))
            # convert to bytes: BytesIO, not StringIO!
            bbuff = BytesIO()
            savetxt(bbuff, val['atom']['data'], fmt='%-8s %12.5f %12.5f %4d %4d')
            # convert bytes to string
            fp.write(bbuff.getvalue().decode().rstrip() + '\n')

            # bonds
            if val.get('bond', {}).get('nentries', 0):
                nqqs  = val['bond'].get('nqqs', 0)
                nqjms = val['bond'].get('nqjms', 0)
                nqqs_existing = val['bond'].get('nqqs_existing', 0)
                if nqqs or nqjms:
                    if nqqs_existing:
                        fp.write(f'# {nqqs} QM-QM (incl. {nqqs_existing} existing QM-QM interactions) and {nqjms}\n')
                        fp.write(f'# QM-(JQ)-MM terms set to zero to prevent calculating their\n')
                        fp.write('# non-bond interactions\n')
                    else:
                        fp.write(f'# {nqqs} QM-QM and {nqjms} QM-(JQ)-MM terms set to zero to\n')
                        fp.write('# prevent calculating their non-bond interactions\n')
                fp.write('bonds {}\n'.format(val['bond']['nentries']))
                # convert to bytes: BytesIO, not StringIO!
                bbuff = BytesIO()
                savetxt(bbuff, val['bond']['data'], fmt='%-8s %8d %8d  %12.7e %12.5f')
                # convert bytes to string
                fp.write(bbuff.getvalue().decode().rstrip() + '\n')

            # angles
            if val.get('angle', {}).get('nentries', 0):
                ndeletions = val['angle'].get('ndeletions', 0)
                if ndeletions:
                    fp.write('# {} Q-Q-Q and Q-Q-M terms removed\n'.format(ndeletions))
                fp.write('angles {}\n'.format(val['angle']['nentries']))
                # convert to bytes: BytesIO, not StringIO!
                bbuff = BytesIO()
                savetxt(bbuff, val['angle']['data'], fmt='%-8s %8d %8d %8d %12.5f %12.5f')
                # convert bytes to string
                fp.write(bbuff.getvalue().decode().rstrip() + '\n')

            # dihedrals
            if val.get('dihedral', {}).get('nentries', 0):
                ndeletions = val['dihedral'].get('ndeletions', 0)
                if ndeletions:
                    fp.write('# {} Q-Q-Q-Q, Q-Q-Q-M, and M-Q-Q-M terms removed\n'.format(ndeletions))
                fp.write('dihedrals {}\n'.format(val['dihedral']['nentries']))
                # convert to bytes: BytesIO, not StringIO!
                bbuff = BytesIO()
                savetxt(bbuff, val['dihedral']['data'], fmt='%-8s %8d %8d %8d %8d %12.5f %12.5f %7.3f %7.5f %7.3f')
                # convert bytes to string
                fp.write(bbuff.getvalue().decode().rstrip() + '\n')

            fp.write('finish\n')

        # vdw
        if field_data.get('vdw', {}).get('nentries', 0):
            fp.write('vdw {}\n'.format(field_data['vdw']['nentries']))
            # convert to bytes: BytesIO, not StringIO!
            bbuff = BytesIO()
            savetxt(bbuff, field_data['vdw']['data'], fmt='%-8s %-8s %-8s  %14.7e  %14.7e')
            # convert bytes to string
            fp.write(bbuff.getvalue().decode().rstrip() + '\n')

        # close
        fp.write('close\n')






