#  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 ...io   import format
from ...base import forcefield

# YL TODO: we currently don't read in or write out any comment

# YL 27/10/2021: raised to module level for convenient reuse
# dtypes of sections
contents2dtypes = { 'bonds'    : [('I','U8'),('J','U8'),('Kb','f8'),('b0','f8')],
                   'angles'   : [('I','U8'),('J','U8'),('K','U8'),('Ktheta','f8'),('Theta0','f8'),('Kub','f8'),('S0','f8')],
                   'dihedrals': [('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'),('dummy','i8'),('psi0','f8')],
                   'cmap'     : None,
                   'nonbonded': [('I','U8'),('dummy','f8'),('epsilon','f8'),('Rmin/2','f8'),('dummy2','f8'),('eps,1-4','f8'),('Rmin/2,1-4','f8')],
                   'nbfix'    : [('I','U8'),('J','U8'),('Emin','f8'),('Rmin','f8')],
                   'hbond'    : [('I','U8'),('J','U8'),('hbond1','f8'),('hbond2','f8')],
                   'end'      : None,
                 }

# CHARMM parameter file
def parse(filename, _regex=None, receiver=None, override_existing=True, prefix='', **kwargs):
    '''Parse a CHARMM parameter file'''

    from numpy import append, arange, argsort, array, char, dtype, delete, genfromtxt, roll, __version__
    from copy  import copy
    from os    import path

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

    print('\n Parsing CHARMM parameter file:\n {}'.format(filename))

    tmp_filename = '._chemsh.prm'

    # mainly for .inp parameters
    synonyms = { 'bonds'    :'bond',
                 'angles'   :'thetas',
                 'dihedrals':'phi',
                 'improper' :'imphi',
               }

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

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

        # line numbers of those start with section names
        # NB: since Python 3.6 dict is ordered
        alphalines = []
        keys = []
        for key in contents2dtypes.keys():
            # skip not-found sections
            try:
                # YL 28/10/2021: try to stretch to the last 'end' because in .str files
                #                there may be multiple sections with a same title
                if key == 'end':
                    alphalines.append(linenums[char.startswith(lines_lower, key)][-1])
                    keys.append(key)
                else:
                    alphalines.append(linenums[char.startswith(lines_lower, key)][0])
                    keys.append(key)
            except:
                # try synonyms
                try:
                    synonym = synonyms[key]
                    alphalines.append(linenums[char.startswith(char.lower(lines), synonym)][0])
                    keys.append(key)
                except:
                    continue

        if 'end' not in keys:
            print("\n >>> ChemShell WARNING: CHARMM PRM parser does not find an \"END\"\n")

        sorting = argsort(alphalines)

        # YL 21/02/2021: shouldn't chop off the last index (for "END")
        alphalines = array(alphalines)[sorting]#[:-1]
        omegalines = roll(alphalines,-1)[:-1]

        # YL 28/09/2021: a block must not begin with 'END'
        if lines_lower[alphalines[0]].startswith('end'):
            alphalines = alphalines[1:]
            omegalines = omegalines[1:]

        sections = dict(zip(zip(alphalines,omegalines), array(keys)[:-1]))

        # loop over sections
        for (ln0, ln1), section in sections.items():

            # section title was included above
            skip_header = 1
            if lines[ln0].strip().endswith('-'):
                skip_header += 1
            # datatype
            dt = dtype(contents2dtypes.get(section, str))
            # skip dummy entries (not in use by CHARMM/NAMD and not supported by DL_POLY/DL_FIELD)
            fields = []
            # number of columns to use
            cols = []
            try:
                for i, name in enumerate(dt.names):
                    if not name.startswith('dummy'):
                        fields.append((name,dt[name]))
                        cols.append(i)
                dt = dtype(fields)
            except TypeError:
                cols = range(0)
                dt = dtype(contents2dtypes.get(section, str))

            fields_dt = copy(dt.names)
            # comment lines and empty lines
            masks = char.startswith(char.strip(lines[ln0:ln1]), '!') + char.isspace(lines[ln0:ln1])

            # remove the trailing comment of each line
            partitioned = char.partition(lines[ln0:ln1][~masks], '!')[:,0]
            # split by white spaces
            # YL 28/07/2023: split is an array of lists (dtype=object), each containing strings
            split = char.split(partitioned)
            # insert ',' as delimiter as genfromtxt() doesn't accept missing data with space as delimiter
            # reason: in genfromtxt() all converters/missing_values/filling_values mechanisms take place only after line splitting
            # see: site-packages/numpy/lib/npyio.py
            numpy_version = __version__.split('.')
            numpy_major_version, numpy_minor_version = int(numpy_version[0]), int(numpy_version[1])
            # YL 28/07/2023: since NumPy 1.25 we can no longer use `char.join(',', split)`
            if numpy_major_version == 1 and numpy_minor_version >= 25:
                joined = array([ ','.join(foo) for foo in split ])
            else:
                joined = char.join(',', split)
            nrows = joined.shape[0]
            # we have to attach zeros to complement the missing columns
            trailing = char.multiply(array([',0.0']*nrows).T, 3)
            joined = char.add(joined, trailing)
            # write out to a temporary .prm file
