## Automatically adapted for numpy.oldnumeric Jul 23, 2007 by 

#
#
#  $Id: test_AD4flex.py,v 1.2 2010/05/26 21:45:39 annao Exp $
#
#

import unittest, glob, os, sys
from string import strip, find, split
import time
from MolKit.molecule import Atom

mv = None
ct = 0
totalCt = 6


class AD4flex_BaseTests(unittest.TestCase):
    """
    setUp + tearDown form a fixture: working environment for the testing code
    """
    
        
    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(trapExceptions=False)
            mv.loadModule('autoflex4Commands', '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()
        for m in self.mv.Mols:
            self.mv.deleteMol(m)


    def tearDown(self):
        """
        clean-up
        """
        #print 'in flex tearDown'
        global ct, totalCt
        #reset flexDict object
        self.mv.flexDict = {}
        #delete any molecules left due to errors
        for m in self.mv.Mols:
            self.mv.deleteMol(m)
        ct = ct + 1
        #print 'ct =', ct
        if ct==totalCt:
            print 'destroying mv'
            self.mv.Exit(0)
            del self.mv       
    
        

class AD4flex_readflexTests(AD4flex_BaseTests):
    
    
    def test_flex4_read_macro_mol(self):
        """readflexTests: check that reading a macromol sets expected flexdict items
        """
        filename = '1ebg_rec.pdbqt'
        self.mv.AD4flex_readMacro(macroFile=filename)
        self.assertEqual(self.mv.flexDict['macrofilename'], filename)
        self.assertEqual(self.mv.flexDict['macroname'], '1ebg_rec')
        self.assertEqual(self.mv.flexDict['macromol'], self.mv.Mols[0])
    

    def test_flex4_choose_macro_mol_name(self):
        """readflexTests: check that choosing a macro sets expected flexDict items
        """
        filename = '1ebg_rec.pdbqt'
        self.mv.readMolecule(filename)
        macro_name = self.mv.Mols[0].name
        self.mv.AD4flex_chooseMacro(macro_name)
        self.assertEqual(self.mv.flexDict['macroname'], macro_name)
        self.assertEqual(self.mv.flexDict['macromol'], self.mv.Mols[0])
        self.assertEqual(self.mv.flexDict['macrofilename'], filename)


    #next 2 tests were added after a bug was fixed    
    def test_flex4_read_macro_delete_it(self):
        """readflexTests: check onRemoveObjFromViewer bug is fixed #1:
                AD4flex_readMacro, delete macro and check flexdict items
        """
        filename = '1ebg_rec.pdbqt'
        self.mv.AD4flex_readMacro(macroFile=filename)
        self.assertEqual(self.mv.flexDict['macrofilename'], filename)
        self.assertEqual(self.mv.flexDict['macroname'], '1ebg_rec')
        self.assertEqual(self.mv.flexDict['macromol'], self.mv.Mols[0])
        self.mv.deleteMol("1ebg_rec")
        self.assertEqual(self.mv.flexDict, {})
    

    def test_flex4_read_ind_read_macro_delete_ind(self):
        """readflexTests: check onRemoveObjFromViewer bug is fixed #2:
                read ind,  AD4flex_readMacro, delete ind and check flexdict items
        """
        self.mv.readMolecule('ind.pdbqt')
        filename = '1ebg_rec.pdbqt'
        self.mv.AD4flex_readMacro(macroFile=filename)
        self.assertEqual(self.mv.flexDict['macrofilename'], filename)
        self.assertEqual(self.mv.flexDict['macroname'], '1ebg_rec')
        self.assertEqual(self.mv.flexDict['macromol'], self.mv.Mols[1])
        self.mv.deleteMol("ind")
        self.assertEqual(self.mv.flexDict['macrofilename'], filename)
        self.assertEqual(self.mv.flexDict['macroname'], '1ebg_rec')
        self.assertEqual(self.mv.flexDict['macromol'], self.mv.Mols[0])
    
 

class AD4flex_formatting_flexResTests(AD4flex_BaseTests):

    def test_flex4_formatting_arg8(self):
        """flex4_formatting_arg8: processing both arg8 residues produces expected result files
        """
        self.mv.readMolecule('hsg1.pdbqt', ask=0, parser=None)
        self.mv.AD4flex_chooseMacro("hsg1")
        self.mv.AD4flex_setResidues("::ARG8")
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:ARG8:CA,CB;", 0)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:ARG8:CA,CB;", 0)
        bnds = self.mv.Mols[0].chains.residues.get("ARG8").atoms.bonds[0]
        sum = 0
        possibleCt = bnds.possibleTors.count(1)
        self.assertEqual(possibleCt, 14)
        activeCt = bnds.activeTors.count(1)
        self.assertEqual(activeCt, 6)
        flex_filename = 'temp_FLEX.pdbqt'
        rigid_filename = 'temp_RIGID.pdbqt'
        for fn in [flex_filename, rigid_filename]:
            if os.path.exists(fn):
                cmd = 'rm -f ' + fn
                os.system(cmd)

        self.mv.AD4flex_writeFlexFile(flex_filename)
        self.mv.AD4flex_writeRigidFile(rigid_filename)
        for testfile,reffile in zip([flex_filename, rigid_filename],\
                ['ref_hsg1_flex.pdbqt', 'ref_hsg1_rigid.pdbqt']):
            testptr = open(testfile) 
            testlines = testptr.readlines()
            refptr = open(reffile)
            reflines = refptr.readlines()
            for refline, testline in zip(reflines, testlines):
                self.assertEqual(refline, testline)
        

    def test_flex4_formatting_TRP(self):
        """ flex4_formatting_TRP: processing a single TRP residue gives expected number of rotatable bonds...
        """
        self.mv.readMolecule('hsg1.pdbqt', ask=0, parser=None)
        self.mv.AD4flex_chooseMacro("hsg1")
        #try to format TRP6 in chain A
        self.mv.AD4flex_setResidues(":A:6")
        self.mv.ADflex_setBondRotatableFlag(":A::CA,CB;", 0)
        bnds = self.mv.Mols.chains.residues.get("6").atoms.bonds[0]
        possibleCt = bnds.possibleTors.count(1)
        self.assertEqual(possibleCt, 4)
        activeCt = bnds.activeTors.count(1)
        self.assertEqual(activeCt, 3)



