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

def parse(filename, _regex=None, receiver=None, prefix='', **kwargs):
    '''Parse a PSF file'''

    from numpy      import count_nonzero, char, fromfile, genfromtxt, loadtxt, set_printoptions, searchsorted, delete, stack, where
    from re         import compile
    from os         import path, remove
    from ...base    import forcefield
    from ...data    import atomicdata, unitconvert
    from ...objects import fragment

    tmp_filename = '_chemsh_tmp.psf'

    # FIXME: not in use
    # find the number of lines of comments beginning with "HEADER/REMARK/TITLE"
    skip_header = 0
    skip_footer = 0

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

    with open(filename, 'r') as fp:
        for lineno, line in enumerate(fp):
            if line.strip().lower().endswith('!ntitle'):
                ntitles = int(line.split()[0])
                break
        titles = fp.readlines()[:ntitles]
    receiver._psf.titles = titles

    with open(filename, 'r') as fp:
        for lineno, line in enumerate(fp):
            if line.strip().lower().endswith('!natom'):
                natoms = int(line.split()[0])
                receiver._psf.natoms = natoms
                break
        atoms = fp.readlines()[:natoms]

    # for the PSF format see:
    # https://www.ks.uiuc.edu/Training/Tutorials/namd/namd-tutorial-unix-html/node23.html                    
    dtype = [
             ('indices' , 'i8'),
             ('segments', 'U4'),
             ('resIDs'  , 'i8'),
             ('residues', 'U5'),
             ('names'   , 'U4'),
             ('types'   , 'U6'),
             ('charges' , 'f8'),
             ('masses'  , 'f8'),
             ('unused'  , 'i8'),
            ]

    abuff = genfromtxt(atoms,
                       skip_header=skip_header,
                       skip_footer=skip_footer,
                       dtype=dtype,
                      )

    # assign more entries (do NOT reorder them!)
    for i in [ 'types', 'charges', 'masses', 'names', 'residues', 'resIDs', 'segments' ]:
        setattr(receiver, i, abuff[i])

    tmp_filename = '_chemsh_tmp.psf'

    # bonds
    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
        for lineno, line in enumerate(fp):
            if line.strip().lower().endswith('!nbond: bonds'):
                nbonds = int(line.split()[0])
                break
        bonds = fp.readlines()[:int(nbonds/4)+1]
        newfp.write(''.join(bonds))
    abuff = fromfile(tmp_filename, dtype=int, sep=' ')
    receiver._psf.bonds = abuff.reshape((nbonds,2))

    # angles
    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
        for lineno, line in enumerate(fp):
            if line.strip().lower().endswith('!ntheta: angles'):
                nangles = int(line.split()[0])
                break
        angles = fp.readlines()[:int(nangles/3)+1]
        newfp.write(''.join(angles))
    abuff = fromfile(tmp_filename, dtype=int, sep=' ')
    receiver._psf.angles = abuff.reshape((nangles,3))

    # dihedrals
    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
        for lineno, line in enumerate(fp):
            if line.strip().lower().endswith('!nphi: dihedrals'):
                ndihedrals = int(line.split()[0])
                break
        dihedrals = fp.readlines()[:int(ndihedrals/2)+1]
        newfp.write(''.join(dihedrals))
    abuff = fromfile(tmp_filename, dtype=int, sep=' ')
    receiver._psf.dihedrals = abuff.reshape((ndihedrals,4))

    # impropers
    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
        for lineno, line in enumerate(fp):
            if line.strip().lower().endswith('!nimphi: impropers'):
                nimpropers = int(line.split()[0])
                break
        impropers = fp.readlines()[:int(nimpropers/2)+1]
        newfp.write(''.join(impropers))
    abuff = fromfile(tmp_filename, dtype=int, sep=' ')
    receiver._psf.impropers = abuff.reshape((nimpropers,4))

# not in use now
#    # donors
#    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
#        for lineno, line in enumerate(fp):
#            if line.strip().lower().endswith('!ndon: donors'):
#                ndonors = int(line.split()[0])
#                break
#        donors = fp.readlines()[:ndonors]
#        newfp.write(''.join(donors))
#
#    # acceptors
#    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
#        for lineno, line in enumerate(fp):
#            if line.strip().lower().endswith('!nacc: acceptors'):
#                nacceptors = int(line.split()[0])
#                break
#        acceptors = fp.readlines()[:nacceptors]
#        newfp.write(''.join(acceptors))
#
#    # nbs
#    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
#        for lineno, line in enumerate(fp):
#            if line.strip().lower().endswith('!nnb'):
#                nnbs = int(line.split()[0])
#                break
#        nbs = fp.readlines()[:nnbs]
#        newfp.write(''.join(nbs))

    # groups
    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
        for lineno, line in enumerate(fp):
            if line.strip().lower().endswith('!ngrp'):
                ngroups = int(line.split()[0])
                break
        groups = fp.readlines()[:ngroups]
        newfp.write(''.join(groups))
    abuff = fromfile(tmp_filename, dtype=int, sep=' ')
    receiver._psf.groups = abuff

    # cross-terms
    with open(filename, 'r') as fp, open(tmp_filename, 'w') as newfp:
        # cross-terms may not exist in the PSF file
        ncterms = -1
        for lineno, line in enumerate(fp):
            if line.strip().lower().endswith('!ncrterm: cross-terms'):
                ncterms = int(line.split()[0])
                break
        # NB: here is number of lines rather than number of cterms!
        cterms = fp.readlines()[:ncterms+1]
        newfp.write(''.join(cterms))
    abuff = fromfile(tmp_filename, dtype=int, sep=' ')
    receiver._psf.cterms = abuff.reshape((int(abuff.size/4),4))

    # remove the tmp file
    remove(tmp_filename)


