#!/usr/bin/env python
from argparse import ArgumentParser
import os,sys


# ==============================================================================
# CLASSES
# ==============================================================================
class ForceField:
    def __init__(self, fname=None, bForgiving=True):
        # if input is file
        if fname is not None:
            self.bForgiving = bForgiving
            self.ff = self.read_file(fname)

    def read_file(self, fname, append=False):
        # check if file exists
        if os.path.isfile(fname)==False:
            return('')

        # we store all ff info in a dict of dicts, where
        # each dict has a tuple for key and a string for value

        # if we are loading a new file/ff then we make new dicts:
        if append is False:
            ff = {}
            ff['atomtypes'] = {}
            ff['pairtypes'] = {}
            ff['bondtypes'] = {}
            ff['constrainttypes'] = {}
            ff['angletypes'] = {}
            ff['dihedraltypes'] = {}
            ff['dihedraltypesX'] = {} # dihedrals with the wildcards
        # otherwise we load the existing parameters
        elif append is True:
            ff = self.ff

        # read file lines
        lines = [x.strip().split(';')[0] for x in open(fname, 'r').readlines()
                 if x.strip().split(';')[0]]
        # remove ifdef lines: this might be an issue!
        lines = [l for l in lines if l[0] not in ('#')]

        # remove * lines
#        lines = [l for l in lines if l[0] not in ('*')]

        section = None
        dihXnum = 1 # wildcard dihedrals will be numbered
        # read each line with ff info
        for line in lines:
            # assess whether we are now parsing a different section
            if line[0] == '[':
                section = line.strip('[] ')
                continue

            # --------------------------------------------
            # cases where 1 atomtype identifies the params
            # --------------------------------------------
            if section in ('atomtypes'):
                el = line.split()
                key = (el[0])
                # check whether key is already presnet in ff
                if key not in ff[section]:
                    ff[section][key] = line
#                else:
                    # silence warning if it is about water atoms - these
                    # warnings are due to the fact we ignore ifdef statements
                    # but since the params are specific for water they are
                    # not an issue
#                    if key not in ['OW', 'HW', 'HT', 'OT', 'OWT3', 'OWT4',
#                                   'HWT3', 'HWT4', 'MWT4']:
#                        print('\n    WARNING! Atom parameters for the '
#                              ' following atoms are already present in the '
#                              'force field object'
#                              ' and will be ignored:\n    %s' % key)

            # --------------------------------------------
            # cases where 2 atomtypes identify the params
            # --------------------------------------------
            elif section in ('pairtypes', 'bondtypes', 'constrainttypes'):
                el = line.split()
                key = (el[0], el[1])
                # check whether key is already present in ff
                if key not in ff[section]:
                    ff[section][key] = line
#                else:
                    # silence some warnings: idem as above. In this case
                    # these are due to protein vsite constraints we do not
                    # care about
#                    if key not in [('MNH3', 'CT1'), ('MNH3', 'CT2'),
#                                   ('MNH3', 'CT3'), ('MNH3', 'MNH3'),
#                                   ('MCH3', 'CT1'), ('MCH3', 'CT2'),
#                                   ('MCH3', 'CT3'), ('MCH3', 'MCH3'),
#                                   ('MCH3', 'S')]:
#                        print('\n    WARNING! %s parameters for the  '
#                              'following atoms are already present in the '
#                              'force field object and will be '
#                              'ignored:\n    %s' % (section, " ".join(key)))
            # --------------------------------------------
            # cases where 3 atomtypes identify the params
            # --------------------------------------------
            elif section in ('angletypes'):
                el = line.split()
                key = (el[0], el[1], el[2])
                # check whether key is already presnet in ff
                if key not in ff[section]:
                    ff[section][key] = line
