#  Copyright (C) 2019 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 sys import stdout
import chemsh
from ...dl_py2f import utils

contents2dtypes = { 'atom'     : [('key','U12'),('type','U36'),('charge','f8')],
                    'bond'     :  None,
                    'angle'    :  None,
                    'dihe'     :  None,
                    'double'   :  None,
                    'triple'   :  None,
                    'impr'     :  None,
                    'cmap'     : [('str','U256')],
                    'donor'    : [('I','U4'),('J','U4')],
                    'acceptor' : [('I','U256')],
                    'ic'       : [('I','U4'),('J','U4'),('K','U4'),('L','U4'),('1-2','f8'),('1-2-3','f8'),('1-2-3-4','f8'),('2-3-4','f8'),('3-4','f8')],
                    'delete'   : [('entry','U8'),('buffer','U2048')],
                    'end'      :  None,
                  }

# CHARMM topology file
def parse(filename, _regex=None, receiver=None, residues_only=False, prefix='', parse_file=False, **kwargs):
    '''Parse a CHARMM topology file'''

    from numpy  import all, append, arange, argwhere, array, bincount, char, copy, delete, dtype, genfromtxt, hstack, isnan, roll, searchsorted, unique, vstack, zeros
    from os     import path, stat
    from pickle import dump, load

    top_entries = {
                    'mass' : [('index','i8'),('type','U4'),('mass','f8')],
                    'decl' : [('item','U4')],
                    'defa' : [('first','U4'),('fter','U4'),('last','U4'),('lter','U4')],
                    'auto' : [('str','U256')],
                  }

    # TODO: not in use
    tmp_receiver = {}

    tmp_filename = '._chemsh.rtf'

    # initialise array of all residue names
    if 'residues' not in receiver:
        receiver.update({'residues':array([], dtype=dtype([('type','U36'),('key','U36'),('mass','f8')]))})

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

    # YL 21/04/2024: parsing some large RTF files may take a long time (such as
    #                top_all36_cgenff.rtf), and thus we write out a picked data
    #                file once the RTF is parsed and next time we don't have to
    #                spend time on parsing it again
    name_rtf = utils.fileutils.getBaseName(path.basename(filename))
    filename_pkl = path.join(path.dirname(filename), f'{name_rtf}.pkl')
    timestamp = stat(filename).st_mtime
    to_parse = True
    if path.isfile(filename_pkl):
        with open(filename_pkl, 'rb') as fp:
            dict_rtf = load(fp)
            # YL 21/04/2024: we check the timestamp because the RTF could be modified/updated
            ts = dict_rtf.get('timestamps', {}).get(filename, 0.0)
            if timestamp == ts:
                to_parse = False
            else:
                to_parse = True
            utils.dictutils.mergeDicts(receiver, dict_rtf)
            return 0

    if to_parse and not parse_file:
        # YL 01/05/2024: we need this recursive run because receiver may be a collection of
        #                multiple RTFs but a PKL must be of single RTF
        dict_rtf = dict()
        parse(filename, fmt='rtf', receiver=dict_rtf, parse_file=True)
        dict_rtf.update({'timestamps':{filename:timestamp}})
        with open(filename_pkl, 'wb') as fp:
            dump(dict_rtf, fp)
        utils.dictutils.mergeDicts(receiver, dict_rtf)
        return 0

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

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

        # line numbers of those start with "RESI", "PRES", and "ATOM"
        pres  = linenums[char.startswith(lines_lower, 'pres')]
        resi  = linenums[char.startswith(lines_lower, 'resi')]
        linenums_resi = linenums[char.startswith(lines_lower, 'resi')+char.startswith(lines_lower, 'pres')]

        ### residues ###
        # a line may start with "ATOMS" (should be excluded)
        atom_masks = char.startswith(lines_lower, 'atom')*~char.startswith(lines_lower, 'atoms')
        # not in use
        bond_masks = char.startswith(lines_lower, 'bond')
        # not in use
        mass_masks = char.startswith(lines_lower, 'mass')

        masks = zeros(shape=(nlines,), dtype=bool)
        for entry in [ 'atom', 'bond', 'double', 'triple', 'angle', 'dihe', 'impr', 'cmap', 'donor', 'acceptor', 'ic', 'delete', 'group' ]:
            masks += char.startswith(lines_lower, entry)

