#  Copyright (C) 2012-2014,2016  Hannes H Loeffler, Julien Michel
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program 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 General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
#  For full details of the license please see the COPYING file
#  that should have come with this distribution.

r"""
Actual functionality and utility functions for the Morph class.
"""


__revision__ = "$Id: util.py 654 2017-11-08 12:06:35Z jmichel $"


import sys
import os
import re
import glob
import math
import itertools
import cPickle as pickle
from collections import OrderedDict, defaultdict

import openbabel as ob

import rdkit.Chem
from rdkit.Chem import ChiralType, rdMolAlign
from rdkit import rdBase

import numpy as np

import Sire.Mol
import Sire.MM
import Sire.Units
import Sire.Maths

# parmed 2.4.0 from AMBER16
from parmed.amber.mask import AmberMask
from parmed.amber.readparm import AmberParm
import parmed.tools.actions as Action

from FESetup import const, errors, logger

from FESetup.munkres import Munkres, print_matrix


class _AtomInfo(object):
    """Simple struct to store Atom info."""

    __slots__ = ['atom', 'index', 'name']

    def __init__(self, atom, index, name):
        self.atom = atom                # Atom: if None then dummy
        self.index = index              # AtomIdx
        self.name = name                # AtomName

    def __str__(self):
        return '%s/%s/%s' % (self.atom, self.index, self.name)


def write_mol2(molecule, outmol2 = '', notypes = False, zz_atoms = [],
               resname = const.LIGAND_NAME, rnum = False):
    """Write a coordinate file in MOL2 format.

    :param molecule: the Sire molecule to be written out
    :type molecule: Sire.Mol
    :param outmol2: output MOL2 file name
    :type outmol2: string
    :param notypes: replace atom types with element names, used for MCSS
    :type notypes: bool
    :param zz_atoms: rename atoms in list to 'zz' to circumvent leap valency check
    :type zz_atoms: list of Sire.Mol.AtomName
    :param resname: the residue name to be used
    :type resname: string
    :param rnum: use residue numbering?
    :type rnum: bool
    :returns: MOL2 data
    :rtype: string
    """

    # Sybyl 7.1 (Mid-2005) requires only num_atoms in data line two of
    # @<TRIPOS>MOLECULE but Leap expects all five fields (and ignores the last
    # two)
    lines = '''\
@<TRIPOS>MOLECULE
%s
%s %s 1 0 0
SMALL
USER_CHARGES
****
generated by FESetup
@<TRIPOS>ATOM
''' % (resname, molecule.nAtoms(),
       molecule.property('connectivity').nConnections())

    for i, atom in enumerate(molecule.atoms() ):
        name = '%s' % atom.name().value()
        coords = atom.property('coordinates')
        charge = atom.property('charge').value()

        if name in zz_atoms:
            atype = 'zz'
        elif notypes:                   # RDKit can't deal with GAFF types
            atype = atom.property('element').symbol()
        else:
            atype = atom.property('ambertype')

        # FXIME: finish this
        if rnum:
            resnum = atom.residue().number().value()
            resname = atom.residue().name().value()
        else:
            resnum = 1

        # assume only 1 substructure present
        lines += '%6d %-4s %10.4f %10.4f %10.4f %-5s %i %s %9.6f\n' % (
            i + 1, name, coords[0], coords[1], coords[2], atype, resnum,
            resname, charge)

    lines += '@<TRIPOS>BOND\n'
    bonds = molecule.property('connectivity').getBonds()

    for i, bond in enumerate(bonds):
        at0 = bond.atom0().value() + 1
        at1 = bond.atom1().value() + 1

        lines += '%6d %6d %6d 1\n' % (i, at0, at1)

    # Sybyl 7.1 (Mid-2005) requires only subst_id, subst_name and root_atom
    # but leap also reads subst_type without actually using it
    lines += '@<TRIPOS>SUBSTRUCTURE\n'

    if rnum:
        for i, res in enumerate(molecule.residues() ):
            for atom in res.atoms():
                first = atom
                break
            
            root = first.index().value() + 1
            lines += '  %i %s %i RESIDUE\n' % (i, res.name().value(), root)
    else:
        lines += '  1 %s 1 TEMP\n' % resname

    if outmol2:
        with open(outmol2, 'w') as output:
            output.write(lines)

    return lines


def create_isotope_map(filename):
    """
    Tag atom pairs as "isotopes" (arbitrary number other than 0) as found
    in the filename.  Two column format: two integers interpreted as
    atom indices (starting with 1).

    :param filename: name of file with user tags
    :type filename: string
    :raises: SetupError
    :returns: isotope map
    :rtype: dict
    """

    isotope_map = {}

    if os.access(filename, os.F_OK):
        with open(filename, 'r') as isomap:
            lcnt = 0

            for line in isomap:
                lcnt += 1

                ltmp = line.lstrip()

                if ltmp.startswith('#') or ltmp == '':
                    continue

                try:
                    idx1, idx2 = line.split()[0:2]
                    isotope_map[int(idx1)] = int(idx2)
                except (ValueError, TypeError) as why:
                    raise errors.SetupError('malformed file %s in '
                                        'line %i: %s' % (filename, lcnt, why) )

    return isotope_map


_fmcs_imp = 'c++'                       # 'python' or 'c++'

if _fmcs_imp == 'c++':
    from rdkit.Chem.rdFMCS import FindMCS, AtomCompare, BondCompare

    # RDKit 2015.03.1 FMCS C++ implentation, seems not to be exactly the
    # same implementation e.g. SMARTS string more specific
    # NOTE: some different parameters and order!
    #       matchChiralTag = False not implemented before 2015.03.1
    _params = dict(maximizeBonds = False, threshold = 1.0,
                   verbose = False, matchValences = False,
                   ringMatchesRingOnly = True, completeRingsOnly = True,
                   bondCompare = BondCompare.CompareAny)
else:
    from rdkit.Chem.MCS import FindMCS

    # defaults are minNumAtoms = 2, maximize = 'bonds',
    # atomCompare = 'elements', bondCompare = 'bondtypes',
    # matchValences = False, ringMatchesRingOnly = False,
    # completeRingsOnly = False, timeout = None, threshold = None
    #
    # when completeRingsOnly = True also ringMatchesRingOnly = True
    #
    # if number of atoms in match < minNumAtoms then smarts will be None
    # completeRingsOnly = True disallows partial rings
    # MCS algorithm is exhaustive so use timeout to limit time
    _params = dict(minNumAtoms = 2, maximize = 'atoms', atomCompare = 'elements',
                   bondCompare = 'any', matchValences = False,
                   ringMatchesRingOnly = True, completeRingsOnly = True,
                   threshold = None)