#                else:
#                    print('\n    WARNING! Angle parameters for the following '
#                          'atoms are already present in the force field object'
#                          ' and will be ignored:\n    %s' % " ".join(key))

            # --------------------------------------------
            # cases where 4 atomtypes identify the params
            # --------------------------------------------
            elif section in ('dihedraltypes'):
                el = line.split()
                dtype = int(el[4])

                # check if the dihedraltype has a wildcard
                if el[0]=='X' or el[1]=='X' or el[2]=='X' or el[3]=='X':
                    if dtype in [4, 9]:
                        # key contains: a1, a2, a3, a4, dihedral type, multiplicity
                        key = [el[0], el[1], el[2], el[3], el[4], el[7]]
                        ff['dihedraltypesX'][dihXnum] = [key,[line]]
                        dihXnum += 1
#                        print ff['dihedraltypesX'][1][0],ff['dihedraltypesX'][1][1]
#                        sys.exit(0)
                    elif dtype in [2]:
                        # key contains: a1, a2, a3, a4, dihedral type, multiplicity
                        key = [el[0], el[1], el[2], el[3], el[4]]
                        ff['dihedraltypesX'][dihXnum] = [key,[line]]
                        dihXnum += 1
                    else:
                        exit('\nBUG: dihedral type not expected/checked/'
                             'implemented yet in class ForceField. '
                             'Script needs to be extended.')
                    continue

                ####### standard dihedrals (no wildcards) ########
                # need different keys bacuse certain torsions have multiple
                # line with different multiplicities
                if dtype in [4, 9]:
                    # key contains: a1, a2, a3, a4, dihedral type, multiplicity
                    key = (el[0], el[1], el[2], el[3], el[4], el[7])
                    # check whether key is already presnet in ff
                    if key not in ff[section]:
                        ff[section][key] = line
#                    else:
#                        print('\n    WARNING! Dihedral parameters for the '
#                              'following atoms are already present in the '
#                              'force field object '
#                              'and will be ignored:\n    %s' % " ".join(key))
                elif dtype in [2]:
                    # key contains: a1, a2, a3, a4, dihedral type, multiplicity
                    key = (el[0], el[1], el[2], el[3], el[4])
                    # check whether key is already presnet in ff
                    if key not in ff[section]:
                        ff[section][key] = line
#                    else:
#                        print('\n    WARNING! Dihedral parameters for the '
#                              'following atoms are already present in the '
#                              'force field object '
#                              'and will be ignored:\n    %s' % " ".join(key))
                else:
                    exit('\nBUG: dihedral type not expected/checked/'
                         'implemented yet in class ForceField. '
                         'Script needs to be extended.')

            # --------------------------------------------
            # unkown cases
            # --------------------------------------------
            else:
                if self.bForgiving==False:
                    exit('ERROR: topology section %s not recognised/supported.' % section)
                else:
                    print('ERROR: topology section %s not recognised/supported.' % section)

        if append is False:
            return ff


# ==============================================================================
# FUNCTIONS
# ==============================================================================
def parse_options():
    parser = ArgumentParser()
    choices = ['prm', 'ff']
    parser.add_argument('-itp', dest='itp',
                        help='Gromacs topology, containing CGenFF parameters '
                        ' for a ligand obtained from the CGenFF program '
                        '(paramchem) and converted with the '
                        '"cgenff_charmm2gmx.py" script.',
                        required=True)
    parser.add_argument('-ff', dest='ff',
                        help='Path to the Charmm36 force '
                        'field, which contains the CGenFF parameters.',
                        required=True)
    parser.add_argument('-prm', dest='prm',
                        help='File with additional parameters that were '
                        'assigned by analogy by the CGenFF program '
                        '(paramchem). (This prm is in gromacs format analogous to .itp)',
                        required=False)
    parser.add_argument('-o', dest='outfname',
                        help='Output itp file which contains all of the force '
                        'field parameters explicitly (excluding the atomtypes'
                        '). Default is "cgenff.itp"',
                        default='cgenff.itp')
    parser.add_argument('--priority', dest='priority',
                        help='This determines the behavior when dealing with '
                        'duplicate parameters. Options are "prm" or "ff" '
                        '(default is "ff"). With "ff" the parameters already in'
                        'the force field take priority and are used; with '
                        '"prm" the parameters present in the prm file are '
                        'used instead in case of duplicate entries. The '
                        'latter option is useful if you have optimized some '
                        'terms specifically for the molecule at hand, e.g. '
                        'with the Force Field Toolkit.',
                        default='ff', choices=choices)
    parser.add_argument('--bForgiving', dest='bForgiving',
                        help='Should the script exit when some ff section is not recognized.',
                        default=False)
    args = parser.parse_args()
    return args


