#  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 sys import stdout

def parse(filename, _regex=None, receiver=None, prefix='', **kwargs):
    ''''''

    from numpy import array, asarray, char, vstack
    from os    import path

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

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

        lines = array(fp.readlines())

        plines = lines[char.startswith(char.lower(lines), 'potential')]

        potentials = vstack(char.split(char.strip(plines)))[:,1]

        for potential in potentials:
            if potential not in receiver:
                receiver.update({potential:{}})


def write(filename, dict_data, **kwargs):
    '''Write out a UDFF file'''

    from io      import BytesIO
    from numpy   import all, argwhere, array, char, dtype, empty, hstack, ndarray, savetxt, unique
    from ...base import forcefield
    from ...data import atomicdata, forcefield as ff
    from ...io   import files

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

        fp.write('# UDFF for DL_FIELD generated by ChemShell\n\n')

        # YL 15/01/2021: loop over potentials (schemes)
        for scheme, dbuff in dict_data.items():

            # pre-defined conventions of names parsing
            scheme_prefix = scheme.partition('_')[0].strip('0123456789').lower()
            try:
                # YL FIXME 16/02/2021: we need the original forcefield to understand some rules, for example names2znums exceptions
                ffObj = ff.getFFObj(scheme_prefix)
                names2znums_exceptions = ffObj.names2znums
            except:
                names2znums_exceptions = {}

            strbuff = 'POTENTIAL {}\n\n'.format(scheme)
    
            # unit
            strbuff += 'UNIT kcal/mol\n\n'
    
            # atom types
            strbuff += '{:22} {:10} {:8} mass   remark\n'.format('ATOM_TYPE', 'ATOM_KEY', 'element')
            atomtypes = []

            len_type    = 22
            len_key     = 10

            residues = dbuff.get('residues', [])
            for residue in residues:
                restype = residue['type']
                atomdata = dbuff.get(restype, {}).get('atom', array([]))
                if atomdata.size == 0:
                    continue
                # for single-atom residues
                elif atomdata.size == 1:
                    atomdata = atomdata[None]
                atomtypes_res = dbuff.get('atomtypes', array([],dtype=atomdata.dtype))['type'] 
                # CHARMM RTF doesn't contain DL_FIELD types
                fmt = '{:<'+'{}'.format(len_type)+'s} {:<'+'{}'.format(len_key)+'s} {:6} {:12.3f} !\n'
                # remove repeated atom types
                # NB: or should use a different key only?
                unique_types, unique_indices = unique(atomdata['key'], return_index=True)
                for atom in atomdata[unique_indices]:
                    # for single-atom residues
                    if atom.shape != ():
                        atom = atom[0]
                    atomtype = atom['type']
                    atomkey  = atom['key']
                    # data from DL_FIELD that contain 'element' and 'mass' fields
                    try:
                        atomelem = atom['element']
                        atommass = atom['mass']
                    except:
                        # for single-atom residues (for example sodium ion)
                        if array([atomtype]).shape == (1,1):
                            atomznum = forcefield.names2znums(array(atomtype), exceptions=names2znums_exceptions)[0]
                        else:
                            atomznum = forcefield.names2znums(array([atomtype]), exceptions=names2znums_exceptions)[0]
                        atomelem = atomicdata.symbols_strarray[atomznum]
                        try:
                            atommass = atomicdata.abundant_masses[atomznum]
                        except IndexError:
                            atommass = 0.0
    
                    # skip those already defined in the DL_FIELD built-in lib
                    if atomtype not in atomtypes_res and atomtype not in atomtypes:
                        # according to the DL_FIELD notations an atom key has to be "XXX_YYY"
                        # (typically from CHARMM PAR and RTF and CHARMM has no "key" for atom definitions)
                        if atomtype.find('_') == -1:
                            fmt = '{:<'+'{}'.format(len_type)+'s} {:<'+'{}'.format(len_key)+'s} {:2} {:12.3f} !\n'
                            strbuff += fmt.format(atomtype+'_'+scheme_prefix, atomtype, atomelem, atommass)
                        else:
                            strbuff += fmt.format(atomtype, atomtype, atomelem, atommass)
                        atomtypes.append(atomtype.strip())
            strbuff += 'END ATOM_TYPE\n\n'
    
    
            # molecule types
            restypes = dbuff.get('residues')['type']
            reskeys  = dbuff.get('residues')['key']
            maxlen_type = char.str_len(char.strip(restypes)).max()
            maxlen_key  = char.str_len(char.strip(reskeys)).max()
            fmt = '{:<'+'{}'.format(maxlen_type+4)+'s} {:<'+'{}'.format(maxlen_key+9)+'s} {:12} {}\n'
            strbuff += fmt.format('MOLECULE_TYPE', 'key', 'mw', 'remark')
            for residue in dbuff.get('residues'):
                restype = residue['type']
                reskey  = residue['key']
                resmass = residue['mass']
                data = dbuff.get(restype, {}).get('atom', [])
                if not isinstance(data, ndarray):
                    continue
                fmt = '{:<'+'{}'.format(maxlen_type+4)+'s} {:<'+'{}'.format(maxlen_key+4)+'s} {:12.3f} {}\n'
                if resmass:
                    strbuff += fmt.format(restype, reskey, resmass, '!')
                # calculate mass of weight if not available
                else:
                    # YL 19/02/2021: data could be 0D array (it's also checked in forcefield.names2znums() now)
                    if data.shape == ():
                        data = data[None]
                    strbuff += fmt.format(restype, reskey, atomicdata.abundant_masses[forcefield.names2znums(data['key'], exceptions=names2znums_exceptions)].sum(), '!')
            strbuff += 'END MOLECULE_TYPE\n\n'
    
            ### topologies ###
    
            # loop over residues
            for residue in dbuff.get('residues', []):
                restype = residue['type']
    
                atomdata = dbuff.get(restype, {}).get('atom', array([]))
                if not atomdata.size:
                    continue
    
                # begin of residue
                totalq = atomdata['charge'].sum()
                strbuff += 'MOLECULE {} {:6d} {:12.6f}\n'.format(restype, atomdata.size, totalq)
    
                # atom
                maxlen = char.str_len(atomdata['key']).max()
                bbuff = BytesIO()
                # could contain only 1 atom
                if atomdata.shape == ():
                    atomdata = atomdata[None]
                if all(char.find(atomdata['type'], '_')>-1):
                    fmt = '%-{0}s %-{0}s %12.6f'.format(maxlen+2)
                # DL_FIELD notation: "XXX_YYY"
                else:
                    fmt = '%-{0}s %{0}s_{1} %12.6f'.format(maxlen+2, scheme_prefix)
                savetxt(bbuff,
                        atomdata[['key','type','charge']],
                        fmt=fmt)
    
                strbuff += bbuff.getvalue().decode()
    
                # bond
                bonds = dbuff.get(restype, {}).get('bond', array([]))
                # CHARMM data
                # YL FIXME: 15/02/2021: there are bond entries in CHARMM with symbols like '+' and '*' preceding atom names
                try:
                    # connect bonds (including double and triple bonds)
                    doubles = dbuff.get(restype, {}).get('double', array([]))
                    triples = dbuff.get(restype, {}).get('triple', array([]))
                    bonds = hstack((bonds, doubles, triples))
                    bonds = bonds.reshape((int(bonds.size/2),2))
                    if bonds.size:
                        charmm_format = False
                        for a in atomdata:
                            connected = hstack((bonds[:,0][argwhere(bonds[:,1]==a[0]).ravel()],
                                                bonds[:,1][argwhere(bonds[:,0]==a[0]).ravel()]))
                            if connected.size:
                                strbuff += 'CONNECT {:>8} > {}\n'.format(a[0], ' '.join(connected))
                                charmm_format = True
                    # raw text from UDFF/DL_FIELD
                    if not charmm_format:
                        bonds = dbuff.get(restype, {}).get('bond', array([]))
                        for b in bonds:
                            strbuff += 'CONNECT {}\n'.format(b)
                # DL_POLY lib data
                except:
                    for b in bonds:
                        strbuff += 'CONNECT {}\n'.format(b)

                # YL 11/10/2023: added patching methods such as "ANGLE DELETE XX YY ZZ"
                patches = dbuff.get(restype, {}).get('patch', array([]))
                for p in patches:
                    strbuff += f'{p}\n'

                # improper
                impropers = dbuff.get(restype, {}).get('impr', array([]))
                # CHARMM data
                try:
                    impropers = impropers.reshape((int(impropers.size/4),4))
                    for i in impropers:
                        strbuff += 'IMPROPER {:8} {:8} {:8} {:8}\n'.format(i[0], i[1], i[2], i[3])
                # DL_POLY lib data
                except:
                    for i in impropers:
                        strbuff += 'IMPROPER {}\n'.format(i)
    
                # rigid
                # TODO: optionally switch off
                rigids = dbuff.get(restype, {}).get('rigid', array([]))
                for r in rigids:
                    strbuff += 'RIGID {}\n'.format(r)
    
                # end of residue
                strbuff += 'END MOLECULE {}\n\n'.format(restype)
    
            ### parameters ###
    
            strbuff += '{:<15} {:<14} {:<14} {}\n'.format('BOND', 'K', 'b0', 'Remark')
            bbuff = BytesIO()
            stdout.flush()
            savetxt(bbuff,
                    dbuff.get('bonds', empty(shape=(0,), dtype=files.prm.contents2dtypes['bonds'])),
                    fmt='%-8s %-8s %12.6f %12.6f !')
            strbuff += bbuff.getvalue().decode()
    
            strbuff += 'END BOND\n\n'
    
            # angle
            strbuff += 'ANGLE           K       theta0   Kub     S0     Remark\n'
            bbuff = BytesIO()
            savetxt(bbuff,
                    dbuff.get('angles', empty(shape=(0,), dtype=files.prm.contents2dtypes['angles'])),
                    fmt='%-8s %-8s %-8s %12.6f %12.6f %12.6f %12.6f !')
            strbuff += bbuff.getvalue().decode()
            strbuff += 'END ANGLE\n\n'
    
            # dihedral
            strbuff += 'DIHEDRAL             Kchi    n   delta  Remark\n'
            bbuff = BytesIO()
            savetxt(bbuff,
                    dbuff.get('dihedrals', empty(shape=(0,), dtype=files.prm.contents2dtypes['dihedrals'])),
                    fmt='%-8s %-8s %-8s %-8s %12.6f %12d %12.6f !')
            strbuff += bbuff.getvalue().decode()
            strbuff += 'END DIHEDRAL\n\n'
    
            # improper
            strbuff += 'IMPROPER               K       psi0    Remark\n'
            bbuff = BytesIO()
            # CHARMM data ('improper' but not 'impropers'!)
            # YL 25/02/2020: no longer a concern since in the PAR parser all dummy columns are omitted
            try:
                savetxt(bbuff,
                        dbuff.get('improper', empty(shape=(0,), dtype=files.prm.contents2dtypes['improper'])),
                        fmt='%-8s %-8s %-8s %-8s %12.6f %12d %12.6f !')
            # DL_POLY parameter data (without dummy column)
            except ValueError:
                savetxt(bbuff,
                        dbuff.get('improper', empty(shape=(0,), dtype=files.prm.contents2dtypes['improper'])),
                        fmt='%-8s %-8s %-8s %-8s %12.6f %12.6f !')
    
            strbuff += bbuff.getvalue().decode()
            strbuff += 'END IMPROPER\n\n'
    
            # vdw
            strbuff += 'VDW   Eps    Rmin/2   Eps_1-4  Rmin/2,1-4  Remark\n'
            bbuff = BytesIO()
            dbuff_nonbonded = dbuff.get('nonbonded', array([], dtype=dtype([('I','U8'),('epsilon','f8'),('Rmin/2','f8'),('eps,1-4','f8'),('Rmin/2,1-4','f8')])))
            savetxt(bbuff,
                    dbuff_nonbonded,
                    fmt='%-8s %12.6f %12.6f %12.6f %12.6f !')
            strbuff += bbuff.getvalue().decode()
            strbuff += 'END VDW\n\n'
    
            # vdw_fix
    #        strbuff += 'VDW_FIX charmm*\n'
    #        strbuff += 'END VDW_FIX\n\n'
    
            # inversion
    #        strbuff += 'INVERSION\n'
    #        strbuff += 'END INVERSION\n\n'
    
            # shell
    #        strbuff += 'SHELL\n'
    #        strbuff += 'END SHELL\n\n'
    
            strbuff += 'END POTENTIAL\n\n'
    
            strbuff += 'TERMINATE\n'
    
            fp.write(strbuff)