def mcss(mol2str_1, mol2str_2, maxtime=60, isotope_map=None, selec=''):
    """
    Maximum common substructure search via RDKit/fmcs.

    :param mol2str_1: first MOL2 string
    :type mol2str_1: string
    :param mol2str_2: second MOL2 string
    :type mol2str_2: string
    :param maxtime: timeout for fmcs in seconds
    :type maxtime: float
    :param isotope_map: explicit user atom mapping
    :type isotope_map: dict
    :param selec: selection method for multiple MCS
    :type selec: string
    :raises: SetupError
    :returns: index map
    :rtype: dict
    """

    # disable warning about no explicit hydrogens
    rdBase.DisableLog('rdApp.warning')
    mol1 = rdkit.Chem.MolFromMol2Block(mol2str_1, sanitize = False,
                                       removeHs = False)
    mol2 = rdkit.Chem.MolFromMol2Block(mol2str_2, sanitize = False,
                                       removeHs = False)
    rdBase.EnableLog('rdApp.warning')

    _params.update(timeout = int(maxtime) )

    # FIXME: test c++ implementation
    if isotope_map:
        if _fmcs_imp == 'c++':
            _params.update(atomCompare = AtomCompare.CompareIsotopes)
        else:
            _params.update(atomCompare = 'isotopes')

        max_idx1 = mol1.GetNumAtoms()
        max_idx2 = mol2.GetNumAtoms()

        icnt = 0

        # NOTE: would it make sense to have multiple atoms of a molecule tagged
        #       as the same isotope?
        for idx1, idx2 in isotope_map.iteritems():
            icnt += 1

            logger.write('Mapping atom index %i to %i' % (idx1, idx2) )

            if idx2 < 1:
                atom1 = mol1.GetAtomWithIdx(idx1-1)
                atom1.SetIsotope(-1 * idx2 + icnt)
            else:
                if idx1 > max_idx1 or idx2 > max_idx2 or idx1 < 0 or idx2 < 0:
                    logger.write('Error: indices out of bounds (%i, %i)' %
                                 (max_idx1, max_idx2) )
                    raise errors.SetupError('Mapping indices out of bounds '
                                            '(%i, %i)' % (max_idx1, max_idx2) )

                # FIXME: guard against non-existing indices
                atom1 = mol1.GetAtomWithIdx(idx1-1)
                atom1.SetIsotope(icnt)

                atom2 = mol2.GetAtomWithIdx(idx2-1)
                atom2.SetIsotope(icnt)
    else:
        if _fmcs_imp == 'c++':
            _params.update(atomCompare = AtomCompare.CompareAny)
        else:
            _params.update(atomCompare = 'any')

        n_chiral1 = len(rdkit.Chem.FindMolChiralCenters(mol1) )
        n_chiral2 = len(rdkit.Chem.FindMolChiralCenters(mol2) )

        if n_chiral1 > 0:
            logger.write('Warning: state 0 has %i chiral center%s. Check if '
                         'configurations are inverted!'
                         % (n_chiral1, 's' if n_chiral1 > 1 else '') )
        if n_chiral2 > 0:
            logger.write('Warning: state 1 has %i chiral center%s. Check if '
                         'configurations are inverted!'
                         % (n_chiral2, 's' if n_chiral2 > 1 else '') )


    mcs = FindMCS( (mol1, mol2), **_params)

    if _fmcs_imp == 'c++':
        smarts = mcs.smartsString
        completed = not mcs.canceled
    else:
        smarts = mcs.smarts
        completed = mcs.completed

    logger.write('Running RDKit/fmcs (%s implementation) with arguments:\n%s' %
                 (_fmcs_imp,
                  ', '.join(['%s=%s' % (k,v) for k,v in _params.iteritems()] ) ) )

    if not smarts:
        raise errors.SetupError('No MCSS match could be found')

    if not completed:
        logger.write('Warning: MCSS timed out after %.2fs' % maxtime)

    p = rdkit.Chem.MolFromSmarts(smarts)

    conv = ob.OBConversion()
    conv.SetInAndOutFormats('mol2', 'mol2')

    # NOTE: this relies on a modified Openbabel MOL2 writer
    conv.AddOption('r', ob.OBConversion.OUTOPTIONS)  # do not append resnum

    obmol1 = ob.OBMol()

    errlev = ob.obErrorLog.GetOutputLevel()
    ob.obErrorLog.SetOutputLevel(0)

    conv.ReadString(obmol1, mol2str_1)

    ob.obErrorLog.SetOutputLevel(errlev)

    # NOTE: experimental!
    if selec == 'spatially-closest':
        m1 = mol1.GetSubstructMatches(p, uniquify=False, maxMatches=100, useChirality=False)
        m2 = mol2.GetSubstructMatches(p, uniquify=False, maxMatches=100, useChirality=False)

        logger.write('Applying spatially-closest algorithm (%s, %s matches)\n' %
                     (len(m1), len(m2) ) )


        # FIXME: is it possible that the smaller one has more then one matches
        #        when uniquify=True?
        if len(m1) < len(m2):
            m1, m2 = m2, m1
            conf1 = mol2.GetConformer()
            conf2 = mol1.GetConformer()
            swapped = True
        else:
            conf1 = mol1.GetConformer()
            conf2 = mol2.GetConformer()
            swapped = False

        # for x in range(m1)...
        #     match1 = m1[x]
        #     for y in range(m2)...
        #         match2 = m2[y]
        #         find sumd
        #         save x, y with smallest sumd
        mind = 999999.0
        minxy = [-1,-1]
        #DEVTHRESHOLD = 0.2**2# That's 0.2 Angstrom
        for x in range(len(m1)):
            match1 = m1[x]
            for y in range(len(m2)):
                match2 = m2[y]
                sumd = 0.0
                for i, idx1 in enumerate(match1):
                    pos1 = conf1.GetAtomPosition(idx1)
                    idx2 = match2[i]
                    pos2 = conf2.GetAtomPosition(idx2)
                    d2 = (pos1.x - pos2.x)**2 + (pos1.y - pos2.y)**2 +\
                         (pos1.z - pos2.z)**2
                    sumd += d2
                #print (x,y,sumd,mind)
                if sumd < mind:
                    mind = sumd
                    minxy = [x,y]
        if swapped:
            mapping = dict(zip(m2[minxy[1]], m1[minxy[0]]))
        else:
            # JM bug 11/17 ? 
            #mapping = dict(zip(m1[minxy[1]], m2[minxy[0]]))
             mapping = dict(zip(m1[minxy[0]], m2[minxy[1]]))
        #print (mapping)

        #dist_sum = []
        #for match1 in m1:
        #    dists = []
        #
        #    for i, idx1 in enumerate(match1):
        #        pos1 = conf1.GetAtomPosition(idx1)
        #        for match2 in m2:
        #            idx2 = match2[i]
        #            pos2 = conf2.GetAtomPosition(idx2)
        #
        #            d = math.sqrt((pos1.x - pos2.x)**2 + (pos1.y - pos2.y)**2 +
        #                          (pos1.z - pos2.z)**2)
        #            #print pos1.x, pos2.x
        #            #print idx1, idx2, d
        #            sumd += d
        #            dists.append(d)
        #    dist_sum.append(sum(dists))
        #    print sum(dists), m1
        #min_idx = dist_sum.index(min(dist_sum) )
        #print (min_idx)
        ## FIXME: assume there is only one match for 2nd molecule
        #if swapped:
        #    mapping = dict(zip(m2[0], m1[min_idx]) )
        #else:
        #    mapping = dict(zip(m1[min_idx], m2[0]) )
        #
        #m1 = m1[min_idx]
        #m2 = m2[0]
    else:
        m1 = mol1.GetSubstructMatch(p)
        m2 = mol2.GetSubstructMatch(p)

        mapping = dict(zip(m1, m2) )

    # FIXME: we may have to reconsider this and understand when rings have
    #        to be assumed "broken"
    #
    # delete atoms from mapping that are also part of an map-external ring
    if False:
        ring_info1 = mol1.GetRingInfo()
        ring_info2 = mol2.GetRingInfo()

        rings1 = ring_info1.AtomRings()
        rings2 = ring_info2.AtomRings()

        map1 = set(mapping.keys())
        map2 = set(mapping.values())

        for ring in rings1:
            if not set(ring).issubset(map1):
                for idx in map1:
                    if idx in ring and ring_info1.NumAtomRings(idx) == 1:
                        del(mapping[idx])

        delete_values = []
        for ring in rings2:
            if not set(ring).issubset(map2):
                for idx in map2:
                    if idx in ring and ring_info2.NumAtomRings(idx) == 1:
                        delete_values.append(idx)

        mapping = {k: v for k, v in mapping.items() if v not in delete_values}


    delete_atoms = []

    for atom in ob.OBMolAtomIter(obmol1):
        idx = atom.GetIdx() - 1

        if idx not in mapping:
            delete_atoms.append(atom)

    obmol1.BeginModify()

    for idx in delete_atoms:
        obmol1.DeleteAtom(idx)

    obmol1.EndModify()

    conv.WriteFile(obmol1, const.MCS_MOL_FILE)

    with open(const.MCS_MAP_FILE, 'wb') as pkl:
        pickle.dump(mapping.keys(), pkl, 0)
        pickle.dump(mapping.values(), pkl, 0)

    return mapping


