#
#  $Id: test_AD4gpf.py,v 1.13 2012/12/06 23:24:18 annao Exp $
#
#

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

mv = None
ct = 0
totalCt = 39
klass = None


class ADgpf4_BaseTests(unittest.TestCase):
    """
    setUp + tearDown form a fixture: working environment for the testing code
    """
    
        
    def startViewer(self):
        global mv
        if mv is None:
            from MolKit import Read
            import Tkinter
            from Pmv.moleculeViewer import MoleculeViewer
            mv = MoleculeViewer(trapExceptions=False)
            mv.loadModule('autogpfCommands', '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
        """
        global klass, mv
        name = self.__class__.__name__
        if not klass:
            klass = name
        elif klass != name:
            print 'SETUP: destroying mv'
            if mv and mv.hasGui:
                mv.Exit(0)
                if hasattr(self, 'mv'):
                    del self.mv
                mv = None
            klass = name
            
        if not hasattr(self, 'mv'):
            self.startViewer()
        for m in self.mv.Mols:
            self.mv.deleteMol(m)


    def tearDown(self):
        """
        clean-up
        """
        #print 'in gpf tearDown'
        global ct, totalCt
        #reset gridparameter object
        from AutoDockTools.GridParameters import GridParameters
        self.mv.gpo = GridParameters()
        self.mv.gpo.vf = self.mv
        #delete any molecules left due to errors
        for m in self.mv.Mols:
            self.mv.deleteMol(m)
        self.mv.GPF_FLEXRES_TYPES = []
        self.mv.GPF_RECEPTOR_TYPES = []
        self.mv.GPF_LIGAND_TYPES = []

        ct = ct + 1
        #print 'ct =', ct
        if ct==totalCt:
            print 'destroying mv'
            self.mv.Exit(0)
            del self.mv       
    
        

class ADgpf4_readgpfTests(ADgpf4_BaseTests):
    
    
    def test_gpf4_read_macro_mol(self):
        """readgpfTest1: read a macro mol and check gpo receptor_filename 
        """
        filename = '1ebg_rec.pdbqt'
        outputfilename = 'new_' + filename
        self.mv.ADgpf4_readMacromolecule(filename, filename=outputfilename)
        self.assertEqual(self.mv.gpo.receptor_filename, filename)
    

    def test_gpf4_read_empty_macro_mol(self):
        """readgpfTests2: read a empty macro mol and check that IOError is raised
        """
        self.assertRaises(IOError,self.mv.ADgpf4_readMacromolecule," ")

    
    def test_gpf4_read_non_existent_macro_mol(self):
        """readgpfTests3: read a non-existent macro mol and check that IOError is raised
        """
        self.assertRaises(IOError, self.mv.ADgpf4_readMacromolecule,"gfdjs.pdbqt")

       
    def test_gpf4_choose_macro_mol_name(self):
        """readgpfTests4: readgpfTests: check chosen macrol name==gpo.receptor_stem
        """
        self.mv.readMolecule('1ebg_rec.pdbqt')
        macro_name = self.mv.Mols[0].name
        self.mv.ADgpf4_chooseMacromolecule(macro_name)
        self.assertEqual(self.mv.gpo.receptor_stem, macro_name)

    
    def test_gpf4_set_map_types_directly(self):
        """readgpfTests5: check for setting map types directly
        """
        teststring = "A B C D E"
        self.mv.ADgpf4_setMapTypes(teststring)
        self.assertEqual(self.mv.gpo['ligand_types']['value'], teststring)
                
 
    def test_gpf4_choose_ligand_name(self):
        """readgpfTests6: choose ligand then check ligand name
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.readMolecule(filename)
        ligand_name = self.mv.Mols[0].name
        self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
        self.assertEqual(self.mv.gpo.ligand.name, ligand_name)


    def test_gpf4_read_formatted_ligand(self):
        """readgpfTests7: read ligand then check ligand name
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.ADgpf4_readFormattedLigand(filename)
        self.assertEqual(self.mv.gpo.ligand_filename, filename)

    
    def test_gpf4_set_map_types_from_ligand(self):
        """readgpfTests8: check for setting map types from ligandfile
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.ADgpf4_readFormattedLigand(filename)
        stored_values = "P C OA HD N"
        #stored_values = "C HD N OA P"
        self.assertEqual(self.mv.gpo['ligand_types']['value'], stored_values)
                

    def test_gpf4_read_empty_formatted_ligand(self):
        """readgpfTests9: check empty ligand name raises IOError
        """
        self.assertRaises(IOError,self.mv.ADgpf4_readFormattedLigand," ")
        
    
    def test_gpf4_read_non_existent_formatted_ligand(self):
        """readgpfTests10: check non_existent ligand name raises IOError
        """
        self.assertRaises(IOError,self.mv.ADgpf4_readFormattedLigand,"hjgf.pdbq ")
       
         
    
#    def test_gpf4_set_up_covalent_map_coords(self):
#        """ check covalent coords  in set up covalent map
#        """
#        self.mv.ADgpf4_readMacromolecule('1ebg_rec.pdbqt')
#        self.mv.ADgpf4_readFormattedLigand('ind.out.pdbq')
#        #c={}
#        #c['covalent_coords']='-3.884'
#        #apply(self.mv.ADgpf4_setUpCovalentMap,(),c)
#        self.mv.ADgpf4_setUpCovalentMap('-22.9870 -0.0600 4.4430',-1.0,'5.0','1000.0','CAONSHZ')
#        self.assertEqual(str(self.mv.gpo['covalent_coords']['value']),'-22.9870 -0.0600 4.4430')

    
    
    def test_gpf4_write_gpf(self):
        """readgpfTests11: check using ref_1ebg.gpf by comparing ref_1ebg.gpf + test_1ebg_rec.gpf
        """
        self.mv.ADgpf4_readMacromolecule('1ebg_rec.pdbqt')
        self.mv.ADgpf4_readFormattedLigand('1ebg_lig.pdbqt')
        self.mv.ADgpf4_writeGPF("test_1ebg_rec.gpf")
        testptr = open("test_1ebg_rec.gpf")
        testlines = testptr.readlines()
        refptr = open("ref_1ebg.gpf")
        reflines = refptr.readlines()
        for refline,testline in zip(reflines,testlines):
            self.assertEqual(refline, testline)


    def test_gpf4_init_Macro_name(self):
        """readgpfTests12: check macro mol name read by init macro
        """
        self.mv.ADgpf4_readMacromolecule("1ebg_rec.pdbqt")
        self.assertEqual(self.mv.gpo.receptor_stem,'1ebg_rec')
        

    def test_gpf4_check_macro_types(self):
        """readgpfTests13: check macro types
        """
        self.mv.ADgpf4_readMacromolecule("1ebg_rec.pdbqt")
        mol = self.mv.Mols[0]
        d = {}
        #build a list of atom autodock types for the mol
        for a in mol.allAtoms:
            d[a.autodock_element] = 1
        types = d.keys() #list of types in macro
        #checking that the gpo has all of those values
        for ch in split(self.mv.gpo['receptor_types']['value']):
            self.assertEqual(ch in types, True)


class ADgpf4_flexResTests(ADgpf4_BaseTests):

    #tests based on Floren's problem
    def test_gpf4_lig_flexres(self):
        """flexResTests: read a ligand file then a flexres file and check written types
        """
        rigid_filename = 'CathepsinB_rigid.pdbqt'
        flexres_filename = 'CathepsinB_flex.pdbqt'
        ligand_filename = 'FGA3.pdbqt'
        output_filename = 'lig_flexres.gpf'
        self.mv.ADgpf4_readMacromolecule(rigid_filename)
        self.mv.ADgpf4_readFormattedLigand(ligand_filename)
        self.mv.ADgpf4_readFormattedFlexRes(flexres_filename)
        self.mv.ADgpf4_writeGPF(output_filename)
        fptr = open(output_filename)
        lines = fptr.readlines()
        types_line = lines[4]
        ref_line ="ligand_types C SA A OA HD N          # ligand atom types\n"
        self.assertEqual(types_line,ref_line)


    def test_gpf4_flexres_lig(self):
        """flexResTests: read a flexres file then a ligand file and check written types
        """
        rigid_filename = 'CathepsinB_rigid.pdbqt'
        flexres_filename = 'CathepsinB_flex.pdbqt'
        ligand_filename = 'FGA3.pdbqt'
        output_filename = 'lig_flexres.gpf'
        self.mv.ADgpf4_readMacromolecule(rigid_filename)
        self.mv.ADgpf4_readFormattedFlexRes(flexres_filename)
        self.mv.ADgpf4_readFormattedLigand(ligand_filename)
        self.mv.ADgpf4_writeGPF(output_filename)
        fptr = open(output_filename)
        lines = fptr.readlines()
        types_line = lines[4]
        ref_line ="ligand_types A C OA HD N SA          # ligand atom types\n"
        self.assertEqual(types_line,ref_line)


    #read receptor
    def test_gpf4_read_rigid_macro_mol(self):
        """flexResTests1: read a rigid receptor and check gpo receptor_filename 
        """
        filename = 'hsg1_rigid.pdbqt'
        self.mv.ADgpf4_readMacromolecule(filename)
        self.assertEqual(self.mv.gpo.receptor_filename,filename)
        args = self.mv.Mols[0].chains.residues.get(lambda x:x.name=='ARG8')
        #check that 5 atoms only are built for the ARG8 residues
        self.assertEqual(len(args.atoms), 10)


    #read ligand
    def test_gpf4_read_flex_mol(self):
        """flexResTests2: read flexible input file and check gpo ligand_filename 
        """
        filename = 'hsg1_flex.pdbqt'
        self.mv.ADgpf4_readFormattedLigand(filename)
        ligand_name = self.mv.Mols[0].name
        self.assertEqual(self.mv.gpo.ligand.name, ligand_name)
        self.assertEqual('ARG8' in self.mv.Mols[0].chains.residues.name, True)
        args = self.mv.Mols[0].chains.residues.get(lambda x:x.name=='ARG8')
        self.assertEqual(len(args), 2)
        self.assertEqual(len(args.atoms.bonds[0])>0, True)
    

    #choose ligand
    def test_gpf4_choose_flex_mol(self):
        """flexResTests3: read in then choose flexible input file. 
        Check ligand_filename in mv.gpo object
        """
        filename = 'hsg1_flex.pdbqt'
        self.mv.readMolecule(filename)
        ligand_name = self.mv.Mols[0].name
        self.assertEqual('ARG8' in self.mv.Mols[0].chains.residues.name, True)
        args = self.mv.Mols[0].chains.residues.get(lambda x:x.name=='ARG8')
        self.assertEqual(len(args), 2)
        self.assertEqual(len(args.atoms.bonds[0])>0, True)
        self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
        self.assertEqual(self.mv.gpo.ligand.name, ligand_name)


    #choose receptor
    def test_gpf4_choose_rigid_mol(self):
        """flexResTests4: readin then choose rigid input file. check gpo macro_filename 
        """
        filename = 'hsg1_rigid.pdbqt'
        self.mv.readMolecule(filename)
        macro_name = self.mv.Mols[0].name
        self.assertEqual('ARG8' in self.mv.Mols[0].chains.residues.name, True)
        args = self.mv.Mols[0].chains.residues.get(lambda x:x.name=='ARG8')
        self.assertEqual(len(args), 2)
        #check that 5 atoms only are built for the ARG8 residues
        self.assertEqual(len(args.atoms), 10)
        self.assertEqual(len(args.atoms.bonds[0])>0, True)
        self.mv.ADgpf4_chooseMacromolecule(macro_name)
        self.assertEqual(self.mv.gpo.receptor_stem, macro_name)


class ADgpf4_setting_flex_types_tests(ADgpf4_BaseTests):

    def setUp(self):
        """
        clean-up
        """
        global klass, mv
        name = self.__class__.__name__
        if not klass:
            klass = name
        elif klass != name:
            print 'SETUP: destroying mv'
            if mv and mv.hasGui:
                mv.Exit(0)
                if hasattr(self, 'mv'):
                    del self.mv
                mv = None
            klass = name
            from AutoDockTools import autoflexCommands
            reload(autoflexCommands)
        if not hasattr(self, 'mv'):
            self.startViewer()
            
        for m in self.mv.Mols:
            self.mv.deleteMol(m)
        self.mv.loadModule('autoflexCommands', 'AutoDockTools')
        self.mv.ADgpf4_readMacromolecule('hsg1.pdbqt')
        #self.mv.ADflex_setResidues("::ARG8")
        #self.mv.ADflex_writeFlexFile('temp_FLEX.pdbqt')
        #self.mv.ADflex_writeRigidFile('temp_RIGID.pdbqt')

    
    def test_gpf4_set_map_types_directly(self):
        """typesTest0: check for setting map types directly
        """
        teststring = "A B C D E"
        self.mv.ADgpf4_setMapTypes(teststring)
        testlist = teststring.split()
        self.assertEqual(self.mv.GPF_LIGAND_TYPES, testlist)
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)
                
 
    def test_gpf4_choose_ligand_name(self):
        """typesTests1: choose ligand then check GPF_LIGAND_TYPES for
        autodock_element of each atom 
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.readMolecule(filename)
        ligand_name = self.mv.Mols[0].name
        self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
        d = {}
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        for k in d.keys():    
            self.assertEqual(k in self.mv.GPF_LIGAND_TYPES, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)


    def test_gpf4_read_formatted_ligand(self):
        """typesTest2: read ligand then check ligand name
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.ADgpf4_readFormattedLigand(filename)
        d = {}
        for a in self.mv.Mols[-1].allAtoms:
            d[a.autodock_element] = 1
        for k in d.keys():    
            print 'testing k=', k
            self.assertEqual(k in self.mv.GPF_LIGAND_TYPES, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)


    def test_gpf4_set_map_types_from_flex_setup_residues(self):
        """typesTest3: check for setting map types from newly 
                       formatted flexible residues
        """
        filename = 'hsg1.pdbqt'
        self.mv.readMolecule(filename)
        macro_name = "hsg1"
        self.mv.ADgpf4_chooseMacromolecule(macro_name)
        self.mv.ADflex_setResidues("::MET46")
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:MET46:CA,CB;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:MET46:CA,CB;", False)
        atom_types =['C','SA']
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)


    def test_gpf4_set_map_types_from_flex_file(self):
        """typesTest4: readFormattedFlexRes
        check for setting map types from flexible residues file
        """
        self.mv.ADgpf4_readFormattedFlexRes("met46_flex.pdbqt")
        atom_types =['C','SA']
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_FLEXRES_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)


    def test_gpf4_set_map_types_by_choosing_flex(self):
        """typesTest5: _chooseFormattedFlexRes
        check for setting map types by choosing flexible residues 'mol'
        """
        filename = 'met46_flex.pdbqt'
        self.mv.readMolecule(filename)
        self.mv.ADgpf4_chooseFormattedFlexRes("met46_flex")
        atom_types =['C','SA']
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(atom_types, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)

 
    def test_gpf4_choose_ligand_then_currently_formatted_flex(self):
        """typesTest6: choose ligand then set flex res from currently
        formatted molecule and check GPF_LIGAND_TYPES for
        autodock_element of each atom 
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.readMolecule(filename)
        ligand_name = self.mv.Mols[0].name
        self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
        d = {}
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        filename = 'hsg1.pdbqt'
        self.mv.readMolecule(filename)
        macro_name = "hsg1"
        self.mv.ADgpf4_chooseMacromolecule(macro_name)
        self.mv.ADflex_setResidues("::MET46")
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:MET46:CA,CB;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:MET46:CA,CB;", False)
        d['C'] = 1
        d['SA'] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)


    def test_gpf4_choose_ligand_then_read_formatted_flex(self):
        """typesTest7: choose ligand then set flex res from file
        and check GPF_LIGAND_TYPES for autodock_element of each atom 
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.readMolecule(filename)
        ligand_name = self.mv.Mols[0].name
        self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
        d = {}
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        self.mv.ADgpf4_readFormattedFlexRes('met46_flex.pdbqt')
        d['C'] = 1
        d['SA'] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        gpo_types  = self.mv.gpo['ligand_types']['value'].split()
        for x in atom_types:
            self.assertEqual(x in gpo_types, True)
        for t in gpo_types:
            self.assertEqual(t in atom_types, True)


    def test_gpf4_choose_ligand_then_choose_formatted_flex(self):
        """typesTest8: choose ligand then choose flex res in viewer
        and check GPF_LIGAND_TYPES for autodock_element of each atom 
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.readMolecule(filename)
        ligand_name = self.mv.Mols[0].name
        self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
        d = {}
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        self.mv.readMolecule("met46_flex.pdbqt")
        self.mv.ADgpf4_chooseFormattedFlexRes('met46_flex.pdbqt')
        d['C'] = 1
        d['SA'] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)

 
    def test_gpf4_read_ligand_then_currently_formatted_flex(self):
        """typesTest9: read ligand then set flex res from currently
        formatted molecule and check GPF_LIGAND_TYPES for
        autodock_element of each atom 
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.ADgpf4_readFormattedLigand(filename)
        d = {}
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        filename = 'hsg1.pdbqt'
        self.mv.readMolecule(filename)
        macro_name = "hsg1"
        self.mv.ADgpf4_chooseMacromolecule(macro_name)
        self.mv.ADflex_setResidues("::MET46")
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:MET46:CA,CB;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:MET46:CA,CB;", False)
        d['C'] = 1
        d['SA'] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)


    def test_gpf4_read_ligand_then_read_formatted_flex(self):
        """typesTest10: read ligand then read flex res file
        and check GPF_LIGAND_TYPES for autodock_element of each atom 
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.ADgpf4_readFormattedLigand(filename)
        d = {}
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        self.mv.ADgpf4_readFormattedFlexRes('met46_flex.pdbqt')
        d['C'] = 1
        d['SA'] = 1
        atom_types = d.keys()
        #msg = 'atom_types =' + str(atom_types)
        #self.mv.warningMsg(msg)
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(atom_types, ' ')
        for x in atom_types:
            gpo_val = self.mv.gpo['ligand_types']['value']
            gpo_types = gpo_val.split()
            #msg = 'testing x ='+ x+ ' is in '+ gpo_val
            #self.mv.warningMsg( msg)
            self.assertEqual(x in atom_types, True)



    def test_gpf4_read_ligand_then_choose_formatted_flex(self):
        """typesTest11: read ligand then choose flex res molecule
        and check GPF_LIGAND_TYPES for autodock_element of each atom 
        """
        filename = "1ebg_lig.pdbqt"
        self.mv.ADgpf4_readFormattedLigand(filename)
        d = {}
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        self.mv.readMolecule("met46_flex.pdbqt")
        self.mv.ADgpf4_chooseFormattedFlexRes('met46_flex.pdbqt')
        d['C'] = 1
        d['SA'] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)

 
 #------------------------------------------------------------------

    def test_gpf4_currently_formatted_flex_then_choose_ligand(self):
        """typesTest12: set flex res from currently formatted molecule 
        then choose ligand and check GPF_LIGAND_TYPES for
        autodock_element of each atom 
        """
        filename = 'hsg1.pdbqt'
        self.mv.readMolecule(filename)
        macro_name = "hsg1"
        self.mv.ADgpf4_chooseMacromolecule(macro_name)
        self.mv.ADflex_setResidues("::MET46")
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:MET46:CA,CB;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:MET46:CA,CB;", False)
        d = {}
        d['C'] = 1
        d['SA'] = 1
        filename = "1ebg_lig.pdbqt"
        self.mv.readMolecule(filename)
        ligand_name = self.mv.Mols[0].name
        self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)


    def test_gpf4_read_formatted_flex_then_choose_ligand(self):
        """typesTest13: read flex res from file then choose ligand 
        and check GPF_LIGAND_TYPES for autodock_element of each atom 
        """
        self.mv.ADgpf4_readFormattedFlexRes('met46_flex.pdbqt')
        d = {}
        d['C'] = 1
        d['SA'] = 1
        filename = "1ebg_lig.pdbqt"
        self.mv.readMolecule(filename)
        ligand_name = self.mv.Mols[0].name
        self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)


    def test_gpf4_choose_formatted_flex_then_choose_ligand(self):
        """typesTest14: choose flex res in viewer then choose ligand 
        and check GPF_LIGAND_TYPES for autodock_element of each atom 
        """
        self.mv.readMolecule("met46_flex.pdbqt")
        self.mv.ADgpf4_chooseFormattedFlexRes('met46_flex.pdbqt')
        d = {}
        d['C'] = 1
        d['SA'] = 1
        filename = "1ebg_lig.pdbqt"
        self.mv.readMolecule(filename)
        ligand_name = self.mv.Mols[0].name
        self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)

 
    def test_gpf4_currently_formatted_flex_then_read_ligand(self):
        """typesTest15: set flex res from currently formatted 
        molecule then read ligand and check GPF_LIGAND_TYPES for
        autodock_element of each atom 
        """
        filename = 'hsg1.pdbqt'
        self.mv.readMolecule(filename)
        macro_name = "hsg1"
        self.mv.ADgpf4_chooseMacromolecule(macro_name)
        self.mv.ADflex_setResidues("::MET46")
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:MET46:CA,CB;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:MET46:CA,CB;", False)
        d = {}
        d['C'] = 1
        d['SA'] = 1
        filename = "1ebg_lig.pdbqt"
        self.mv.ADgpf4_readFormattedLigand(filename)
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)


    def test_gpf4_read_formatted_flex_then_read_ligand(self):
        """typesTest16: read flex res file then read ligand 
        and check GPF_LIGAND_TYPES for autodock_element of each atom 
        """
        self.mv.ADgpf4_readFormattedFlexRes('met46_flex.pdbqt')
        d = {}
        d['C'] = 1
        d['SA'] = 1
        filename = "1ebg_lig.pdbqt"
        self.mv.ADgpf4_readFormattedLigand(filename)
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)


    def test_gpf4_choose_formatted_flex_then_read_ligand(self):
        """typesTest17: choose flex res molecule then read ligand 
        and check GPF_LIGAND_TYPES for autodock_element of each atom 
        """
        self.mv.readMolecule("met46_flex.pdbqt")
        self.mv.ADgpf4_chooseFormattedFlexRes('met46_flex.pdbqt')
        d = {}
        d['C'] = 1
        d['SA'] = 1
        filename = "1ebg_lig.pdbqt"
        self.mv.ADgpf4_readFormattedLigand(filename)
        for a in self.mv.Mols[0].allAtoms:
            d[a.autodock_element] = 1
        atom_types = d.keys()
        for x in self.mv.GPF_FLEXRES_TYPES:
            self.assertEqual(x in atom_types, True)
        teststring = join(self.mv.GPF_LIGAND_TYPES, ' ')
        self.assertEqual(self.mv.gpo['ligand_types']['value'],teststring)



class ADgpf4_formatting_flexResTests(ADgpf4_BaseTests):

    def setUp(self):
        """
        clean-up
        """
        global klass, mv
        name = self.__class__.__name__
        if not klass:
            klass = name
        elif klass != name:
            print 'SETUP: destroying mv'
            if mv and mv.hasGui:
                mv.Exit(0)
                if hasattr(self, 'mv'):
                    del self.mv
                mv = None
            klass = name
            
        if not hasattr(self, 'mv'):
            self.startViewer()
        for m in self.mv.Mols:
            self.mv.deleteMol(m)
        self.mv.loadModule('autoflexCommands', 'AutoDockTools')
        self.mv.readMolecule('hsg1.pdbqt', ask=0, parser=None)
        self.mv.ADflex_chooseMacro("hsg1")
        from MolKit.molecule import Atom
        self.mv.setIcomLevel(Atom, KlassSet=None)
        self.mv.ADflex_setResidues("::ARG8")
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:ARG8:NE,CZ;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:ARG8:CZ,NH1;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:ARG8:CZ,NH2;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:ARG8:NE,CZ;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:ARG8:CZ,NH1;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:ARG8:CZ,NH2;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:B:ARG8:CA,CB;", False)
        self.mv.ADflex_setBondRotatableFlag("hsg1:A:ARG8:CA,CB;", False)
        self.mv.ADflex_writeFlexFile('temp_FLEX.pdbqt')
        self.mv.ADflex_writeRigidFile('temp_RIGID.pdbqt')


    #choose ligand
    #def test_gpf4_choose_flex_mol(self):
    #    """choose new flexible file. check gpo ligand_filename 
    #    """
    #    ligand_name = 'ind'
    #    #oldval = self.mv.gpo['ligand_types']['value']
    #    written_ligand_file = "temp_FLEX.pdbqt"
    #    self.mv.ADgpf4_chooseFormattedLigand(ligand_name)
    #    self.assertEqual(self.mv.gpo.ligand_filename, written_ligand_file)


    #choose receptor
    def test_gpf4_choose_rigid_mol(self):
        """choose new rigid file. check gpo macro_filename 
        """
        macro_name = "hsg1"
        written_macro_name = "temp_RIGID"
        self.mv.ADgpf4_chooseMacromolecule(macro_name)
        self.assertEqual(self.mv.gpo.receptor_stem, written_macro_name)
        self.assertEqual(self.mv.gpo['gridfld']['value'], written_macro_name+".maps.fld")
        self.assertEqual(self.mv.gpo['dsolvmap']['value'], written_macro_name+".d.map")
        self.assertEqual(self.mv.gpo['elecmap']['value'], written_macro_name+".e.map")