def list2file(lst, outfname):
    with open(outfname, 'w') as f:
        for l in lst:
            f.write('{line}\n'.format(line=l))


def get_atoms_map(itp):
    # read the atoms from the itp file
    lines = [x.strip().split(';')[0] for x in open(itp, 'r').readlines()
             if x.strip().split(';')[0]]

    # map atom ids to their atomtypes
    section = ''
    id2type_map = {}
    for line in lines:
        # assess whether we are now parsing a different section
        if line[0] == '[':
            section = line.strip('[] ')
            continue
        if section == 'atoms':
            el = line.split()
            a_id = el[0]
            a_type = el[1]
            id2type_map[a_id] = a_type

    return id2type_map

def check_wildcard_dihedral( key, at ):
    found = True
    for k,a in zip(key,at):
        if k=='X':
            continue
        elif k==a:
            continue
        else:
            found=False
            return(found)
    return(found)

def fill_params(itp, a_map, ff):
    # read the atoms from the itp file
    lines = [x.strip().split(';')[0] for x in open(args.itp, 'r').readlines()
             if x.strip().split(';')[0]]

    out = []

    section = ''
    for line in lines:
        if line.startswith('#'):
            continue

        # assess whether we are now parsing a different section
        if line[0] == '[':
            section = line.strip('[] ')
            out.append(line)
            continue

        # -----
        # Bonds
        # -----
        if section == 'bonds':
            el = line.split()
            a1 = a_map[el[0]]
            a2 = a_map[el[1]]
            key = (a1, a2)

            # order does not matter: check both
            try:
                params = ff['bondtypes'][key]
            except:
                key = key[::-1]
                params = ff['bondtypes'][key]

            # remove the atomtypes from the ff string
            paramsplit = params.split()
            params = ""
            for p in paramsplit:
                if p not in key:
                    params += " "+str(p)
#            for k in key:
                #params = params.replace(k, '')
		#print k," --- VG ---",params

            # create new topology line
            newline = '{a1:<5}{a2:<5}{params}'.format(a1=el[0],
                                                      a2=el[1],
                                                      params=params)
            out.append(newline)

        # ------
        # Angles
        # ------
        elif section == 'angles':
            el = line.split()
            a1 = a_map[el[0]]
            a2 = a_map[el[1]]
            a3 = a_map[el[2]]
            key = (a1, a2, a3)

            # order does not matter: check both
            try:
                params = ff['angletypes'][key]
            except:
                key = key[::-1]
                params = ff['angletypes'][key]

            # remove the atomtypes from the ff string
            paramsplit = params.split()
            params = ""
            for p in paramsplit:
                if p not in key:
                    params += " "+str(p)