def split_system(mols):
    """Create new Sire molecule with first residue (=ligand) deleted.

    :param mols: system to be split
    :type mols: Sire.Mol.Molecules
    :returns: ligand, rest of system
    :rtype: Sire.Mol.Molecule, Sire.Mol.MoleculeGroup
    """

    moln = mols.molNums()
    moln.sort()

    lig = mols.at(moln[0]).molecule()

    rest = Sire.Mol.Molecules(mols)
    rest.remove(moln[0])

    return lig, rest


def map_atoms(lig_initial, lig_final, timeout, isotope_map = None,
              mcs_sel = ''):
    """
    Compute the atom mapping between initial and final state using MCSS.
    Creates lig_morph, appends to atom_map and reverse_atom_map.

    atom_map contains the AtomIdx (including newly created dummies where
    necessary) of state0 mapping to state1, reverse_atom_map does the opposite.
    lig_morph contains the Atom corresponding to the AtomName as in atom_map.

    :param lig_initial: the initial state molecule
    :type lig_initial: Sire.Mol.Molecule
    :param lig_final: the final state molecule
    :type lig_final: Sire.Mol.Molecule
    :param timeout: fmcs timeout in seconds
    :type timeout: float
    :param isotope_map: explicit user atom mapping
    :type isotope_map: dict
    :raises: SetupError
    :returns: morph molecule, forward map, reverse map
    :rtype: Sire.Mol.CutGroup, OrderedDict of Sire.Mol.AtomName to
      Sire.Mol.AtomName, OrderedDict of Sire.Mol.AtomName to Sire.Mol.AtomName
    """

    # make all atoms carbons to ensure consideration of hydrogens in MCSS
    # backend code (RDKit) is not able to deal with GAFF types anyway
    mol1 = write_mol2(lig_initial, notypes = True)
    mol2 = write_mol2(lig_final, notypes = True)

    # JM 10/16
    if (mcs_sel == 'shapealign' and len(isotope_map) > 0):
        logger.write("user provided mappings override shape align mode")
    elif mcs_sel == 'shapealign':
        logger.write("Will map atoms using shape align mode")
        # Use O3Align to maximise shape overlay
        m1 = rdkit.Chem.MolFromMol2Block(mol1, removeHs=False)
        m2 = rdkit.Chem.MolFromMol2Block(mol2, removeHs=False)
        pyO3A = rdMolAlign.GetO3A(m2, m1)
        score = pyO3A.Align()
        #print (score)
        #rdkit.Chem.MolToPDBFile(m1,"test1.pdb")
        #rdkit.Chem.MolToPDBFile(m2,"test2.pdb")
        # JM 10/16 cannot find atomatic types from supplied Mol2 ! 
        #rdkit.Chem.SetAromaticity(m1)
        #rdkit.Chem.SetAromaticity(m2)
        #print (rdkit.Chem.MolToSmiles(m1))
        #print (rdkit.Chem.MolToSmiles(m2))
        #import pdb; pdb.set_trace()
        #print (os.getcwd())
        # add to isotope map pairs of closest atoms
        # create 2D distance matrix
        n1 = m1.GetNumAtoms()
        n2 = m2.GetNumAtoms()
        # FIXME do we need n1 >= n2 ? 
        m1conf = m1.GetConformer()
        m2conf = m2.GetConformer()
        dists = np.array( np.zeros([n1,n2]), dtype=np.float64 )
        for i in range (0,n1):
            ipos = m1conf.GetAtomPosition(i)
            for j in range(0,n2):
                jpos = m2conf.GetAtomPosition(j)
                dij2 = (jpos.x - ipos.x)**2 + (jpos.y - ipos.y)**2 +\
                       (jpos.z - ipos.z)**2
                #print (i,j,dij2)
                #import pdb ; pdb.set_trace()
                dists[i][j] = dij2
                #dists[j][i] = dij2
        #print (dists)
        #Kuhn-Munkres algorithm for useful solving the rectangular Assignment Problem
        mu = Munkres()
        indexes = mu.compute(dists.tolist())
        for (i,j) in indexes:
            isotope_map[i+1] = j+1
        #import pdb ; pdb.set_trace()
        #tomap = []
        #for i in range(0,n1):
        #    tomap.append(i)
        #mapped = {}
        #mapped_dist = {}
        #for j in range(0,n2):
        #    mapped[j] = -1
        #    mapped_dist[j] = -1
        #it = 0
        #MAXIT = 1000
        #while (len(tomap) > 0):
        #    it += 1
        #    if (it > MAXIT):
        #        logger.write("Exceeed MAXIT to map atoms in shapealign mode. This sounds like a bug.")
        #        break
        #    #print (len(tomap))
        #    i = tomap.pop()
        #    j = dists[i].argmin()
        #    dij = dists[i].min()
        #    if mapped[j] < 0:
        #        mapped[j] = i
        #        mapped_dist[j] = dij
        #    else:
        #        if (dij < mapped_dist[j]):
        #            tomap.append(mapped[j])
        #            mapped[j] = i
        #            mapped_dist[j] = dij
        #    #import pdb ; pdb.set_trace()
        #    #sys.exit(-1)
        ##print (mapped)
        ## if spatially closest atom is within X ANG of next closest atom
        ## do not add to map (ambiguous) ?
        #for j,i in mapped.items():
        #    if (i < 0):
        #        continue
        #    isotope_map[i+1] = j+1
        #import pdb ; pdb.set_trace()
        #sys.exit(-1)
    #print (indexes)
    #print (isotope_map)
    #import pdb ; pdb.set_trace()
    #sys.exit(-1)
    index_map = mcss(mol1, mol2, timeout, isotope_map, mcs_sel)

    if not index_map:
        raise errors.SetupError('MCSS error')


    # NOTE: lig_morph = Sire.Mol.Molecule(lig_initial) would create a new
    #       molecule including all properties, but in adding new atoms below
    #       these properties would need to be created first, see parm_conn()

    lig_morph = Sire.Mol.Molecule('morph').edit()  # MolEditor
    lig_morph = lig_morph.renumber()    # MolEditor

    for res in lig_initial.residues():
        lig_morph = lig_morph.add(res.number() ) # ResStructureEditor
        lig_morph = lig_morph.rename(res.name() ) # ResStructureEditor
        lig_morph = lig_morph.molecule()  # MolStructureEditor

    lig_morph = lig_morph.add(Sire.Mol.CGName('0') ) # CGStructureEditor

    atoms_final = lig_final.atoms()     # Selector_Atom_
    dummy_count = lig_final.nAtoms()

    atom_map = OrderedDict()
    reverse_atom_map = OrderedDict()

    # create a morph molecule using the atoms in the initial molecule
    for atom_i in lig_initial.atoms():
        iname = atom_i.name()
        iidx = atom_i.index()
        iinfo = _AtomInfo(atom_i, iidx, iname)

        resnum = atom_i.residue().number()
 
        lig_morph = lig_morph.add(iname)  # AtomStructureEditor
        lig_morph = lig_morph.reparent(resnum) # AtomStructureEditor
        lig_morph = lig_morph.cutGroup()  # CGStructureEditor

        i = atom_i.index()

        try:
            jidx = Sire.Mol.AtomIdx(index_map[iidx.value()] )
            atom_j = atoms_final.select(jidx)
            jname = atom_j.name()
        except KeyError:
            # NOTE: Sire indexes are 0 based so increment dummy_count aftewards
            jidx = Sire.Mol.AtomIdx(dummy_count)

            dummy_count += 1

            atom_j = None
            jname = Sire.Mol.AtomName('DU%s' % dummy_count)

        jinfo = _AtomInfo(atom_j, jidx, jname)

        atom_map[iinfo] = jinfo
        reverse_atom_map[jinfo] = iinfo


    # atoms not mapped onto atoms of the final molecule are dummies
    dummy_count = lig_initial.nAtoms()

    for atom_f in atoms_final:
        jidx = atom_f.index()

        if not search_by_index(jidx, reverse_atom_map):
            dc = dummy_count
            dummy_count += 1
            name = 'DU%s' % dummy_count
            dummy_name = Sire.Mol.AtomName(name)

            # FIXME: use resnum of the residue the dummy is bound to
            lig_morph = lig_morph.add(dummy_name) # AtomStructureEditor
            lig_morph = lig_morph.reparent(resnum) # AtomStructureEditor
            lig_morph = lig_morph.cutGroup() # CGStructureEditor

            iinfo = _AtomInfo(None, Sire.Mol.AtomIdx(dc), dummy_name)
            jinfo = _AtomInfo(atom_f, jidx, atom_f.name())

            atom_map[iinfo] = jinfo
            reverse_atom_map[jinfo] = iinfo

    # molecule has no properties yet, see parm_conn()
    lig_morph = lig_morph.molecule().commit()

    return lig_morph, atom_map, reverse_atom_map