###class AD4flex_hingeTests(AD4flex_BaseTests):


###    def test_flex4_met46_lys55_hinge_log(self):
###        """format hinge met46_lys55: __call__ gives expected hinge
###        """
###        cmd = self.mv.AD4flex_setHinge
###        cmd.selectionBase.set("")
###        cmd.atomOne = None
###        cmd.atomTwo = None
###        cmd.atoms = []
###        self.mv.AD4flex_readMacro('./1hvr_rec.pdbqt', log=0)
###        self.mv.AD4flex_setHinge(("1hvr_rec:A:MET46:CA", "1hvr_rec:A:LYS55:CA"), "1hvr_rec:A:MET46:C;1hvr_rec:A:MET46:O;1hvr_rec:A:MET46:CB;1hvr_rec:A:MET46:CG;1hvr_rec:A:MET46:SD;1hvr_rec:A:MET46:CE;1hvr_rec:A:ILE47:N;1hvr_rec:A:ILE47:CA;1hvr_rec:A:ILE47:C;1hvr_rec:A:ILE47:O;1hvr_rec:A:ILE47:CB;1hvr_rec:A:ILE47:CG1;1hvr_rec:A:ILE47:CG2;1hvr_rec:A:ILE47:CD1;1hvr_rec:A:ILE47:H;1hvr_rec:A:GLY48:N;1hvr_rec:A:GLY48:CA;1hvr_rec:A:GLY48:C;1hvr_rec:A:GLY48:O;1hvr_rec:A:GLY48:H;1hvr_rec:A:GLY49:N;1hvr_rec:A:GLY49:CA;1hvr_rec:A:GLY49:C;1hvr_rec:A:GLY49:O;1hvr_rec:A:GLY49:H;1hvr_rec:A:ILE50:N;1hvr_rec:A:ILE50:CA;1hvr_rec:A:ILE50:C;1hvr_rec:A:ILE50:O;1hvr_rec:A:ILE50:CB;1hvr_rec:A:ILE50:CG1;1hvr_rec:A:ILE50:CG2;1hvr_rec:A:ILE50:CD1;1hvr_rec:A:ILE50:H;1hvr_rec:A:GLY51:N;1hvr_rec:A:GLY51:CA;1hvr_rec:A:GLY51:C;1hvr_rec:A:GLY51:O;1hvr_rec:A:GLY51:H;1hvr_rec:A:GLY52:N;1hvr_rec:A:GLY52:CA;1hvr_rec:A:GLY52:C;1hvr_rec:A:GLY52:O;1hvr_rec:A:GLY52:H;1hvr_rec:A:PHE53:N;1hvr_rec:A:PHE53:CA;1hvr_rec:A:PHE53:C;1hvr_rec:A:PHE53:O;1hvr_rec:A:PHE53:CB;1hvr_rec:A:PHE53:CG;1hvr_rec:A:PHE53:CD1;1hvr_rec:A:PHE53:CD2;1hvr_rec:A:PHE53:CE1;1hvr_rec:A:PHE53:CE2;1hvr_rec:A:PHE53:CZ;1hvr_rec:A:PHE53:H;1hvr_rec:A:ILE54:N;1hvr_rec:A:ILE54:CA;1hvr_rec:A:ILE54:C;1hvr_rec:A:ILE54:O;1hvr_rec:A:ILE54:CB;1hvr_rec:A:ILE54:CG1;1hvr_rec:A:ILE54:CG2;1hvr_rec:A:ILE54:CD1;1hvr_rec:A:ILE54:H;1hvr_rec:A:LYS55:N;1hvr_rec:A:LYS55:H;")
###        hinge_list = self.mv.flexDict['hinge_list']
###        self.assertEqual(len(hinge_list), 1)
###        hinge = hinge_list[0]
###        self.assertEqual(len(hinge), 2)
###        (atomOne, atomTwo), atoms = hinge
###        self.assertEqual(atomOne.full_name(), '1hvr_rec:A:MET46:CA')
###        self.assertEqual(atomTwo.full_name(), '1hvr_rec:A:LYS55:CA')
###        self.assertEqual(len(atoms), 67)


