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

from ...io   import format
from ...base import forcefield
from .       import dlp_field

# AMBER parameter/topology file
# see: https://ambermd.org/FileFormats.php#topo.cntrl
def parse(filename, _regex=None, receiver=None, prefix='', **kwargs):
    '''Parse an AMBER parameter/topology (prmtop) file'''

    from numpy import absolute, arange, argsort, char, empty, genfromtxt, inf, pi, unique
    from os    import path

    delimiters = []
    dtype      = None
    entry      = ''
    row        = 0
    tbuff      = []

    if not receiver:
        receiver = {}
        receiver.update({'vdw':{}})
        receiver.update({'unit':'kcal'})
        receiver.update({'ntypes':1})

    # YL 05/11/2021: let AMBER PRMTOP share same datatype with DL_POLY's FIELD
    contents2dtypes = dlp_field.contents2dtypes

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

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

        for line in fp:

            line_lower = line.strip().lower()

            if line_lower.startswith('%version'):
                pass

            elif line_lower.startswith('%flag'):

                # data of the previous block because we can't collect the whole contents until the block ends
                if row:
                    try:
                        abuff = genfromtxt(tbuff,
                                           dtype=dtype,
                                           delimiter=delimiters,
                                          )
                    except:
                        abuff = line

                # title
                if entry == 'title':
                    title = ''.join(abuff.ravel()).strip()
                    receiver.update({title:{}})
                    receiver[title].update({'nmolecules':1})
                    receiver[title].update({'title':title})
                    receiver[title].update({'atom':{}})
                    receiver[title].update({'bond':{}})
                    receiver[title].update({'angle':{}})
                    receiver[title].update({'dihedral':{}})

                # number of data points of each term
                elif entry == 'pointers':
                    natom,    ntypes, nbonh,  mbona,  ntheth, mtheta, \
                    nphih,    mphia,  nhparm, nparm,  nnb,    nres  , \
                    nbona,    ntheta, nphia,  numbnd, numang, nptra , \
                    natyp,    nphb,   ifpert, nbper,  ngper,  ndper , \
                    mbper,    mgper,  mdper,  ifbox,  nmxrs,  ifcap , \
                    numextra, ncopy = abuff.ravel()[:32]

                    # number of total entries
                    receiver[title]['atom'].update({'nentries':natom})
                    receiver[title]['bond'].update({'nentries':nbonh+mbona})
                    receiver[title]['angle'].update({'nentries':ntheth+mtheta})
                    receiver[title]['dihedral'].update({'nentries':nphih+mphia})

                    nvdws = int(natyp*(natyp+1)/2)
                    receiver['vdw'].update({'nentries':nvdws})
                    receiver['vdw'].update({'data':empty(shape=(nvdws,), dtype=contents2dtypes['vdw'])})

                # atom type indices
                elif entry == 'atom_type_index':
                    # we use this list of indices to remove the duplicates
                    vdw_atom_type_index = abuff.ravel()[:natom]
                    unique_vdw_atom_type_index, unique_vdw_atom_type_indices = unique(vdw_atom_type_index, return_index=True)

                # atom names
                elif entry == 'atom_name':
                    pass

                # atom types
                elif entry == 'amber_atom_type':
                    all_atom_types = abuff.ravel()[:natom]
                    unique_all_atom_types, unique_all_atom_types_indices = unique(all_atom_types, return_index=True)
                    unique_types_vdw_indices = vdw_atom_type_index[unique_all_atom_types_indices]
                    receiver[title]['atom']['data']['name'] = abuff.ravel()[:natom]

                # atom charges
                elif entry == 'charge':
                    # charge block comes before amber_atom_type
                    receiver[title]['atom'].update({'data':empty(shape=(natom,), dtype=contents2dtypes['atom'])})
                    receiver[title]['atom']['data']['nrepeats'] = 1
                    # NB: note the conversion factor!
                    receiver[title]['atom']['data']['charge'] = abuff.ravel()[:natom]/18.2223
                    
                # atom masses
                elif entry == 'mass':
                    receiver[title]['atom']['data']['mass'] = abuff.ravel()[:natom]

                # bond force constant: kcal/(mol*Angstrom**2)
                elif entry == 'bond_force_constant':
                    bond_force_constant = abuff.ravel()[:numbnd] # numbnd is the number of bond types

                # bond equil value
                elif entry == 'bond_equil_value':
                    bond_equil_value = abuff.ravel()[:numbnd] # numbnd is the number of bond types

                # bond entries containing hydrogen: I, J, and index in bond_force_constant and bond_equil_value
                elif entry == 'bonds_inc_hydrogen':
                    receiver[title]['bond'].update({'data':empty(shape=(nbonh+mbona,), dtype=contents2dtypes['bond'])})
                    receiver[title]['bond']['data']['type'] = 'harm'
                    receiver[title]['bond']['data']['I'][:nbonh] = (abuff.ravel()[0::3][:nbonh]/3+1).astype(int)
                    receiver[title]['bond']['data']['J'][:nbonh] = (abuff.ravel()[1::3][:nbonh]/3+1).astype(int)
                    # the original indexing in AMBER starts from 1
                    indices = abuff.ravel()[2::3][:nbonh] - 1
                    # NB: in DL_POLY the coefficient is k/2 (see the manual)
                    receiver[title]['bond']['data']['Kb'][:nbonh] = bond_force_constant[indices]*2.0
                    receiver[title]['bond']['data']['b0'][:nbonh] = bond_equil_value[indices]

                # bond entries not containing hydrogen: I, J, and index in bond_force_constant and bond_equil_value
                elif entry == 'bonds_without_hydrogen':
                    receiver[title]['bond']['data']['I'][nbonh:] = (abuff.ravel()[0::3][:mbona]/3+1).astype(int)
                    receiver[title]['bond']['data']['J'][nbonh:] = (abuff.ravel()[1::3][:mbona]/3+1).astype(int)
                    # the original indexing in AMBER starts from 1
                    indices = abuff.ravel()[2::3][:mbona] - 1
                    # NB: in DL_POLY the coefficient is k/2 (see the manual)
                    receiver[title]['bond']['data']['Kb'][nbonh:] = bond_force_constant[indices]*2.0
                    receiver[title]['bond']['data']['b0'][nbonh:] = bond_equil_value[indices]

                # angle force constant: kcal/(mol*radian**2)
                elif entry == 'angle_force_constant':
                    angle_force_constant = abuff.ravel()[:numang] # numang is the number of angle types

                # angle equil value
                elif entry == 'angle_equil_value':
                    angle_equil_value = abuff.ravel()[:numang] # numang is the number of angle types

                # angle entries with hydrogen: I, J, and index in angle_force_constant and angle_equil_value
                elif entry == 'angles_inc_hydrogen':
                    receiver[title]['angle'].update({'data':empty(shape=(ntheth+mtheta,), dtype=contents2dtypes['angle'])})
                    receiver[title]['angle']['data']['type'] = 'harm'
                    receiver[title]['angle']['data']['I'][:ntheth] = (abuff.ravel()[0::4][:ntheth]/3+1).astype(int) 
                    receiver[title]['angle']['data']['J'][:ntheth] = (abuff.ravel()[1::4][:ntheth]/3+1).astype(int)
                    receiver[title]['angle']['data']['K'][:ntheth] = (abuff.ravel()[2::4][:ntheth]/3+1).astype(int)
                    # the original indexing in AMBER starts from 1
                    indices = abuff.ravel()[3::4][:ntheth] - 1
                    # NB: in DL_POLY the coefficient is k/2 (see the manual)
                    receiver[title]['angle']['data']['Ktheta'][:ntheth] = angle_force_constant[indices]*2.0
                    receiver[title]['angle']['data']['Theta0'][:ntheth] = angle_equil_value[indices]*180.0/pi

                # angle entries without hydrogen: I, J, and index in angle_force_constant and angle_equil_value
                elif entry == 'angles_without_hydrogen':
                    receiver[title]['angle']['data']['I'][ntheth:] = (abuff.ravel()[0::4][:mtheta]/3+1).astype(int)
                    receiver[title]['angle']['data']['J'][ntheth:] = (abuff.ravel()[1::4][:mtheta]/3+1).astype(int)
                    receiver[title]['angle']['data']['K'][ntheth:] = (abuff.ravel()[2::4][:mtheta]/3+1).astype(int)
                    # the original indexing in AMBER starts from 1
                    indices = abuff.ravel()[3::4][:mtheta] - 1
                    # NB: in DL_POLY the coefficient is k/2 (see the manual)
                    receiver[title]['angle']['data']['Ktheta'][ntheth:] = angle_force_constant[indices]*2.0
                    receiver[title]['angle']['data']['Theta0'][ntheth:] = angle_equil_value[indices]*180.0/pi

                # dihedral force constant: kcal/mol
                elif entry == 'dihedral_force_constant':
                    dihedral_force_constant = abuff.ravel()[:nptra] # nptra is the number of dihedral types

                # dihedral periodicity
                elif entry == 'dihedral_periodicity':
                    dihedral_periodicity = abuff.ravel()[:nptra] # nptra is the number of dihedral types

                # dihedral phase
                elif entry == 'dihedral_phase':
                    dihedral_phase = abuff.ravel()[:nptra] # nptra is the number of dihedral types

                # 1-4 electrostatic scaling constants
                elif entry == 'scee_scale_factor':
                    dihedral_14e = abuff.ravel()[:nptra] # nptra is the number of dihedral types
                    masks_zero = dihedral_14e == 0.0
                    dihedral_14e[masks_zero] = inf

                # 1-4 vdW scaling constants
                elif entry == 'scnb_scale_factor':
                    dihedral_14vdw = abuff.ravel()[:nptra] # nptra is the number of dihedral types
                    masks_zero = dihedral_14vdw == 0.0
                    dihedral_14vdw[masks_zero] = inf

                # dihedral entries including hydrogen: I, J, and index in dihedral_force_constant and dihedral_equil_value
                elif entry == 'dihedrals_inc_hydrogen':
                    receiver[title]['dihedral'].update({'data':empty(shape=(nphih+mphia,), dtype=contents2dtypes['dihedral'])})
                    receiver[title]['dihedral']['data']['type'] = 'cos'
                    # NB: up to here they are still raw indices of coordinates and will be converted to atom indices at the bottom
                    receiver[title]['dihedral']['data']['I'][:nphih] = abuff.ravel()[0::5][:nphih]
                    receiver[title]['dihedral']['data']['J'][:nphih] = abuff.ravel()[1::5][:nphih]
                    receiver[title]['dihedral']['data']['K'][:nphih] = abuff.ravel()[2::5][:nphih]
                    receiver[title]['dihedral']['data']['L'][:nphih] = abuff.ravel()[3::5][:nphih]
                    # the original indexing in AMBER starts from 1
                    indices = abuff.ravel()[4::5][:nphih] - 1
                    receiver[title]['dihedral']['data']['Kchi'][:nphih]  = dihedral_force_constant[indices]
                    receiver[title]['dihedral']['data']['n'][:nphih]     = dihedral_periodicity[indices]
                    receiver[title]['dihedral']['data']['delta'][:nphih] = dihedral_phase[indices]*180.0/pi
                    receiver[title]['dihedral']['data']['14e'][:nphih]   = 1.0/dihedral_14e[indices]
                    receiver[title]['dihedral']['data']['14vdw'][:nphih] = 1.0/dihedral_14vdw[indices]

                # dihedral entries without hydrogen: I, J, and index in dihedral_force_constant and dihedral_equil_value
                elif entry == 'dihedrals_without_hydrogen':
                    # NB: up to here they are still raw indices of coordinates and will be converted to atom indices at the bottom
                    receiver[title]['dihedral']['data']['I'][nphih:] = abuff.ravel()[0::5][:mphia]
                    receiver[title]['dihedral']['data']['J'][nphih:] = abuff.ravel()[1::5][:mphia]
                    receiver[title]['dihedral']['data']['K'][nphih:] = abuff.ravel()[2::5][:mphia]
                    receiver[title]['dihedral']['data']['L'][nphih:] = abuff.ravel()[3::5][:mphia]
                    # the original indexing in AMBER starts from 1
                    indices = abuff.ravel()[4::5][:mphia] - 1
                    receiver[title]['dihedral']['data']['Kchi'][nphih:]  = dihedral_force_constant[indices]
                    receiver[title]['dihedral']['data']['n'][nphih:]     = dihedral_periodicity[indices]
                    receiver[title]['dihedral']['data']['delta'][nphih:] = dihedral_phase[indices]*180.0/pi
                    receiver[title]['dihedral']['data']['14e'][nphih:]   = 1.0/dihedral_14e[indices]
                    receiver[title]['dihedral']['data']['14vdw'][nphih:] = 1.0/dihedral_14vdw[indices]

                # nonbonded parameter indices
                elif entry == 'nonbonded_parm_index':
                    nonbonded_indices = abuff.ravel()[:ntypes**2]

                # Lennard-Jones r**12 terms
                elif entry == 'lennard_jones_acoef':
                    lennard_jones_12 = abuff.ravel()[:nvdws]

                # Lennard-Jones r**6 terms
                elif entry == 'lennard_jones_bcoef':
                    lennard_jones_6 = abuff.ravel()[:nvdws]

                # H-bond r**12 terms
                elif entry == 'hbond_acoef':
                    hbond_12 = abuff.ravel()[:nphb]

                # H-bond r**10 terms
                elif entry == 'hbond_bcoef':
                    hbond_10 = abuff.ravel()[:nphb]


                entry = line_lower.split('flag')[1].strip()

            elif line_lower.startswith('%format'):
                format = line_lower.split('(')[1].split(')')[0]
                if 'a' in format:
                    nnumbers = int(format.split('a')[0])
                    ndigits  = int(format.split('a')[1])
                    dtype = 'U8'
                elif 'i' in format:
                    nnumbers = int(format.split('i')[0])
                    ndigits  = int(format.split('i')[1])
                    dtype = 'i8'
                elif 'e' in format:
                    nnumbers = int(format.split('e')[0])
                    ndigits  = int(format.split('e')[1].split('.')[0])
                    dtype = 'f8'

                # delimiters for genfromtxt()
                delimiters = [ ndigits ]*nnumbers
                # row index of the current block
                row   =  0
                tbuff = []

                continue

            else:

                row   +=  1
                tbuff += [line]
             
        # we have to deal with non-bonded terms afterall because the "amber_atom_type" block is the last to appear
        # nonbonded_parm_index constructed from unique_all_atom_types: index = ICO(ntypes*(IAC(i)-1)+IAC(j))
        counter = 0
        for i, unique_atom_type_i in enumerate(unique_all_atom_types):
            for j, unique_atom_type_j in enumerate(unique_all_atom_types):
                if i <= j:
                    counter += 1
                    # map to vdW types (NB: unique_types_vdw_indices start from 1)
                    ii = unique_types_vdw_indices[i] - 1
                    jj = unique_types_vdw_indices[j] - 1
                    index_index = ntypes*(ii)+jj+1
                    index_parameter = nonbonded_indices[index_index-1]
                    # Lennard-Jones
                    if index_parameter > 0:
                        # NB: receiver['vdw']['data'][i]['field'] doesn't return a reference
                        receiver['vdw']['data'][counter-1] =  unique_all_atom_types[i], \
                                                              unique_all_atom_types[j], \
                                                             '12-6', \
                                                              lennard_jones_12[index_parameter-1], \
                                                              lennard_jones_6[index_parameter-1]
                    # H-bond
                    elif index_parameter < 0:
                        # NB: receiver['vdw']['data'][i]['field'] doesn't return a reference
                        receiver['vdw']['data'][counter-1] =  unique_all_atom_types[i], \
                                                              unique_all_atom_types[j], \
                                                             'hbnd', \
                                                              hbond_12[-index_parameter-1], \
                                                              hbond_10[-index_parameter-1]
                    else:
                        print(f"\n WARNING: file {filename} may be problematic about non-bonded interactions!")

