#############################################################################
#
# Author: Ruth HUEY, Michel F. SANNER
#
# Copyright: M. Sanner TSRI 2000
#
#############################################################################


# $Header: /mnt/raid/services/cvs/python/packages/share1.5/AutoDockTools/energyCalculator.py,v 1.5.10.1 2016/02/11 09:24:08 annao Exp $
#
# $Id: energyCalculator.py,v 1.5.10.1 2016/02/11 09:24:08 annao Exp $
#
#

"""
This class calculates the AutoDock internal energy for a set of atoms.
"""

import numpy
from math import sqrt
from AutoDockTools.energyConstants import Rij, epsij
from MolKit.molecule import AtomSet


        
class EnergyCalculator:
    """This class knows how to calculate the energy of its set of atoms
    the Autodock way.  

    self.data: the list of atoms for the calculation
    self.reset: if 1, recalculates all the energies each time
    self.lowLim: number to be flag that calculation hasn't been done
    self.energy_type: internal, some day we may have others...
    
    """
    def __init__(self, data, reset=0, lowLim=-999999.0, energy_type='internal'):
        """data is a set of atoms
        """
        self.data = data
        self.data_ids = []
        for a in data:
            self.data_ids.append(id(a))
        #have to make sure that atom.number is in order
        self.data.sort()
        # save the pair-wise energies for reuse by get_energy
        # code: -1 if not set, 0 if within 1-4 distance, otherwise pairwise
        # energy
        self.energy_matrix = numpy.zeros([len(data), len(data)]) + lowLim
        self.lowLim = lowLim
        self.setRotatable()
        self.setNeighbors()
        self.weedbonds()
        self.reset = reset
        self.NBCUTOFF = 8.0


    def setRotatable(self):
        #find the bonds which are in the torsion Map and label rotatable
        mols = self.data.top.uniq()
        for m in mols:
            if not hasattr(m, 'torTree'):
                continue
            l = []
            m.allAtoms.bonds[0].rotatable = 0
            for n in m.torTree.torsionMap:
                ats = m.allAtoms.get(lambda x, n=n: x.number-1 in n.bond)
                #ats = AtomSet([n.a,n.b])
                b = ats.bonds[0][0]
                b.rotatable = 1


    def weedbonds(self):
        mols = self.data.top.uniq()
        for m in mols:
            if not hasattr(m, 'torTree'):
                continue
            l = []
            atL = m.torTree.rootNode.atomList
            #print 'setting 0 for rootNode atomList:', atL
            for i in atL:
                for j in atL:
                    self.energy_matrix[i][j] = 0
            for n in m.torTree.torsionMap:
                atL = n.atomList
                atL.extend(n.bond)
                #print 'setting 0 for atomList:', atL
                for i in atL:
                    for j in atL:
                        self.energy_matrix[i][j] = 0


    def setNeighbors(self):
        """
        detect 1-2, 1-3 and 1-4 neighbors and set their entries in
        self.energy_matrix to 0
        FIX THIS: if bond is rotatable, don't reset the 1-4 pair
        """
        
        self.nb = {}
        for a1 in self.data:
            ind1 = self.data.index(a1)
            nbd =  {}
            nbd[ind1] = 0
            #nblist = self.nb[ind1] = [ind1]
            self.energy_matrix[ind1, ind1] = 0
            #set 1-2 interactions to 0
            for b in a1.bonds:
                a2 = b.atom1
                if id(a2)==id(a1): a2 = b.atom2
                if a2.number > a1.number:
                    #then do something here
                    ind2 = self.data.index(a2)
                    nbd[ind2] = 0
                    self.energy_matrix[ind1, ind2] = 0
                    self.energy_matrix[ind2, ind1] = 0
                #set 1-3 interactions to 0
                for b2 in a2.bonds:
                    a3 = b2.atom1
                    if id(a3)==id(a2): a3 = b2.atom2
                    if id(a3)==id(a1): continue
                    if a3.number > a1.number:
                        #then do something here
                        ind3 = self.data.index(a3)
                        nbd[ind3] = 0
                        self.energy_matrix[ind1, ind3] = 0
                        self.energy_matrix[ind3, ind1] = 0
                    #set 1-4 interactions to 0
                    #if rotatable, skip 1-4s for b2
                    if b2.rotatable:
                        #print 'skipping ', b2.atom1.name, '-', b2.atom2.name
                        continue
                    for b3 in a3.bonds:
                        #if hasattr(b3, 'activeTors') and b3.activeTors: continue
                        a4 = b3.atom1
                        if id(a4)==id(a3): a4 = b3.atom2
                        if id(a4)==id(a2): continue
                        if id(a4)==id(a1): continue
                        if a4.number > a1.number:
                            #then do something here
                            ind4 = self.data.index(a4)
                            nbd[ind4] = 0
                            self.energy_matrix[ind1, ind4] = 0
                            self.energy_matrix[ind4, ind1] = 0
            self.nb[ind1] = nbd.keys()


    def dist(self, a, b):
        """return distance between two atoms, a and b.
        """
        d = numpy.array(b.coords) - numpy.array(a.coords)
        return sqrt(numpy.sum(d*d))

                        

    #def _get_energy_custom(self, a, b):
    #    """return energy between two atoms, a and b.
    #    """
    #    ax = self.data.index(a)
    #    bx = self.data.index(b)