#            for k in key:
#                params = params.replace(k, '')

            # create new topology line
            newline = '{a1:<5}{a2:<5}{a3:<5}{params}'.format(a1=el[0],
                                                             a2=el[1],
                                                             a3=el[2],
                                                             params=params)
            out.append(newline)

        # ---------
        # Dihedrals
        # ---------
        elif section == 'dihedrals':
            el = line.split()
            a1 = a_map[el[0]]
            a2 = a_map[el[1]]
            a3 = a_map[el[2]]
            a4 = a_map[el[3]]
            dtype = el[4]

            # get largest multiplicity in ff
            max_pn = max([int(x[-1]) for x in ff['dihedraltypes'].keys()])

            # -----------------------------
            # dihedrals with multiplicity
            # -----------------------------
            if int(dtype) in [4, 9]:
                found = False
                # check various multiplicities
                for pn in range(1, max_pn+1):
                    # check both orders
                    try:
                        key = (a1, a2, a3, a4, dtype, str(pn))
                        params = ff['dihedraltypes'][key]
                        found = True
                    except:
                        try:
                            key = (a4, a3, a2, a1, dtype, str(pn))
                            params = ff['dihedraltypes'][key]
                            found = True
                        except:
                            # if missing it's because there is no param
                            # with this multiplicity. This can however
                            # be dangerous, so we use "found" to make sure
                            # the term is not skipped (i.e. at least one
                            # torsion needs to be found)
                            continue

                    if found==True:
                        # remove the atomtypes from the ff string
                        paramsplit = params.split()
                        params = ""
                        for p in paramsplit:
                            if p not in [a1,a2,a3,a4]:
                                params += " "+str(p)
                        # create new topology line
                        newline = ('{a1:<5}{a2:<5}{a3:<5}{a4:<5}'
                                   '{params}'.format(a1=el[0], a2=el[1],
                                                     a3=el[2], a4=el[3],
                                                     params=params))
                        out.append(newline)


		############################################################
		# if not yet found a dihedral, search through the wildcards #
                # the search depends on the dihedral number #
                # the first match (or matches of different multiplicities will be used #
                if found==False:
                    dihXparams = []
                    prev = 0
                    for d in ff['dihedraltypesX']:
                        # first check that the dihedral is of the correct type
                        if int(ff['dihedraltypesX'][d][0][4]) not in [4,9]:
                            continue
                        # in case another dihedral matching the wildcard is found,
                        # but it is not directly after the previous dihedral
                        if found==True and d!=prev+1:
                            break
                        if check_wildcard_dihedral( ff['dihedraltypesX'][d][0], [a1,a2,a3,a4] )==True:
                            params = ff['dihedraltypesX'][d][1]
                            dihXparams.append(params)
                            found = True
                            prev = d
                        elif check_wildcard_dihedral( ff['dihedraltypesX'][d][0], [a4,a3,a2,a1] )==True:
                            params = ff['dihedraltypesX'][d][1]
                            dihXparams.append(params)
                            found = True
                            prev = d

                    # remove the atomtypes from the ff string
                    for params in dihXparams:
                        paramsplit = params[0].split()
                        params = ""
                        for p in paramsplit:
                            if p not in [a1,a2,a3,a4]:
                                if p!='X':
                                    params += " "+str(p)
                        # create new topology line
                        newline = ('{a1:<5}{a2:<5}{a3:<5}{a4:<5}'
                                   '{params}'.format(a1=el[0], a2=el[1],
                                                     a3=el[2], a4=el[3],
                                                     params=params))
                        out.append(newline)


                # check whether at least one torsional term was found
                if found==False:
                    print('\n    ERROR! Could not find dihedral parameters:\n %s' % " ".join(key))