def search_by_index(query, atom_map):
    """Helper function to search dict of _AtomInfo for a key.
    This is necessary because differently created AtomIdx with the same name
    do not have the same hash, i.e. idx1 == idx2 is True but
    id(idx1) == id(idx2) is False.

    :param query: query key
    :type query: Sire.Mol.AtomIdx
    :param atom_map: the atom map to be searched
    :type atom_map: dict of _AtomInfo to _AtomInfo
    :returns: atom name
    :rtype: Sire.Mol.AtomIdx or None
    """

    for key, value in atom_map.items():
        if query == key.index:
            return value.index

    return None

def search_atom(query, atom_map):
    """Helper function to search _AtomInfo for a key.

    :param query: query key, index because Atom apparently not equal
    :type query: Sire.Mol.AtomIdx
    :param atom_map: the atom map to be searched
    :type atom_map: dict of _AtomInfo to _AtomInfo
    :returns: atom name
    :rtype: Sire.Mol.Atom or None
    """

    for key, value in atom_map.items():
        if query == key.index:
            return value.atom

    return None

def search_atominfo(query, atom_map):
    """Helper function to search _AtomInfo for a key.

    :param query: query key, index because Atom apparently not equal
    :type query: Sire.Mol.AtomIdx
    :param atom_map: the atom map to be searched
    :type atom_map: dict of _AtomInfo to _AtomInfo
    :returns: atom name
    :rtype: _AtomInfo or None
    """

    for key, value in atom_map.items():
        if query == key.index:
            return value

    return None