###    def test_flex4_formatting_met46_lys55_hinge_with_saved_sets(self):
###        """format hinge met46_lys55: formatting with saved_sets gives expected hinge
###        """
###        cmd = self.mv.AD4flex_setHinge
###        cmd.selectionBase.set("")
###        cmd.atomOne = None
###        cmd.atomTwo = None
###        cmd.atoms = []
###        self.mv.readMolecule('1hvr_rec.pdbqt', ask=0, parser=None)
###        self.mv.saveSet("1hvr_rec:A:MET46,LYS55:CA", 'hinge1')
###        self.mv.saveSet("1hvr_rec:A:ILE54:N;1hvr_rec:A:ILE54:CA;1hvr_rec:A:ILE54:C;1hvr_rec:A:ILE54:O;1hvr_rec:A:ILE54:CB;1hvr_rec:A:ILE54:CG1;1hvr_rec:A:ILE54:CG2;1hvr_rec:A:ILE54:CD1;1hvr_rec:A:ILE54:H;1hvr_rec:A:PHE53:N;1hvr_rec:A:PHE53:CA;1hvr_rec:A:PHE53:C;1hvr_rec:A:PHE53:O;1hvr_rec:A:PHE53:CB;1hvr_rec:A:PHE53:CG;1hvr_rec:A:PHE53:CD1;1hvr_rec:A:PHE53:CD2;1hvr_rec:A:PHE53:CE1;1hvr_rec:A:PHE53:CE2;1hvr_rec:A:PHE53:CZ;1hvr_rec:A:PHE53:H;1hvr_rec:A:GLY52:N;1hvr_rec:A:GLY52:CA;1hvr_rec:A:GLY52:C;1hvr_rec:A:GLY52:O;1hvr_rec:A:GLY52:H;1hvr_rec:A:GLY51:N;1hvr_rec:A:GLY51:CA;1hvr_rec:A:GLY51:C;1hvr_rec:A:GLY51:O;1hvr_rec:A:GLY51:H;1hvr_rec:A:ILE50:N;1hvr_rec:A:ILE50:CA;1hvr_rec:A:ILE50:C;1hvr_rec:A:ILE50:O;1hvr_rec:A:ILE50:CB;1hvr_rec:A:ILE50:CG1;1hvr_rec:A:ILE50:CG2;1hvr_rec:A:ILE50:CD1;1hvr_rec:A:ILE50:H;1hvr_rec:A:GLY49:N;1hvr_rec:A:GLY49:CA;1hvr_rec:A:GLY49:C;1hvr_rec:A:GLY49:O;1hvr_rec:A:GLY49:H;1hvr_rec:A:GLY48:N;1hvr_rec:A:GLY48:CA;1hvr_rec:A:GLY48:C;1hvr_rec:A:GLY48:O;1hvr_rec:A:GLY48:H;1hvr_rec:A:ILE47:N;1hvr_rec:A:ILE47:CA;1hvr_rec:A:ILE47:C;1hvr_rec:A:ILE47:O;1hvr_rec:A:ILE47:CB;1hvr_rec:A:ILE47:CG1;1hvr_rec:A:ILE47:CG2;1hvr_rec:A:ILE47:CD1;1hvr_rec:A:ILE47:H;1hvr_rec:A:LYS55:N;1hvr_rec:A:LYS55:H;1hvr_rec:A:MET46:C;1hvr_rec:A:MET46:O", 'h1', log=0, comments='No description')
###        self.mv.AD4flex_setHinge(self.mv.sets['hinge1'],self.mv.sets['h1']) 
###        hinge_list = self.mv.flexDict['hinge_list']
###        self.assertEqual(len(hinge_list), 1)
###        hinge = hinge_list[0]
###        self.assertEqual(len(hinge), 2)
###        (atomOne, atomTwo), atoms = hinge
###        self.assertEqual(atomOne.full_name(), '1hvr_rec:A:MET46:CA')
###        self.assertEqual(atomTwo.full_name(), '1hvr_rec:A:LYS55:CA')
###        self.assertEqual(len(atoms), 63)