#            tmpfp.writelines('\n'.join(joined))

            # YL 27/09/2021: remove lines starting with special commands and atom definitions,
            #                especially for .str (stream) files
            abuff = char.lower(joined)
            masks = char.startswith(abuff, '*')
            masks += char.startswith(abuff, 'if')
            masks += char.startswith(abuff, 'set')
            masks += char.startswith(abuff, 'read')
            masks += char.startswith(abuff, 'atoms')
            masks += char.startswith(abuff, 'mass')
            # YL 28/10/2021: 'wrnlev', 'bomlev', and 'end' may appear within these data lines
            #                because .str may contain multiple sections with a same title
            masks += char.startswith(abuff, 'wrnlev')
            masks += char.startswith(abuff, 'bomlev')
            masks += char.startswith(abuff, 'end')
            # YL 28/10/2021: the section name should not appear between these data lines except
            #                for the first line; we do this because .str files may contain multiple
            #                sections with a same title
            masks_section = char.startswith(abuff, section)
            masks_section[0] = False
            masks += masks_section
            joined = joined[~masks]

            # YL 19/02/2021: skip empty section
            if not joined.shape[0] - skip_header:
                continue

            # read in data
            try:
                abuff = genfromtxt(joined,
                                   skip_header=skip_header,
                                   dtype=dt,
                                   usecols=cols,
                                   delimiter=',',
                                   comments='!',
                                   invalid_raise=True)
                # YL 01/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
            
            except ValueError:
                abuff = joined

            # initialise array of all residue names
            if section not in receiver:
                receiver.update({section:array([], dtype=dt)})

            # YL 16/02/2021: override existing entries by deleting
            if override_existing and receiver[section].dtype.fields:
                receiver[section] = forcefield.removeRepeats(receiver[section], abuff)

            if isinstance(receiver, dict):
                try:
                    receiver.update({section:append(receiver[section], abuff)})
                except:
                    print('\n WARNING: Section \"{}\" not parsed from CHARMM parameter file \"{}\"\n'.format(filename))

    return 0


def write(filename, dict_data, append=False, fill_with_zeros=False):
    '''Writes out a .prm file in the CHARMM format using the provided dictionary of forcefield data'''

    if append:
        fp = open(filename, 'a')
    else:
        fp = open(filename, 'w')

    fp.write('\n\n')
    for key, val in dict_data.items():
        fp.write(f'{key.upper()}\n')
        dtype = val.dtype
        try:
            lbuff = [':<8' if dtype[i].name.startswith('str')
                           else ':>14.6f' if dtype[i].name.startswith('float')
                           else ':>3d' if dtype[i].name.startswith('int')
                           else ':' for i in range(len(dtype.names))]
            fmt = '{' + '} {'.join(lbuff) + '}\n'
            # YL 28/09/2021: we need the dummy columns for Tcl-ChemShell
            if key == 'improper':
                ltmp = fmt.split()
                ltmp.insert(5,'0')
                fmt = ' '.join(ltmp) + '\n'
            elif key == 'nonbonded':
                ltmp = fmt.split()
                ltmp.insert(1,'0.0')
                ltmp.insert(4,'0.0')
                fmt = ' '.join(ltmp) + '\n'
            # YL 22/10/2021: keep a copy for special treatment of angles (see below)
            fmt_copy = fmt
            
            for i in val:
                # YL 22/10/2021: angles must be handled specially because if there are data, even 0.0,
                #                for the Kub and S0 columns, Tcl-ChemShell considers this angle Urey-Bradley;
                #                on the other hand, DL_FIELD needs dummy data (0.0) for the Kub and S0 columns,
                #                therefore we provide a keyword option fill_with_zeros to control different versions
                if not fill_with_zeros and key == 'angles':
                    if i[-1] == i[-2] == 0.0:
                        fmt = ' '.join(fmt_copy.split()[:-2]) + '\n'
                    else:
                        fmt = fmt_copy
                # YL 22/10/2021: similar is nonbonded
                if not fill_with_zeros and key == 'nonbonded':
                    if i[-1] == i[-2] == 0.0:
                        # YL 22/10/2021: remove the above added dummy 0.0 column
                        fmt = ' '.join(fmt_copy.split()[:-3]) + '\n'
                    else:
                        fmt = fmt_copy
                # write to file
                fp.write(fmt.format(*i))
        except:
            pass
        fp.write('\n')

    fp.write('END\n')

    fp.close()
    


