def parm_conn(lig_morph, atoms_initial, lig_initial, lig_final, atom_map,
              reverse_atom_map):
    """
    Assign initial non-bonded force field parameters and atom types/elements
    for all atoms including the dummy atoms.  Also assign coordinates as they
    are needed for computing of the connectivities.

    :param lig_morph: the morph molecule
    :type lig_morph: Sire.Mol.Molecule
    :param atoms_initial: the set of initial atoms
    :type atoms_initial: Sire.Mol.Selector_Atom
    :param lig_initial: the initial state molecule
    :type lig_initial: Sire.Mol.Molecule
    :param lig_final: the final state molecule
    :type lig_final: Sire.Mol.Molecule
    :param atom_map: the forward atom map
    :type atom_map: dict of _AtomInfo to _AtomInfo
    :param reverse_atom_map: the reverse atom map
    :type reverse_atom_map: dict of _AtomInfo to _AtomInfo
    :raises: SetupError
    :returns: morph molecule, morph connectivity, final state connectivit
    :rtype: Sire.Mol.Molecule, Sire.Mol.Connectivity, Sire.Mol.Connectivity
    """

    ### assign non-bonded force field parameters, also need coordinates
    ### for connectivity calculation below
    lig_morph = lig_morph.edit()
    atom_num = 0

    for info in atom_map:
        if not info.atom:
            element = Sire.Mol.Element('DU')
            coordinates = Sire.Maths.Vector(0, 0, 0)
            charge = 0.0 * Sire.Units.mod_electron
            mass = 0.0
            lj = Sire.MM.LJParameter(0.0 * Sire.Units.angstrom,
                                     0.0 * Sire.Units.kcal_per_mol)
            ambertype = const.DUMMY_TYPE
            new = lig_morph.atom(info.index) # AtomEditor
        else:
            try:
                base = atoms_initial.select(info.index)
            except UserWarning as error:     # could be "anything"...
                raise errors.SetupError('%s not found in reference: %s'
                                        % (info.name, error) )

            element = Sire.Mol.Element(info.name.value())
            coordinates = base.property('coordinates')
            charge = base.property('charge')
            mass = base.property('mass').value()
            lj = base.property('LJ')
            ambertype = '%s' % base.property('ambertype')
            new = lig_morph.atom(info.index) # AtomEditor

        # full setup for initial state otherwise Sire.Mol.Connectivity below
        # will fail (initial_* are not required)
        new.setProperty('element', element)
        new.setProperty('initial_element', element)
        new.setProperty('coordinates', coordinates)
        new.setProperty('charge', charge)
        new.setProperty('initial_charge', charge)
        new.setProperty('initial_mass', mass)
        new.setProperty('LJ', lj)
        new.setProperty('initial_LJ', lj)
        new.setProperty('ambertype', ambertype)
        new.setProperty('initial_ambertype', ambertype)

        lig_morph = new.molecule()

    lig_morph = lig_morph.commit()


    ### create connectivity for dummy atoms
    con_morph = Sire.Mol.Connectivity(lig_morph).edit()

    # We remove all bonds since there is no reason to expect the covalent
    # bond hunter worked correctly for the dummy atoms
    con_morph.disconnectAll()

    con_initial = lig_initial.property('connectivity')
    con_final = lig_final.property('connectivity')

    for i, f in atom_map.iteritems():
        if not i.atom:
            bonded_indices = []

            # no need to check f because dummy i must map to real atom f
            for bonded_idx in con_final.connectionsTo(f.index):
                reversed_idx = search_by_index(bonded_idx, reverse_atom_map)

                if reversed_idx is None:
                    raise errors.SetupError('reversed name')

                bonded_indices.append(reversed_idx)
        else:
            bonded_indices = con_initial.connectionsTo(i.index)

        for bonded_idx in bonded_indices:
            con_morph.connect(i.index, bonded_idx)

    con_morph = con_morph.commit()
    lig_morph = lig_morph.edit() \
        .setProperty('connectivity', con_morph) \
        .commit()

    # NOTE: bonded parameters are assigned in make_pert_file() in pertfile.py

    return lig_morph, con_morph, con_final