###    def test_flex4_formatting_met46_lys55_hinge_via_gui(self):
###        """hinge met46_lys55: formatting with saved_sets gives expected hinge 
###        """
###        cmd = self.mv.AD4flex_setHinge
###        cmd.selectionBase.set("")
###        cmd.atomOne = None
###        cmd.atomTwo = None
###        cmd.atoms = []
###        self.mv.readMolecule('1hvr_rec.pdbqt', ask=0, parser=None)
###        #select the hinge atoms
###        self.mv.select("1hvr_rec:A:MET46,LYS55:CA")
###        cmd = self.mv.AD4flex_setHinge
###        cmd.selectHingeMethod.set('cursel')
###        cmd.updateHinge()
###        cmd.selectionBase.set('between')
###        cmd.updateBase()
###        cmd.setHinge_cb()
###        hinge_list = self.mv.flexDict['hinge_list']
###        self.assertEqual(len(hinge_list), 1)
###        hinge = hinge_list[0]
###        self.assertEqual(len(hinge), 2)
###        (atomOne, atomTwo), atoms = hinge
###        self.assertEqual(atomOne.full_name(), '1hvr_rec:A:MET46:CA')
###        self.assertEqual(atomTwo.full_name(), '1hvr_rec:A:LYS55:CA')
###        self.assertEqual(len(atoms), 67)



if __name__ == '__main__':
    test_cases = [
        'AD4flex_readflexTests',
        'AD4flex_formatting_flexResTests',
        #'AD4flex_hingeTests',
            ]
    unittest.main( argv=([__name__ ,] + test_cases) )
    #unittest.main( argv=([__name__ ,'-v'] + test_cases) )
    #unittest.main()