class ADgpf4_write_Covalent_GPF_Tests(ADgpf4_BaseTests):
    def setUp(self):
        ADgpf4_BaseTests.setUp(self)
        from AutoDockTools import autogpfCommands
        reload(autogpfCommands)

    def test_gpf4_write_covalent_gpf(self):
        """covalent_map: (1)read rec and lig files,(2)setup covalent atom
        (3) setup map location, (4)write gpf (5)FINALLY compare with ref gpf
        """
        self.mv.ADgpf4_readMacromolecule('cov_rec.pdbqt', preserve_input_charges=1, log=0)
        self.mv.ADgpf4_readFormattedLigand('cov_lig.pdbqt', log=0)
        self.mv.ADgpf4_readFormattedFlexRes('cov_flex.pdbqt', log=0)
        #self.mv.select("LR_cov:d:****:C;",log=0)
        self.mv.ADgpf_setGpo( covalentmap=1, covalent_coords='20.0000 19.0000 18.0000', 
                ligand_types='C SA HD OA N A Z', covalent_half_width='5.0', 
                covalent_energy_barrier='1000.0')
        self.mv.ADgpf_selectCenter("cov_flex:A:PHE154:CD2", log=0)
        self.mv.ADgpf_setGpo(gridcenterAuto=0, gridcenter=[16.284, 73.749, 28.198], log=0, npts=[60, 60, 60])
        self.mv.ADgpf4_writeGPF('test_cov.gpf', log=0)
        #then compare
        fptr = open("ref_cov.gpf")
        reflines = fptr.readlines()
        fptr.close()
        fptr = open("test_cov.gpf")
        testlines = fptr.readlines()
        fptr.close()
        for rline,tline in zip(reflines, testlines):
            self.assertEqual(rline,tline)


if __name__ == '__main__':
    #unittest.main()
    test_cases = [
        'ADgpf4_BaseTests',
        'ADgpf4_readgpfTests',
        'ADgpf4_flexResTests',
        'ADgpf4_setting_flex_types_tests',
        'ADgpf4_formatting_flexResTests',
        'ADgpf4_write_Covalent_GPF_Tests',
            ]
    unittest.main( argv=([__name__ ,] + test_cases) )