#        _masks = char.startswith(lines_lower, 'bond')
        atom = linenums[atom_masks]
        linenums_entries = linenums[masks]
        nresidues = linenums_resi.size

        # each residue's starting and ending lines of "ATOM"
        # "PRES" can appear after the last occurance of "ATOM"
        sorting = searchsorted(atom, resi)
        sorting_resi = searchsorted(linenums_entries, linenums_resi)

        # IndexError won't happen unless the RTF file is wrong
        lindices = atom[sorting]
        try:
            rindices = atom[append(sorting[1:]-1, atom.size-1)]
        except IndexError:
            if chemsh.CHEMSH_DEBUG > 2:
                print(' ChemShell RTF Parser: file {} is not a standard RTF! Returning...'.format(filename))
            return 1
        lindices_resi = linenums_resi
        rindices_resi = linenums_entries[append(sorting_resi[1:]-1, linenums_entries.size-1)]

        # find leading "PRES" lines after "RESI" ones (a residue may contain multiple patch entries)
        if pres.size:
            sorting = searchsorted(pres, resi)
            try:
                rpres = pres[unique(sorting)]
            except IndexError:
                # if there are more RESI entries beyond the last PRES
                if unique(sorting)[-1] == pres.size:
                    rpres = pres[unique(sorting)[:-1]]
                # if the above is not the case there must be something wrong!
                else:
                    raise
            # replace rindeces those are preceded by a "PRES"
            try:
                rindices[searchsorted(rindices,rpres)] = rpres
            # extreme case that there is only one RESI
            except IndexError:
                rindices = append(rindices, rpres)

        # construct a masks array using the starting and ending indices
        # shape of masks: (nresidues, nlines)
        masks_old = (lindices[:,None]<=linenums)*(rindices[:,None]>=linenums)
        masks = (lindices_resi[:,None]<=linenums)*(rindices_resi[:,None]>=linenums)

        # resnames
        try:
            resnames = vstack(char.split(lines[resi], maxsplit=2).tolist())[:,1]
        # YL 20/06/2023: it's possible this RTF file doesn't contain any residues (but only patches instead, for example)
        except ValueError:
            resnames = array([], dtype='U')
        for iresi, resname in enumerate(resnames):
            alines = lines[masks_old[iresi]*atom_masks]
            try:
                atomnames = vstack(char.split(alines, maxsplit=2).tolist())[:,1]
            # it is possible a residue entry has no atom defined
            except ValueError:
                atomnames = array([], dtype=alines.dtype)
            tmp_receiver.update({resname:atomnames})

        # loop over top-level entries
        for entry, val in top_entries.items():
            dt = dtype(val)
            txt = lines[char.startswith(lines_lower,entry)]
            if len(txt):
                abuff = genfromtxt(txt,
                                   dtype=dt,
                                   usecols=tuple(range(1,len(dt.names)+1)),
                                   comments='!',
                                   invalid_raise=True)
            # no record
            else:
                abuff = array([], dtype=dt)
            receiver.update({entry:abuff})

        # YL 19/02/2021: we attach two dummy columns first and then split to make sure the shape
        #                is consistent to be used by vstack() because there might be missing
        #                columns in some RTF files
        abuff = char.add(char.add(lines[linenums_resi], '0.00'), '! dummy')
        try:
            resnames = vstack(char.split(abuff, maxsplit=2).tolist())[:,1]
        # YL 20/06/2023: it's possible this RTF file doesn't contain any residues (but only patches instead, for example)
        except ValueError:
            resnames = array([], dtype='<U1')

        # by now we have got all of the residues names
        # make a dummy dict containing the residue names if only the latters are wanted
        if residues_only:
            receiver.update(dict(zip(resnames,[None]*resnames.size)))
            return -1

        # loop over residues (including protonated ones)
        for iresi, resname in enumerate(resnames):
            # create an item of the current residue in result dict
            receiver.update({resname:{}})
            # complement the residue name entries to match DL_POLY/DL_FIELD format
            receiver['residues'] = append(receiver['residues'],
                                          array([(resname,resname,0.0)],
                                                dtype=receiver['residues'].dtype))
            for section, val in contents2dtypes.items():
                dt = dtype(val)
                # number of columns to use
                try:
                    ncols = len(dt.names)
                except:
                    ncols = 0
                skip_header = 0
                # relevant lines only
                txt = lines[masks[iresi]*char.startswith(lines_lower,section)]
                # skip empty entry
                if not len(txt):
                    continue
                try:
                    abuff = genfromtxt(txt,
                                       skip_header=skip_header,
                                       dtype=dt,
                                       usecols=tuple(range(1,ncols+1)),
                                       comments='!',
                                       invalid_raise=True)

                    # YL 28/09/2023: treat deletions as a special case
                    if section == 'delete':
                        # remove comments
                        txt = char.strip(char.partition(array(txt), '!')[:,0])
                        txt = char.strip(char.partition(txt, section.upper())[:,2])
                        abuff = array([], dtype=contents2dtypes['delete'])
                        for i_del, to_del in enumerate(['ATOM','BOND','DOUBLE','TRIPLE','ANGLE','DIHE','IMPR']):
                            indices_startswith = char.startswith(txt, to_del)
                            buffer = char.partition(txt[indices_startswith], to_del)
                            if buffer.shape != (0,):
                                buffer = ' '.join(char.strip(buffer[:,2]))
                                abuff = hstack((abuff, array((to_del, buffer), dtype=contents2dtypes['delete'])))
                                
                    # treat as string if got [nan nan nan...]
                    try:
                        if all(isnan(abuff)):
                            raise ValueError
                    except TypeError:
                        pass
                # read in the whole string if failed (for uncertain length of rows such as BOND)
                except ValueError:
                    abuff = char.partition(txt, section.upper()).shape
                    # remove comments
                    txt = char.strip(char.partition(array(txt), '!')[:,0])
                    # remove leading keys and join
                    txt = ' '.join(char.partition(txt, section.upper())[:,2])
                    # split again
                    abuff = array(txt.split())

                # update the result dict
                receiver[resname].update({section:abuff})