# YL 14/06/2021: in Tcl-ChemShell it is "klm/power" used to emulate the AMBER TIP3P model but here we can fully trust the "hbnd" scheme
        # AMBER TIP3P model
#        masks_water = (char.strip(char.lower(receiver['vdw']['data']['A']))=='hw')*(char.strip(char.lower(receiver['vdw']['data']['B']))=='ow') \
#                    + (char.strip(char.lower(receiver['vdw']['data']['A']))=='ow')*(char.strip(char.lower(receiver['vdw']['data']['B']))=='hw')
#        indices_klm = arange(nvdws)[masks_water]
#        # NB: receiver['vdw']['data'][i]['field'] doesn't return a reference
#        for i in indices_klm:
#            receiver['vdw']['data'][i] = receiver['vdw']['data'][i]['A'], \
#                                         receiver['vdw']['data'][i]['B'], \
#                                        'klm', \
#                                         receiver['vdw']['data'][i]['a'], \
#                                         12.0

        # ignore end-group interactions to prevent double counting of 1-4
        masks_k = receiver[title]['dihedral']['data']['K'] < 0
        indices_k = arange(nphih+mphia)[masks_k]
        receiver[title]['dihedral']['data']['14e'][indices_k]   = 0.0
        receiver[title]['dihedral']['data']['14vdw'][indices_k] = 0.0
        # turn coordinate indices to positive and atom indices
        receiver[title]['dihedral']['data']['I'] = (absolute(receiver[title]['dihedral']['data']['I'])/3+1).astype(int)
        receiver[title]['dihedral']['data']['J'] = (absolute(receiver[title]['dihedral']['data']['J'])/3+1).astype(int)
        receiver[title]['dihedral']['data']['K'] = (absolute(receiver[title]['dihedral']['data']['K'])/3+1).astype(int)
        receiver[title]['dihedral']['data']['L'] = (absolute(receiver[title]['dihedral']['data']['L'])/3+1).astype(int)



    return receiver


