#
#
#  $Id: test_PyAutoDock.py,v 1.6 2007/09/27 19:05:18 annao Exp $ 
#  $Header: /opt/cvs/python/packages/share1.5/AutoDockTools/Tests/test_PyAutoDock.py,v 1.6 2007/09/27 19:05:18 annao Exp $
#
#


import unittest, glob, os
import numpy.oldnumeric as Numeric
from string import find, split
from MolKit import molecule, protein
#molecule.bhtreeFlag = 0
#protein.bhtreeFlag = 0

mv = None
ct = 0
totalCt = 10
class PyAutoDock_Tests(unittest.TestCase):
    """
    setUp + tearDown form a fixture: working environment for the testing code
    """

    lines = [
    "at#  elem       estat       hbond         vdw         desolv",
    "  1    N      -0.0099      -0.0001      -0.2204       0.0000",
    "  2    C      -0.0481       0.0000      -0.2896      -0.0976",
    "  3    C      -0.0540       0.0000      -0.2754      -0.0983",
    "  4    C      -0.0385       0.0000      -0.3023      -0.1049",
    "  5    O      -0.0026      -0.0014      -0.2562       0.2360",
    "  6    N       0.0535      -0.0008      -0.2344       0.0000",
    "  7    N       0.1733      -0.0001      -0.2173       0.0000",
    "  8    C      -0.0742       0.0000      -0.3414      -0.1175",
    "  9    C      -0.0517       0.0000      -0.3405      -0.1182",
    " 10    C      -0.1279       0.0000      -0.4233      -0.1189",
    " 11    C      -0.1292       0.0000      -0.4071      -0.1187",
    " 12    C      -0.0325       0.0000      -0.3249      -0.1129",
    " 13    C      -0.0429       0.0000      -0.3252      -0.1121",
    " 14    C      -0.0125       0.0000      -0.3426      -0.1163",
    " 15    C      -0.0638       0.0000      -0.2889      -0.1015",
    " 16    O       0.0047      -0.0029      -0.2253       0.2360",
    " 17    N       0.1447      -0.0009      -0.2360       0.0000",
    " 18    C      -0.0283       0.0000      -0.3450      -0.1054",
    " 19    C      -0.0295       0.0000      -0.4043      -0.1014",
    " 20    C      -0.0122       0.0000      -0.4316      -0.1092",
    " 21    A       0.0028       0.0000      -0.4923      -0.0194",
    " 22    A      -0.0010       0.0000      -0.5748      -0.0230",
    " 23    A      -0.0002       0.0000      -0.5267      -0.0254",
    " 24    A      -0.0001       0.0000      -0.4778      -0.0237",
    " 25    A      -0.0011       0.0000      -0.3777      -0.0203",
    " 26    A       0.0021       0.0000      -0.3725      -0.0180",
    " 27    H      -0.0246      -0.0033      -0.0199       0.1180",
    " 28    H      -0.1724      -0.2910      -0.0330       0.1180",
    " 29    O       0.7385      -0.0006      -0.4405       0.2360",
    " 30    H      -0.5705      -0.4923      -0.0176       0.1180",
    " 31    A       0.0077       0.0000      -0.3121      -0.0168",
    " 32    A      -0.0012       0.0000      -0.3468      -0.0177",
    " 33    A      -0.0001       0.0000      -0.3358      -0.0160",
    " 34    A       0.0001       0.0000      -0.3722      -0.0165",
    " 35    A       0.0000       0.0000      -0.2975      -0.0134",
    " 36    A       0.0000       0.0000      -0.3137      -0.0139",
    " 37    C      -0.0033       0.0000      -0.3657      -0.1035",
    " 38    A      -0.0013       0.0000      -0.2973      -0.0132",
    " 39    A       0.0462       0.0000      -0.3079      -0.0125",
    " 40    A       0.0007       0.0000      -0.2723      -0.0126",
    " 41    N      -0.2328      -0.3911      -0.1800       0.0000",
    " 42    A       0.0460       0.0000      -0.1824      -0.0076",
    " 43    A       0.0032       0.0000      -0.1860      -0.0090",
    " 44    C      -0.0017       0.0000      -0.3774      -0.1296",
    " 45    C      -0.0039       0.0000      -0.4594      -0.1376",
    " 46    C      -0.0036       0.0000      -0.4036      -0.1370",
    " 47    C       0.0001       0.0000      -0.4630      -0.1401",
    " 48    O       0.0935      -0.2694      -0.4128       0.2360",
    " 49    H      -0.0474      -0.4145      -0.0474       0.1180",
    ]

    estat_weight = 0.1146 # Autogrid3 weight
    dsolv_weight = 0.1711 # Autogrid3 weight
    hbond_weight = 0.0656 # Autogrid3 weight
    vdw_weight = 0.1485 # Autogrid3 weight
    #data_file = open('evaluator_results')
    #lines = data_file.readlines()
    #data_file.close()
    estat_data = Numeric.array([float(split(l)[2]) for l in lines[1:]])
    hbond_data = Numeric.array([float(split(l)[3]) for l in lines[1:]])
    vdw_data   = Numeric.array([float(split(l)[4]) for l in lines[1:]])
    dsolv_data = Numeric.array([float(split(l)[5]) for l in lines[1:]])
    ad305_data= estat_data+hbond_data+vdw_data+dsolv_data
        
    def startViewer(self):
        global mv
        #print 'in test_AD startViewer'
        if mv is None:
            from MolKit import Read
            import Tkinter
            from Pmv.moleculeViewer import MoleculeViewer
            mv = MoleculeViewer(withShell=False)
            #mv = MoleculeViewer(trapExceptions = False)
            mv.loadModule('pyAutoDockCommands', 'AutoDockTools')
            mv.loadModule('deleteCommands', 'Pmv')
            mv.loadModule('selectionCommands', 'Pmv')
            #change warningMsg format
            mv.setUserPreference(('warningMsgFormat', 'printed'), log = 0)
        self.mv = mv


    def setUp(self):
        """
        clean-up
        """
        if not hasattr(self, 'mv'):
            self.startViewer()
        if "pytest_hsg1" not in self.mv.Mols.name:
            mv.readMolecule("pytest_hsg1.pdbqs")
        if "docked" not in self.mv.Mols.name:
            mv.readMolecule("docked.pdbqs")
        self.ligAts = mv.Mols.get('docked')[0].allAtoms
        self.lenLig = len(self.ligAts)


    def tearDown(self):
        """
        clean-up
        """
        global ct, totalCt
        ct = ct + 1
        #print 'ct =', ct
        if ct==totalCt:
            print 'destroying mv'
            self.mv.Exit(0)
            del self.mv


    def test_VDW_pyAutoDock_vs_evaluator(self):
        """ compare AD3 vdw energies in ADT with stored values from PyAutoDock"""
        self.mv.PyAD_calcVDWEnergies("pytest_hsg1:::", "docked:::", self.vdw_weight)
        for i, a in zip(range(self.lenLig), self.ligAts):
            self.assertAlmostEqual(a.vdw_energy, self.vdw_data[i], 2)


    def test_HB_pyAutoDock_vs_evaluator(self):
        """ compare AD3 hbond energies in ADT with stored values from PyAutoDock"""
        self.mv.PyAD_calcHBONDEnergies("pytest_hsg1:::", "docked:::", self.hbond_weight)
        for i, a in zip(range(self.lenLig), self.ligAts):
            self.assertAlmostEqual(a.hb_energy, self.hbond_data[i], 1)


    def test_DSOLV_pyAutoDock_vs_evaluator(self):
        """ compare AD3 dsolv energies in ADT with stored values from PyAutoDock"""
        self.mv.PyAD_calcDSOLVEnergies("pytest_hsg1:::", "docked:::", self.dsolv_weight)
        for i, a in zip(range(self.lenLig), self.ligAts):
            self.assertAlmostEqual(a.dsolv_energy, self.dsolv_data[i], 3)


    def test_ESTAT_pyAutoDock_vs_evaluator(self):
        """ compare AD3 estat energies in ADT with stored values from PyAutoDock"""
        self.mv.PyAD_calcESTATEnergies("pytest_hsg1:::", "docked:::", self.estat_weight)
        for i, a in zip(range(self.lenLig), self.ligAts):
            #print i, ':', a.name, ':', a.estat_energy, ':', self.estat_data[i]
            self.assertAlmostEqual(a.estat_energy, self.estat_data[i], 2)


    def test_AD305_pyAutoDock_vs_evaluator(self):
        """ compare sum of AD305 energies computed in ADT with stored values from PyAutoDock
        """
        self.mv.PyAD_calcAD3Energies("pytest_hsg1:::", "docked:::")
        max_err = -20
        max_pErr = -20
        for i, a in zip(range(self.lenLig), self.ligAts):
            #apparently oxygen 29 required this slop in value comparison
            #print i, ':', a.name, ':', a.ad305_energy, ':', self.ad305_data[i]
            self.assertAlmostEqual(a.ad305_energy, self.ad305_data[i], 1)
            err = abs(a.ad305_energy-self.ad305_data[i])
            if err>max_err:
                #print "resetting err for ", a.name, ':',  err
                max_err = err
            self.assertEqual(err < .02, True)
            percent_err = abs(err/a.ad305_energy)*100.
            if percent_err>max_pErr:
                #print "reseting for ", a.name, ':', percent_err
                max_pErr= percent_err
            self.assertEqual(percent_err < 10, True)
        #print "max_pErr=", max_pErr
        #print "max_err=", max_err
            