#                # special treatment because CHARMM's atom names are useless to DL_FIELD while types are missing
#                if section == 'atom':
#                    receiver[resname][section]['type'] = char.add(receiver[resname][section]['type'], '_CHARMM')

        ### patch residues ###
        pres_dict = dict()

        # each patch's starting and ending lines of "atom"
        # NB: a PRES may not contain ATOM entry and has to be removed
        # YL: this is a special pure-NumPy algorithm for finding the indices where the occurence is greater than 1
        #     1) indices of inserting `pres` into `atom`
        #     2) indices to reconstruct `sorting`
        #     3) find out where elements are repeated
        sorting = searchsorted(atom, pres)
        _, lb = unique(sorting, return_inverse=True)
        nodiff = argwhere((lb-roll(lb,-1))==0).ravel()
        todel = unique(hstack((nodiff,nodiff+1)))

        try:
            pres = delete(pres, arange(pres.size)[todel])
        # again, extreme case that there's only one PRES
        except IndexError:
            pres = delete(pres, arange(pres.size)[0])

        if not pres.size:
            return 0

        # IndexError won't happen unless the RTF file is wrong
        sorting = searchsorted(atom, pres)
        lindices = atom[sorting]
        rindices = atom[append(sorting[1:]-1, atom.size-1)]

        # construct a masks array using the starting and ending indices
        # shape of masks: (nresidues, nlines)
        masks = (lindices[:,None]<=linenums)*(rindices[:,None]>=linenums)*atom_masks

        # pres names
        presnames = vstack(char.split(lines[pres], maxsplit=2).tolist())[:,1]

        for ipres, presname in enumerate(presnames):
            alines = lines[masks[ipres]]
            atomnames = vstack(char.split(alines, maxsplit=2).tolist())[:,1]
            tmp_receiver.update({presname:atomnames})


    return 0


# _topol_abspaths must be a list of absolute path(s)
def getResidueDicts(_topol_abspaths, residues_only=False):
    '''Return a dict of dicts containing mapping from topology files to parsed residue-atoms data'''

    _rtf_res_dicts = dict()
    for itop, topol_abspath in enumerate(_topol_abspaths):
        # save RTF res/atom names in a dict
        rtf_res_dict = dict()
        parse(topol_abspath, receiver=rtf_res_dict, residues_only=residues_only)
        _rtf_res_dicts.update({topol_abspath:rtf_res_dict})

    return _rtf_res_dicts


# YL 20/09/2023
def write(filename, dict_data):
    '''Writer of RTF data'''

    from numpy import array, char

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

        # loop over residues
        for resname, dict_residue_data in dict_data.items():
            fp.write(f'! {resname}\n')
            data = dict_residue_data.get('atom', array([], dtype=contents2dtypes['atom']))
            fp.write(f'RESI {resname:8s} {data["charge"].sum():>.6f}\n')
            # write atoms
            for atom in data:
                fp.write(f'ATOM {atom["key"]:6s} {atom["type"]:6s} {atom["charge"]:>.6f}\n')
            # write bonds
            for entry in [ 'bond', 'double', 'triple' ]:
                data = dict_residue_data.get(entry, array([], dtype=contents2dtypes[entry]))
                lbuff = list(zip(data[::2],data[1:][::2]))
                for bond in lbuff:
                    if not any(char.find(bond,'+')+1) and not any(char.find(bond,'-')+1):
                        fp.write(f'BOND {bond[0]:6s} {bond[1]:6s}\n')
            # write angles
            data = dict_residue_data.get('angle', array([], dtype=contents2dtypes['angle']))
            lbuff = list(zip(data[::3],data[1:][::3],data[2:][::3]))
            for angle in lbuff:
                fp.write(f'ANGLE {angle[0]:6s} {angle[1]:6s} {angle[2]:6s}\n')
            # write dehedrals/impropers
            for entry in [ 'dihe', 'impr' ]:
                data = dict_residue_data.get(entry, array([], dtype=contents2dtypes[entry]))
                lbuff = list(zip(data[::4],data[1:][::4],data[2:][::4],data[3:][::4]))
                for dihedral in lbuff:
                    fp.write(f'{entry.upper()} {dihedral[0]:6s} {dihedral[1]:6s} {dihedral[2]:6s} {dihedral[3]:6s}\n')
            fp.write('\n\n')
            fp.flush()
            stdout.flush()