def dummy_coords(lig_morph, con_morph, atoms_initial, lig_initial, lig_final,
                 atom_map, reverse_atom_map, connect_final, zz_atoms,
                 dummy_atoms):
    """
    Create coordinates for dummy atoms from initial state.
    Prepares the initial state for make_pert_file().

    :param lig_morph: the morph molecule
    :type lig_morph: Sire.Mol.Molecule
    :param con_morph: the connectivity of the morph
    :type con_morph: Sire.Mol.Connectivity
    :param atoms_initial: set of initial atoms
    :type atoms_initial: Sire.Mol.Selector_Atom
    :param lig_initial: the initial state molecule
    :type lig_initial: Sire.Mol.Molecule
    :param lig_final: the final state molecule
    :type lig_final: Sire.Mol.Molecule
    :param atom_map: the forward atom map
    :type atom_map: dict of _AtomInfo to _AtomInfo
    :param reverse_atom_map: the reverse atom map
    :type reverse_atom_map: dict of _AtomInfo to _AtomInfo
    :param con_final: the connectivity of the final state
    :type con_final: Sire.Mol.Connectivity
    :param zz_atoms: rename atoms in list to 'zz' to circumvent leap valency check
    :type zz_atoms: list of Sire.Mol.AtomName
    :param dummy_atoms: the atom indexes of the dummy atoms
    :type dummy_atoms: list of Sire.Mol.AtomIdx
    :raises: SetupError
    :return: morph molecule, initial state molecule, final state molecule,
       zz map
    :rtype: Sire.Mol.Molecule, Sire.Mol.Molecule, Sire.Mol.Molecule,
       list of Sire.Mol.AtomName
    """

    # Fix coordinates of dummies: needs to be done for every coord set
    # For each dummy, use connectivity to find bond/angle/dihedral
    # Select dihedral that does not involve other dummies
    # If not possible, do first dummies with 0 other dummies, then dummies with
    # 1 other dummy, then with 2, then with 3

    if dummy_atoms:
        logger.write('Computing coordinates of added dummy atoms')
        dummies = dummy_atoms[:]
        toprocess = len(dummies)
    else:
        toprocess = 0

    while toprocess:
        match = None

        # iterate over all dummies and ensure that e.g. in C1-C2-C3-DU1-DU2
        # DU1 gets coordinates before DU2
        for dummy in dummies:
            dihs = con_morph.getDihedrals(dummy)

            for dih in dihs:
                if ( (dih.atom1() not in dummies) and
                     (dih.atom2() not in dummies) and
                     (dih.atom3() not in dummies) ):
                    at1 = lig_morph.select(dih.atom1() )  # Atom
                    at2 = lig_morph.select(dih.atom2() )
                    at3 = lig_morph.select(dih.atom3() )

                    # True:  at1.index() == atom_map.keys()[n]
                    # False: id(at1.index()) == id(atom_map.keys()[n])
                    at1f = search_atominfo(at1.index(), atom_map)
                    at2f = search_atominfo(at2.index(), atom_map)
                    at3f = search_atominfo(at3.index(), atom_map)

                    # dummies in final state? try next dihedral
                    if not at1f.atom or not at2f.atom or not at3f.atom:
                        continue

                    at1_index = at1f.index
                    at2_index = at2f.index
                    at3_index = at3f.index

                    names = filter(lambda n: n,
                                   (at1_index, at2_index, at3_index) )

                    if len(names) < 3:
                        raise errors.SetupError('BUG: %s not found in atom map'
                                                % ', '.join(at1_index,
                                                            at2_index,
                                                            at3_index) )

                    if (at1 and at2 and at3):
                        match = dih
                        break

            if match:
                break

        if not match:
            # Could happen in systems with less than 4 atoms !
            raise errors.SetupError('Current setup code does not support '
                                    'molecules without dihedrals.')

        at0 = lig_morph.select(dih.atom0() )
        at0_index = search_by_index(at0.index(), atom_map)

        if not at0_index:
            raise errors.SetupError('BUG: %s not found in atom map' % at0_index)

        logger.write('This dummy is in the initial torsion:\n\t%s %s %s %s' %
                     (at0.name(), at1.name(), at2.name(), at3.name() ) )

        # NOTE: tleap will not create bonds if the internal maximum bond valence
        # is exceeded (H=1, O=4, all others=8, absolute maximum is 8 even if
        # check switched off).  The solution taken here is to change the atom
        # type of the hydrogen bonded to a dummy atom to 'zz'.  However, we will
        # have to set the correct forcefield parameters, for any dof involving
        # 'zz' in the pert file. The parm files will have the correct bonding
        # but atom types are wrong!

        at1n = at1.name().value()

        if at1n not in zz_atoms:
            ambertype = '%s' % at1.property('ambertype')

            # Is this a sufficient test for hydrogens?
            if ambertype.upper().startswith('H') and at1n.startsWith('H'):
                zz_atoms.append(at1n)

        # Now find what the bond length/angle/distance was in the final
        # topology for this dummy atom
        try:
            at0f = lig_final.select(at0_index)
            at1f = lig_final.select(at1_index)
            at2f = lig_final.select(at2_index)
            at3f = lig_final.select(at3_index)
        except StopIteration as error:
            error_type = re.search(const.RE_SIRE_ERROR_STR,
                                   str(error)).group(1)

            # FIXME: reorder final state such that there are no dummies at the
            #        beginning
            if (error_type == 'SireError::invalid_index'):
                print (lig_final.nAtoms(), at0_index, at1_index, at2_index,
                       at3_index)
                raise errors.SetupError(
                    'BUG: index out of range: lig_morph has more atoms than '
                    'final state.')
            else:
                raise error

        logger.write('Final state equivalents:\n\t%s %s %s %s\n' %
                     (str(at0f), str(at1f), str(at2f), str(at3f) ) )


        # A problem here is that because the does not follow 3D geometry, some
        # bonded atoms could still end up too close to the dummy.  How to pick
        # the "best" one ?
        # --> We check if at1 has 1 atom that is bonded to that is not at0f or
        # at2f and also not in dummies. If so, we use that atom to make an
        # improper
        bonded_at1f = connect_final.getBonds(at1f.index() )
        alternates3f = []

        for altbond in bonded_at1f:
            altf = lig_final.select(altbond.atom1() )

            if (altf != at0f and altf != at2f):
                alti = lig_morph.select(search_by_index(altf.index(),
                                                       reverse_atom_map) )

                if alti.index() not in dummies:
                    alternates3f.append(altf)
                    logger.write(altf)

        # More than one alternative...There is a risk that we put the dummy
        # on top of the second atom
        if len(alternates3f) > 0:
            at3f = alternates3f[0]
            at3 = lig_morph.select(search_by_index(at3f.index(),
                                                  reverse_atom_map) )

        q1 = at1.property('coordinates')
        q2 = at2.property('coordinates')
        q3 = at3.property('coordinates')

        q0f = at0f.property('coordinates')
        q1f = at1f.property('coordinates')
        q2f = at2f.property('coordinates')
        q3f = at3f.property('coordinates')

        bond = Sire.Maths.Vector.distance(q0f, q1f)
        angle = Sire.Maths.Vector.angle(q0f, q1f, q2f)
        dihedral = Sire.Maths.Vector.dihedral(q0f, q1f, q2f, q3f)

        coords = Sire.Maths.Vector.generate(bond, q1, angle, q2, dihedral, q3)

        # What if MORE THAN 2 alternates?
        #
        # Another problem is that we may have multiple equivalent choices. It
        # would be sensible to position the morph atoms such that the RMSD to
        # the equivalent atoms in the final ligand is minimised, otherwise we
        # may actually change the binding mode of the ligand and create sampling
        # problems.

        if len(alternates3f) == 2:
            # Check if we overlapped dummy with alternates[1]
            alt3f1c = lig_morph.select(
                search_by_index(alternates3f[1].index(),
                               reverse_atom_map)).property('coordinates')

            if Sire.Maths.Vector.distance(coords, alt3f1c) < 0.90:
                coords = Sire.Maths.Vector.generate(bond, q1, angle, q2,
                                                    dihedral +
                                                    120 * Sire.Units.degrees,
                                                    q3)

        lig_morph = lig_morph.edit()\
            .atom(dummy).setProperty('coordinates', coords).molecule().commit()

        dummies.remove(dummy)
        toprocess -= 1

    return lig_morph, lig_initial, lig_final, zz_atoms


def _get_dihedrals(dihedrals, idx_list):
    """Helper function to extract proper and improper dihedrals."""

    propers = defaultdict(list)
    impropers = {}

    for dihedral in dihedrals:
        i1 = dihedral.atom1.idx
        i2 = dihedral.atom2.idx
        i3 = dihedral.atom3.idx
        i4 = dihedral.atom4.idx

        if i1 not in idx_list or i2 not in idx_list or i3 not in idx_list or \
               i4 not in idx_list:
            continue

        i1 += 1
        i2 += 1
        i3 += 1
        i4 += 1

        per = dihedral.type.per
        phi_k = dihedral.type.phi_k
        phase = dihedral.type.phase
        scee = dihedral.type.scee
        scnb = dihedral.type.scnb

        if dihedral.signs[1] < 0:
            ordered = tuple(sorted((i1, i2, i3, i4) ) )
            impropers[ordered] = (i1, i2, i3, i4, per, phi_k, phase, scee, scnb)
        else:
            propers[i1, i2, i3, i4].append( [per * dihedral.signs[0], phi_k,
                                             phase, scee, scnb] )

    return propers, impropers


def _add_proper(idx, terms, parm):
    """Helper function to recreate parmtop proper dihedral table."""

    act = Action.deleteDihedral(parm, '@%i @%i @%i @%i' % idx)
    act.execute()

    for term in terms:
        per = abs(term[0])
        phi_k = term[1]
        phase = term[2]
        scee = term[3]
        scnb = term[4]