def write(filename, fragObj, **kwargs):
    '''Writer of PSF'''

    from numpy import arange, array2string, dtype, ndarray, savetxt
    from io    import BytesIO
    from sys   import maxsize, stdout
    from math  import ceil

    strbuff = ''

    strbuff += 'PSF CMAP\n'

    strbuff += '\n'

    try:
        strbuff += '{:>8} {}\n'.format(len(fragObj._psf.titles), '!NTITLE')
    except:
        print(f'\n >>> ERROR: {fragObj} does not contain PSF information')
        stdout.flush()
        exit(999)

    strbuff += ''.join(fragObj._psf.titles)

    strbuff += '\n{:>8} {}\n'.format(fragObj.natoms, '!NATOM')

    # convert to bytes: BytesIO, not StringIO!
    bbuff = BytesIO()
    descr = [('indices' , '<i8'),
             ('segments', '<U8'),
             ('resIDs'  , '<i8'),
             ('residues', '<U8'),
             ('names'   , '<U8'),
             ('types'   , '<U8'),
             ('charges' , '<f8'),
             ('masses'  , '<f8'),
            ]
    dtp = dtype(descr)
    abuff = ndarray((fragObj.natoms,), dtype=dtp)
    abuff['indices'] = arange(fragObj.natoms)+1
    for i in [ 'segments', 'resIDs', 'residues', 'names', 'types', 'charges', 'masses' ]:
        abuff[i][:] = getattr(fragObj, i)[:]
    fmt = '%8d %4s %-4d %-4s %-4s %-5s %9.6f      %8.4f           0'
    savetxt(bbuff,
            abuff,
            fmt=fmt)
    # convert bytes to string
    outstr = bbuff.getvalue().decode().rstrip()
    strbuff += outstr

    # bonds
    strbuff += '\n\n{:>8} {}\n'.format(fragObj._psf.nbonds    , '!NBOND: bonds')
    strbuff += ' '
    strbuff += array2string(fragObj._psf.bonds.ravel(),
                            separator=' ',
                            formatter={'int':lambda x:'{:>7d}'.format(x)},
                            # Yl 27/11/2021: make sure all numbers are printed instead of "..."
                            threshold=maxsize,
                            max_line_width=65)[1:-1] + '\n'

    # angles
    strbuff += '\n{:>8} {}\n'.format(fragObj._psf.nangles   , '!NTHETA: angles')
    strbuff += ' '
    strbuff += array2string(fragObj._psf.angles.ravel(),
                            separator=' ',
                            formatter={'int':lambda x:'{:>7d}'.format(x)},
                            threshold=maxsize,
                            max_line_width=74)[1:-1] + '\n'

    # dihedrals
    strbuff += '\n{:>8} {}\n'.format(fragObj._psf.ndihedrals, '!NPHI: dihedrals')
    strbuff += ' '
    strbuff += array2string(fragObj._psf.dihedrals.ravel(),
                            separator=' ',
                            formatter={'int':lambda x:'{:>7d}'.format(x)},
                            threshold=maxsize,
                            max_line_width=65)[1:-1] + '\n'

    # impropers
    strbuff += '\n{:>8} {}\n'.format(fragObj._psf.nimpropers, '!NIMPHI: impropers')
    strbuff += ' '
    strbuff += array2string(fragObj._psf.impropers.ravel(),
                            separator=' ',
                            formatter={'int':lambda x:'{:>7d}'.format(x)},
                            threshold=maxsize,
                            max_line_width=65)[1:-1] + '\n'

    # TODO
    # donors
    strbuff += '\n{:>8} {}\n'.format(0, '!NDON: donors')

    # TODO
    # acceptors
    strbuff += '\n{:>8} {}\n'.format(0, '!NACC: acceptors')

    # TODO
    # non-bonded
    strbuff += '\n{:>8} {}\n'.format(0, '!NNB')

    # groups
    strbuff += '\n{:>8}{:>8} {}\n'.format(1, 0, '!NGRP')
    strbuff += '{:>8}{:>8}{:>8}\n'.format(0,0,0)

    # cross-terms
    # NB: here is number of lines rather than number of cterms!
    strbuff += '\n{:>8} {}\n'.format(ceil(fragObj._psf.ncterms/4), '!NCRTERM: cross-terms')
    strbuff += ' '
    strbuff += array2string(fragObj._psf.cterms.ravel(),
                            separator=' ',
                            formatter={'int':lambda x:'{:>7d}'.format(x)},
                            max_line_width=65)[1:-1] + '\n'


    try:
        with open(filename, 'w') as fp:
            fp.write(strbuff)
    except:
        filename.write(strbuff)