class PyAutoDock4_Tests(unittest.TestCase):
    """
    setUp + tearDown form a fixture: working environment for the testing code
    """
    lines = [
        "at#  at_name   estat     hbond       vdw    dsolvLig  dsolvRec",
        "  1    N1      -0.0153    0.0000   -0.2341   -0.0270    0.0375",
        "  2    C1      -0.0697    0.0000   -0.3043    0.0873    0.0664",
        "  3    C2      -0.0816    0.0000   -0.2893    0.0343    0.0704",
        "  4    C3      -0.0697    0.0000   -0.3176    0.0588    0.0610",
        "  5    O1       0.0352   -0.0021   -0.2691    0.0224    0.0289",
        "  6    N2       0.0946    0.0000   -0.2501    0.1109    0.0424",
        "  7    N3       0.2671    0.0000   -0.2303    0.1297    0.0459",
        "  8    C8      -0.1093    0.0000   -0.3587    0.0103    0.0620",
        "  9    C9      -0.0816    0.0000   -0.3584    0.0935    0.0486",
        " 10    C10     -0.1956    0.0000   -0.4447    0.0110    0.0874",
        " 11    C11     -0.1998    0.0000   -0.4278    0.0034    0.0894",
        " 12    C12     -0.0505    0.0000   -0.3413   -0.0461    0.0703",
        " 13    C13     -0.0633    0.0000   -0.3416   -0.0257    0.0683",
        " 14    C14     -0.0206    0.0000   -0.3599   -0.0449    0.0539",
        " 15    C21     -0.0999    0.0000   -0.3035    0.0465    0.0712",
        " 16    O3       0.0451   -0.0044   -0.2367    0.0225    0.0313",
        " 17    N4       0.1862    0.0000   -0.2527    0.1056    0.0595",
        " 18    C22     -0.0450    0.0000   -0.3624    0.0088    0.0883",
        " 19    C23     -0.0558    0.0000   -0.4247    0.0096    0.1116",
        " 20    C24     -0.0265    0.0000   -0.4534   -0.0316    0.1099",
        " 21    A25      0.0101    0.0000   -0.5171   -0.0003    0.0970",
        " 22    A26     -0.0019    0.0000   -0.6038   -0.0297    0.0902",
        " 23    A27     -0.0002    0.0000   -0.5533   -0.0385    0.0794",
        " 24    A28     -0.0002    0.0000   -0.5019   -0.0360    0.0648",
        " 25    A29     -0.0024    0.0000   -0.3967   -0.0250    0.0676",
        " 26    A30      0.0058    0.0000   -0.3913   -0.0122    0.0813",
        " 27    H4      -0.0465   -0.0018   -0.0250    0.1140    0.0000",
        " 28    H32     -0.1729   -0.4870   -0.0441    0.1198    0.0000",
        " 29    O2       1.1325   -0.0008   -0.4627    0.1187    0.0651",
        " 30    H21     -0.8732   -0.4870   -0.0280    0.2014    0.0000",
        " 31    A15      0.0120    0.0000   -0.3279    0.0059    0.0603",
        " 32    A16     -0.0014    0.0000   -0.3644   -0.0235    0.0729",
        " 33    A20     -0.0004    0.0000   -0.3528   -0.0216    0.0618",
        " 34    A17      0.0002    0.0000   -0.3910   -0.0252    0.0812",
        " 35    A18      0.0000    0.0000   -0.3125   -0.0211    0.0785",
        " 36    A19      0.0001    0.0000   -0.3295   -0.0216    0.0676",
        " 37    C31     -0.0039    0.0000   -0.3848    0.0940    0.0546",
        " 38    A32     -0.0014    0.0000   -0.3125   -0.0153    0.0564",
        " 39    A33      0.0451    0.0000   -0.3235    0.0321    0.0628",
        " 40    A36      0.0009    0.0000   -0.2862   -0.0154    0.0519",
        " 41    N5      -0.2140   -0.3009   -0.1891    0.0405    0.0423",
        " 42    A34      0.0419    0.0000   -0.1916    0.0179    0.0514",
        " 43    A35      0.0031    0.0000   -0.1954   -0.0087    0.0434",
        " 44    C4      -0.0037    0.0000   -0.3965   -0.0698    0.0509",
        " 45    C5      -0.0074    0.0000   -0.4826   -0.0636    0.0634",
        " 46    C6      -0.0078    0.0000   -0.4240   -0.0625    0.0440",
        " 47    C7      -0.0026    0.0000   -0.4864   -0.0644    0.0383",
        " 48    O4       0.2186   -0.3993   -0.4336    0.1001    0.0710",
        " 49    H35     -0.1014   -0.4421   -0.0696    0.1722    0.0000",
    ]

    vdw_weight4 = 0.1560 # Autogrid4 weight
    hbond_weight4 = 0.0974 # Autogrid4 weight
    estat_weight4 = 0.1465 # Autogrid4 weight
    dsolv_weight4 = 0.1159 # Autogrid4 weight
    #data_file4 = open('evaluator4_results')
    #lines = data_file4.readlines()
    #data_file4.close()
    estat_data4 = Numeric.array([float(split(l)[2]) for l in lines[1:]])
    hbond_data4 = Numeric.array([float(split(l)[3]) for l in lines[1:]])
    vdw_data4   = Numeric.array([float(split(l)[4]) for l in lines[1:]])
    dsolv_data4 = Numeric.array([(float(split(l)[5])+float(split(l)[6])) for l in lines[1:]])
    ad4_data= estat_data4+hbond_data4+vdw_data4+dsolv_data4
        

    def startViewer(self):
        global mv
        #print 'in test_AD startViewer'
        if mv is None:
            from MolKit import Read
            import Tkinter
            from Pmv.moleculeViewer import MoleculeViewer
            mv = MoleculeViewer(withShell=False)
            #mv = MoleculeViewer(trapExceptions = False)
            mv.loadModule('pyAutoDockCommands', 'AutoDockTools')
            mv.loadModule('deleteCommands', 'Pmv')
            mv.loadModule('selectionCommands', 'Pmv')
            #change warningMsg format
            mv.setUserPreference(('warningMsgFormat', 'printed'), log = 0)
            mv.readMolecule("hsg1.pdbqt")
            mv.readMolecule("docked4.pdbqt")
        self.mv = mv


    def setUp(self):
        """
        clean-up
        """
        if not hasattr(self, 'mv'):
            self.startViewer()
        if "hsg1" not in self.mv.Mols.name:
            mv.readMolecule("hsg1.pdbqt")
        if "docked4" not in self.mv.Mols.name:
            mv.readMolecule("docked4.pdbqt")
        self.ligAts = mv.Mols.get("docked4")[0].allAtoms
        self.lenLig = len(self.ligAts)


    def tearDown(self):
        """
        clean-up
        """
        global ct, totalCt
        ct = ct + 1
        #print 'ct =', ct
        if ct==totalCt:
            print 'destroying mv'
            self.mv.Exit(0)
            del self.mv


    def test_VDW_pyAutoDock4_vs_evaluator4(self):
        """ compare vdw4 energies in ADT with stored values from PyAutoDock"""
        self.mv.PyAD4_calcVDWEnergies("hsg1:::", "docked4:::", self.vdw_weight4)
        max_err = -20
        max_pe = 0.
        max_pe2 = 0.
        for i, a in zip(range(self.lenLig), self.ligAts):
            eval_data = self.vdw_data4[i]
            adt_data = a.ad4_vdw_energy
            #print i, ':', a.name, ':',adt_data, ':', eval_data,
            self.assertAlmostEqual(adt_data, eval_data, 1)
            err = abs(adt_data-eval_data)
            #print "err=", err
            if err>max_err:
                max_err = err
                #print "max_err set to ", max_err, ' for ',  a.name
            self.assertEqual(err < .020, True)
            percent_err = abs(err/eval_data)*100.
            if percent_err>max_pe:
                max_pe=percent_err
            #print "percent_err=", percent_err, ':', abs(adt_data)>.075
            if abs(adt_data)>.06:
                if percent_err>max_pe2:
                    max_pe2=percent_err
                self.assertEqual(percent_err < 3, True)
        #print "max_pe=", max_pe
        #print "max_pe2=", max_pe2


    def test_HB_pyAutoDock4_vs_evaluator4(self):
        """ compare hbond energies4  in ADT with stored values from PyAutoDock"""
        self.mv.PyAD4_calcHBONDEnergies("hsg1:::", "docked4:::", self.hbond_weight4)
        for i, a in zip(range(self.lenLig), self.ligAts):
            adt_energy = a.ad4_hb_energy
            #adt_sum_energy = a.sum_hb_energy
            eval_energy = self.hbond_data4[i]
            #print i, ':', a.name,':', adt_energy, '-', eval_energy
            #self.assertAlmostEqual(a.ad4_hb_energy, self.hbond_data4[i], 1)
            err = abs(adt_energy - eval_energy)
            #print "err=", err
            self.assertEqual(err < .10, True)
            #if a.ad4_hb_energy!=0.0:
            if adt_energy!=0.0:
                percent_err = abs(err/adt_energy)*100.
                #print "percent_err=", percent_err
                if abs(adt_energy)>.075:
                    self.assertEqual(percent_err < 30, True)


    def test_DSOLV_pyAutoDock4_vs_evaluator4(self):
        """ compare dsolv energies in ADT4 with stored values from PyAutoDock"""
        self.mv.PyAD4_calcDSOLVEnergies("hsg1:::", "docked4:::", self.dsolv_weight4)
        for i, a in zip(range(self.lenLig), self.ligAts):
            #print i, ':', a.name, ':', a.ad4_dsolv_energy, ':', self.dsolv_data4[i]
            self.assertAlmostEqual(a.ad4_dsolv_energy, self.dsolv_data4[i], 3)


    def test_ESTAT_pyAutoDock4_vs_evaluator4(self):
        """ compare estat energies in ADT with stored values from PyAutoDock"""
        self.mv.PyAD4_calcESTATEnergies("hsg1:::", "docked4:::", self.estat_weight4)
        max_err = -20
        for i, a in zip(range(self.lenLig), self.ligAts):
            err = abs(a.ad4_estat_energy-self.estat_data4[i])
            #print i, ':', a.name, ':', a.ad4_estat_energy, ':', self.estat_data4[i]
            if err>max_err:
                max_err = err
                ##print "max_err set to ", max_err, ' for ',  a.name
            self.assertAlmostEqual(a.ad4_estat_energy, self.estat_data4[i], 2)
        ##print 'max_err=', max_err


    def test_AD4_pyAutoDock4_vs_evaluator4(self):
        """ compare sum of AD4 energies computed in ADT with stored values from PyAutoDock
        """
        self.mv.PyAD4_calcAD4Energies("hsg1:::", "docked4:::")
        max_err = -20
        max_Perr = -20
        for i, a in zip(range(self.lenLig), self.ligAts):
            #apparently oxygen 29 required this slop in value comparison
            #print i, ':', a.name, ':', a.ad4_energy, ':', self.ad4_data[i]
            #print round(a.ad4_energy,1)==round(self.ad4_data[i],1)
            #self.assertAlmostEqual(a.ad4_energy, self.ad4_data[i], 1)
            err = abs(a.ad4_energy-self.ad4_data[i])
            #print "err=", err
            if err>max_err:
                max_err = err
                #print "reset max_err for ", a.name
            self.assertEqual(err < .3, True)
            percent_err = abs(err/a.ad4_energy)*100.
            if a.ad4_energy>.05:
                if percent_err>max_Perr:
                    max_Perr = percent_err
                    #print ">>> reset max_Perr for ", a.name, " percent_err=", percent_err
                #print 'percent_err=', percent_err
                self.assertEqual(percent_err < 5, True)
        #print "max_err=", max_err
        #print "max_Perr=", max_Perr
            


if __name__ == '__main__':
    test_cases = [
        'PyAutoDock_Tests',
        'PyAutoDock4_Tests',
        #'PyAutoDock_CalcClosestPairs',
            ]
    unittest.main( argv=([__name__ ,] + test_cases) )
    #unittest.main( argv=([__name__ , '-v'] + test_cases) )








        

   