#       if term[0] < 1:
#           dtype = 'multiterm'
#       else:
#           dtype = 'normal'

        # FIXME: also add scee and scnb from unperturbed state if necessary,
        #        e.g. dummies in carbohydrates have 1.2/2.0 instead of the
        #        1.0/1.0 for the GLYCAM force field
        act = Action.addDihedral(parm, '@%i @%i @%i @%i %f %f %f %f %f '
                               'type normal' %
                               (idx + (phi_k, per, phase, scee, scnb) ) )

        act.execute()


def _add_improper(impropers0, impropers1, parm):
    """Helper function to recreate parmtop improper dihedral table."""

    for k, v in impropers0.iteritems():
        if k not in impropers1:
            i1, i2, i3, i4 = v[:4]
            per, phi_k, phase, scee, scnb = v[4:]

            act = Action.deleteDihedral(parm, '@%i @%i @%i @%i' % (i1, i2, i3, i4) )
            act.execute()

            act = Action.addDihedral(parm, '@%i @%i @%i @%i %f %f %f %f %f '
                                   'type improper' % (i1, i2, i3, i4,
                                                      phi_k, per, phase,
                                                      scee, scnb) )
            act.execute()


def patch_parmtop(parm0_fn, parm1_fn, mask0, mask1, copy_dih=True):
    """
    Patch two AMBER parmtop files to include bond, angle and dihedral functions
    from the respective other state.  These parameters would otherwise be all
    zero because leap/parmchk create those as zero for dummy atoms.  It is
    assumed that no bonds are broken.

    Conceptually this does the same as make_pert_file.  Eventually both should
    be merged into one with probably the format in make_pert_file used as a
    master file for all perturbations.  Also, dihedrals including dummies should
    be allowed to just be zero instead of copying them from the other state.

    :param parm0_fn: filename of the parmtop file for state0
    :type parm0_fn: string
    :param parm1_fn: filename of the parmtop file for state1
    :type parm1_fn: string
    :param mask0: first mask used when one or two parmtops
    :type mask0: string
    :param mask1: second mask only used when one parmtop
    :type mask1: string
    :param copy_dih: copy dihedral from other state if periodicity is zero
    :type copy_dih: bool
    :raises: SetupError
    """

    if parm1_fn != '':
        parm0 = AmberParm(parm0_fn)
        parm1 = AmberParm(parm1_fn)

        if not const.DUMMY_TYPE in parm0.parm_data['AMBER_ATOM_TYPE'] and \
               not const.DUMMY_TYPE in parm1.parm_data['AMBER_ATOM_TYPE']:
            return

        if parm0.ptr('natom') != parm1.ptr('natom'):
            raise errors.SetupError('BUG: parms of different lengths')

        logger.write('Patching parmtops %s, %s with mask %s\n' %
                     (parm0_fn, parm1_fn, mask0) )

        mask_str = mask0
        idx_list = [i for i in AmberMask(parm0, mask_str).Selected()]
        idx_list2 = idx_list

        pmemd = False
    else:
        parm0 = AmberParm(parm0_fn)
        parm1 = parm0

        logger.write('Patching parmtop %s with masks %s, %s\n' %
                     (parm0_fn, mask0, mask1) )

        mask_str0 = mask0
        mask_str1 = mask1
        idx_list =  [i for i in AmberMask(parm0, mask_str0).Selected()]
        idx_list2 = [i for i in AmberMask(parm0, mask_str1).Selected()]

        pmemd = True

    atn0 = parm0.parm_data['ATOMIC_NUMBER']
    atn1 = parm1.parm_data['ATOMIC_NUMBER']

    mass0 = parm0.parm_data['MASS']
    mass1 = parm1.parm_data['MASS']

    for idx, idx2 in zip(idx_list, idx_list2):
        if atn0[idx] < 0:
            atn0[idx] = atn1[idx2]

        if mass0[idx] == 0.0:
            mass0[idx] = mass1[idx2]

        if atn1[idx2] < 0:
            atn1[idx2] = atn0[idx]

        if mass1[idx2] == 0.0:
            mass1[idx2] = mass0[idx]

        offset = idx2 - idx             # FIXME: check if really constant

    for b0 in itertools.chain(parm0.bonds_inc_h, parm0.bonds_without_h):
        idx1 = b0.atom1.idx
        idx2 = b0.atom2.idx

        if idx1 not in idx_list or idx2 not in idx_list:
            continue

        for b1 in itertools.chain(parm1.bonds_inc_h, parm1.bonds_without_h):
            i1 = b1.atom1.idx
            i2 = b1.atom2.idx

            if i1 not in idx_list2 or i2 not in idx_list2:
                continue

            if (idx1+offset, idx2+offset) == (i1, i2) or \
                   (idx1+offset, idx2+offset) == (i2, i1):
                k0 = b0.type.k
                k1 = b1.type.k

                if k0 == k1 == 0.0:
                    raise errors.SetupError('BUG: bonds of both states are '
                                            'zero: %i %i' % (idx1, idx2) )

                if k0 == 0.0:
                    req1 = b1.type.req

                    act = Action.setBond(parm0, '@%i @%i %f %f' %
                                       (idx1+1, idx2+1, k1, req1) )
                    act.execute()

                if k1 == 0.0:
                    req0 = b0.type.req

                    act = Action.setBond(parm1, '@%i @%i %f %f' %
                                       (i1+1, i2+1, k0, req0) )
                    act.execute()

    for a0 in itertools.chain(parm0.angles_inc_h, parm0.angles_without_h):
        idx1 = a0.atom1.idx
        idx2 = a0.atom2.idx
        idx3 = a0.atom3.idx

        if idx1 not in idx_list or idx2 not in idx_list or idx3 not in idx_list:
            continue

        for a1 in itertools.chain(parm1.angles_inc_h, parm1.angles_without_h):
            i1 = a1.atom1.idx
            i2 = a1.atom2.idx
            i3 = a1.atom3.idx

            if i1 not in idx_list2 or i2 not in idx_list2 or i3 not in idx_list2:
                continue

            if (idx1+offset, idx2+offset, idx3+offset) == (i1, i2, i3) or \
                   (idx1+offset, idx2+offset, idx3+offset) == (i3, i2, i1):
                k0 = a0.type.k
                k1 = a1.type.k

                # this may happen when dummies are created for atoms that
                # could otherwise be represented as real atoms, e.g.
                # when breaking rings
                if k0 == k1 == 0.0:
                    pass
                else:
                    if k0 == 0.0:
                        theteq1 = a1.type.theteq

                        act = Action.setAngle(parm0, '@%i @%i @%i %f %f' %
                                        (idx1+1, idx2+1, idx3+1, k1, theteq1) )
                        act.execute()

                    if k1 == 0.0:
                        theteq0 = a0.type.theteq

                        act = Action.setAngle(parm1, '@%i @%i @%i %f %f' %
                                        (i1+1, i2+1, i3+1, k0, theteq0) )
                        act.execute()

    # NOTE: dihedrals can be all zero, can be multiterm in the other state,
    #       dummy propers will have per = 0, end-groups may be excluded in
    #       single term but if multi-term always excluded
    #       impropers may be missing and atoms differently ordered
    dihedrals0 = itertools.chain(parm0.dihedrals_inc_h,
                                 parm0.dihedrals_without_h)
    dihedrals1 = itertools.chain(parm1.dihedrals_inc_h,
                                 parm1.dihedrals_without_h)

    propers0, impropers0 = _get_dihedrals(dihedrals0, idx_list)
    propers1, impropers1 = _get_dihedrals(dihedrals1, idx_list2)

    propers_coll0 = defaultdict(list)
    propers_coll1 = defaultdict(list)

    # NOTE: a given periodicity only occurs once
    for idx0, idx1 in zip(sorted(propers0), sorted(propers1) ):
        terms0 = propers0[idx0]
        terms1 = propers1[idx1]

        if terms0 == terms1:
            continue

        l0 = len(terms0)
        l1 = len(terms1)

        per0s = [abs(per[0]) for per in terms0]
        per1s = [abs(per[0]) for per in terms1]

        if l0 == 1 and terms0[0][0] == 0.0:  # dummy in state0
            if terms1[0][0] == 0.0:
                raise errors.SetupError('BUG: both states have zero dihedrals')

            if copy_dih:
                propers_coll0[idx0] = terms1
                propers_coll1[idx1] = terms1
            else:
                terms0[0][0] = terms1[0][0]
                propers_coll0[idx0] = terms0
                propers_coll1[idx1] = terms1
        elif l1 == 1 and terms1[0][0] == 0.0:  # dummy in state1
            if terms0[0][0] == 0.0:
                raise errors.SetupError('BUG: both states have zero dihedrals')

            if copy_dih:
                propers_coll0[idx0] = terms0
                propers_coll1[idx1] = terms0
            else:
                terms1[0][0] = terms0[0][0]
                propers_coll0[idx0] = terms0
                propers_coll1[idx1] = terms1
        elif per0s == per1s:            # same pers in both states
            propers_coll0[idx0] = terms0
            propers_coll1[idx1] = terms1
        else:                           # different terms
            propers_coll0[idx0] = terms0
            propers_coll1[idx1] = terms1

            for per0 in per0s:
                if per0 not in per1s:
                    propers_coll1[idx1].append( [per0, 0.0, 0.0,
                                                 terms1[0][3], terms1[0][4] ] )

            for per1 in per1s:
                if per1 not in per0s:
                    propers_coll0[idx0].append( [per1, 0.0, 0.0,
                                                 terms0[0][3], terms0[0][4] ] )

            propers_coll0[idx0].sort(key = lambda p: abs(p[0] ) )
            propers_coll1[idx1].sort(key = lambda p: abs(p[0] ) )
 

    # fix multitermS for proper 1,4 pairs
    for terms0, terms1 in zip(propers_coll0.values(), propers_coll1.values() ):
         if len(terms0) > 1:
             for term0 in terms0[:-1]:
                 term0[0] = abs(term0[0])
             terms0[-1][0] = -abs(terms0[-1][0])

         if len(terms1) > 1:
             for term1 in terms1[:-1]:
                 term1[0] = abs(term1[0])
             terms1[-1][0] = -abs(terms1[-1][0])

            
    # reconstruct dihedral table in parmtop because parmed only adds dihedrals
    # FIXME: filter when PK=0.0 on both sides? in parmtop for 1,4 only?
    #        if so filter on GROMACS side
    for idx0, idx1 in zip(sorted(propers_coll0), sorted(propers_coll1) ):
        _add_proper(idx0, propers_coll0[idx0], parm0)
        _add_proper(idx1, propers_coll1[idx1], parm1)

    _add_improper(impropers0, impropers1, parm1)
    _add_improper(impropers1, impropers0, parm0)

    #parm0.overwrite = True
    parm0.write_parm(parm0_fn)

    if not pmemd:
        #parm1.overwrite = True
        parm1.write_parm(parm1_fn)