#                              ' following atoms are already present in the '
#                              'force field object'
#                              ' and will be ignored:\n    %s' % key)
                    sys.exit(0)
                #assert found is True

            # -----------------------------
            # dihedrals without multiplicity
            # -----------------------------
            elif int(dtype) in [2]:
                found = False
                el = line.split()
                a1 = a_map[el[0]]
                a2 = a_map[el[1]]
                a3 = a_map[el[2]]
                a4 = a_map[el[3]]
                dtype = el[4]

                # order does not matter: check both
                try:
                    key = (a1, a2, a3, a4, dtype)
                    params = ff['dihedraltypes'][key]
                    found = True
                except:
                    try:
                        key = (a4, a3, a2, a1, dtype)
                        params = ff['dihedraltypes'][key]
                        found = True
                    except:
                        found = False

                if found==True:
                    # remove the atomtypes from the ff string
                    paramsplit = params.split()
                    params = ""
                    for p in paramsplit:
                        if p not in [a1,a2,a3,a4]:
                            params += " "+str(p)
                    # create new topology line
                    newline = ('{a1:<5}{a2:<5}{a3:<5}{a4:<5}'
                               '{params}'.format(a1=el[0], a2=el[1],
                                                 a3=el[2], a4=el[3],
                                                 params=params))
                    out.append(newline)

                ############################################################
                # if not yet found a dihedral, search through the wildcards #
                # the search depends on the dihedral number #
                # the first match will be the correct dihedral #
                if found==False:
                    dihXparams = ''
                    for d in ff['dihedraltypesX']:
                        # first check that the dihedral is of the correct type
                        if int(ff['dihedraltypesX'][d][0][4]) not in [2]:
                            continue
                        if check_wildcard_dihedral( ff['dihedraltypesX'][d][0], [a1,a2,a3,a4] )==True:
                            dihXparams = ff['dihedraltypesX'][d][1]
                            found = True
                            break
                        elif check_wildcard_dihedral( ff['dihedraltypesX'][d][0], [a4,a3,a2,a1] )==True:
                            dihXparams = ff['dihedraltypesX'][d][1]
                            found = True
                            break

                    # remove the atomtypes from the ff string
                    paramsplit = dihXparams[0].split()
#                    print paramsplit
#                    sys.exit(0)
                    params = ""
                    for p in paramsplit:
                        if p not in [a1,a2,a3,a4]:
                            if p!='X':
                                params += " "+str(p)
                    # create new topology line
                    newline = ('{a1:<5}{a2:<5}{a3:<5}{a4:<5}'
                               '{params}'.format(a1=el[0], a2=el[1],
                                                 a3=el[2], a4=el[3],
                                                 params=params))
                    out.append(newline)

            # ----------------------
            # catch unkown dihedrals
            # ----------------------
            else:
                exit('\nBUG: dihedral type %s not expected/checked/'
                     'implemented yet in func "fill_params".'
                     ' Script needs to be extended.' % dtype)

        # -----------------------------------------------------------
        # Just copy the other lines, e.g. the pairs, moleculetype etc.
        # -----------------------------------------------------------
        else:
            out.append(line)

    return out


# ==============================================================================
# MAIN
# ==============================================================================
def main(args):
    print('\n Reading force field information...')

    if args.priority == 'ff':
        if 'itp' in args.ff:
            ffparams = ForceField(fname='{path}'.format(path=args.ff), bForgiving=args.bForgiving)
            ffparams.read_file(fname='{path}'.format(path=args.ff),append=True)
        else:
            ffparams = ForceField(fname='{path}/ffnonbonded.itp'.format(path=args.ff), bForgiving=args.bForgiving)
            ffparams.read_file(fname='{path}/ffbonded.itp'.format(path=args.ff),append=True)
        ffparams.read_file(fname='{prm}'.format(prm=args.prm),append=True)
    elif args.priority == 'prm':
        ffparams = ForceField(fname='{prm}'.format(prm=args.prm), bForgiving=args.bForgiving)
        ffparams.read_file(fname='{path}/ffnonbonded.itp'.format(path=args.ff),
                              append=True)
        ffparams.read_file(fname='{path}/ffbonded.itp'.format(path=args.ff),
                           append=True)

    print('\n Filling in the topology file...')
    # map atom ids to their atomtypes
    id2type_map = get_atoms_map(itp=args.itp)
    print(id2type_map)
    # put all ff parameters explicitly in the itp file
    newitp = fill_params(itp=args.itp, a_map=id2type_map, ff=ffparams.ff)
    # save the new itp
    list2file(lst=newitp, outfname=args.outfname)
    print(' New topology file saved as "{fname}"\n'.
          format(fname=args.outfname))

# call main
if __name__ == "__main__":
    args = parse_options()
    main(args)