#
#        if self.energy_matrix[ax][bx] >= 0.0:
#            # return previously saved energy
#            return self.energy_matrix[ax][bx]
#        else:
#            # compute, save, and return energy
#            #FIX THIS
#            dist = a.getRMSD_custom(b.getCoords())
#            self.energy_matrix[ax][bx] = self.energy_matrix[bx][ax] = dist
#            return dist


    def get_energy(self, a, b):
        """return energy between two atoms, a and b.
        """
        ax = self.data.index(a)
        bx = self.data.index(b)

        e = self.energy_matrix[ax][bx]
        if e > self.lowLim and not self.reset:
            # return previously saved energy
            return e
        elif e==0:
            # if getNeighbors set this to 0
            return e
        else:
            # compute, save, and return  energy
            energy = self.nbenergy(a,b)
            self.energy_matrix[ax][bx] = self.energy_matrix[bx][ax] = energy
            return energy


    def nbenergy(self, a, b):
        r = self.dist(a,b)
        if r>=self.NBCUTOFF:
            #print 'skipping ', a.number, '-', b.number
            return 0.0
        key = 'lj'+ a.element + b.element
        r_eq = Rij[key]
        ep = epsij[key]
        n = 12
        m = 6
        energy = (m/(n-m))*ep*r_eq**n/(r**n) - (n/(n-m))*ep*r_eq**m/r**m
        #energy = m*epsilon*r_eq**n/(r**n)/(n-m) - n*epsilon*r_eq**m/r**m/(n-m)
        return  energy


    def get_internal_energy(self):
        intEnergy = 0
        datalen = len(self.data)
        for i in range(datalen):
            a = self.data[i]
            for j in range(i+1, datalen):
                b = self.data[j]
                eij = self.get_energy(a,b)
                intEnergy = intEnergy + eij
        return intEnergy
                
    
    def get_subset_internal_energy(self, ats):
        subset_intEnergy = 0
        for a1 in ats:
            if id(a1) not in self.data_ids:
                print 'warning', a1.name, ' not in self.data'
            for a2 in ats:
                if id(a2) not in self.data_ids:
                    print 'warning', a2.name, ' not in self.data'
                eij = self.get_energy(a1, a2)
                subset_intEnergy = subset_intEnergy + eij
        return subset_intEnergy


    def reset_matrix(self):
        self.energy_matrix = numpy.zeros([len(self.data), len(self.data)]) + self.lowLim
        self.weedbonds()
        self.setNeighbors()



    def print_nb_matrix(self):
        n = len(self.data)
        s = 'Atom: ID:'
        for i in range(1,n+1):
            s = s + '%2d'%i
        print s
        s = '_'*n
        print s
        em = self.energy_matrix
        for i in range(n):
            at = self.data[i]
            s = '  %2s     %2d'%(at.name,at.number)
            for j in range(n):
                if em[i][j]!=0:
                    s = s + '|x'
                else:
                    s = s + '|_'
            print s


    def print_half_nb_matrix(self):
        n = len(self.data)
        em = self.energy_matrix
        for i in range(n):
            s = '' + i*'_'
            for j in range(i, n):
                if em[i][j]!=0:
                    s = s + 'x'
                else:
                    s = s + '_'
            print s