def transfer_charges(mol0, mol1, atom_map, dum=False):
    """
    Transfer charges from the final state to the initial state.

    :param mol0: molecule 0
    :type mol0: Sire.Mol.CutGroup
    :param mol1: molecule 1
    :type mol1: Sire.Mol.Molecule
    :param atom_map: the forward atom map
    :type atom_map: dict of _AtomInfo to _AtomInfo
    :returns: molecule with new charges
    :rtype: Sire.Mol.Molecule
    """

    fdummies = False

    for finfo in atom_map.values():
        if not finfo.atom:
            fdummies = True
            break

    if fdummies:
        mol_m = Sire.Mol.Molecule(mol0)
    else:
        mol_m = Sire.Mol.Molecule(mol1)

    mol = mol_m.edit()              # MolEditor

    for iinfo in atom_map:
        new = mol.atom(iinfo.index)
        new.setProperty('charge', 0.0 * Sire.Units.mod_electron)
        mol = new.molecule()

    for iinfo, finfo in atom_map.items():
        new = mol.atom(iinfo.index) # AtomEditor

        if not finfo.atom:
            charge = 0.0 * Sire.Units.mod_electron
        else:
            if fdummies:
                # FIXME: is this really correct?
                if dum:
                    idx = iinfo.index
                else:
                    idx = finfo.index

                base = mol1.select(idx)
                charge = base.property('charge')
                new = mol.atom(idx) # AtomEditor
            else:
                if iinfo.atom:
                    base = mol0.select(iinfo.index)
                    charge = base.property('charge')
                else:
                    charge = 0.0 * Sire.Units.mod_electron

        new.setProperty('charge', charge)
 
        mol = new.molecule()        # MolEditor


    return mol.commit()


def zero_charges(mol1, atom_map):
    """
    Set charges to zero.

    :param mol1: molecule to be modified
    :type mol1: Sire.Mol.Molecule
    :param atom_map: the forward atom map
    :type atom_map: dict of _AtomInfo to _AtomInfo
    """

    mol_m = Sire.Mol.Molecule(mol1)
    mol = mol_m.edit()                  # MolEditor

    for finfo in atom_map.values():
        new = mol.atom(finfo.index)     # AtomEditor

        if not finfo.atom:
            new.setProperty('charge', 0.0 * Sire.Units.mod_electron)
            mol = new.molecule()        # MolEditor

    return mol.commit()
