#
#
#  $Id: test_ADgpf.py,v 1.11 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 = 42


class ADgpf_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('autotorsCommands', 'AutoDockTools')
            mv.loadModule('autogpfCommands', 'AutoDockTools')
            mv.loadModule('autodpfCommands', 'AutoDockTools')
            #8/18:
            mv.loadModule('autoanalyzeCommands', '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 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)

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

class ADgpf_readgpfTests(ADgpf_BaseTests):
    
    
    def test_gp_read_macro_mol(self):
        """reads a macro mol and checks gpo receptor_filename 
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        self.assertEqual(self.mv.gpo.receptor_filename,"hsg1.pdbqs")
    

    def test_gp_read_empty_macro_mol(self):
        
        """reads a empty macro mol and checks that IOError is raised
        """
        self.assertRaises(IOError,self.mv.ADgpf_readMacromolecule," ")

    
    def test_gp_read_non_existent_macro_mol(self):
        """reads a non-exixtent macro mol and checks that IOError is raised
        """
        self.assertRaises(IOError,self.mv.ADgpf_readMacromolecule,"gfdjs.pdbqs")

       
    def test_gp_choose_macro_mol_name(self):
        """checks chosen macrol name =mv.Mols[0]
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        self.mv.ADgpf_chooseMacromolecule("hsg1")
        self.assertEqual(self.mv.Mols[0].name,"hsg1")


    def test_gp_choose_macro_mol_widget(self):
        
        """checks widget of choose macro mol
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        self.mv.ADgpf_chooseMacromolecule("hsg1")
        c=mv.ADgpf_initLigand
        self.assertEqual(c.form.root.winfo_ismapped(),1)
        c.Close_cb()
            

    def test_gp_add_solvation_parameters_comparison(self):
        
        """compares two files in add solvation parameters after writing
        """
        
        mol=self.mv.readMolecule('hsg1.pdb')[0]
        self.mv.setIcomLevel(Atom, KlassSet=None)
        self.mv.selectFromString(mols='',chains='',res= 'HOH*',atoms='*', negate=0, silent=1)
        self.mv.deleteAtomSet(self.mv.getSelection())
        self.mv.add_hGC("hsg1:::", polarOnly=1, renumber=1)
        testname='test_hsg1.pdbqs'
        self.mv.ADgpf_addSolvationParameters(mol,testname)
        testptr=open('test_hsg1.pdbqs')
        testlines=testptr.readlines()
        if find(testlines[0], 'REMARK')>-1:
            testlines = testlines[1:]
        realptr=open('ref.pdbqs')
        reallines=realptr.readlines()
        if find(reallines[0], 'REMARK')>-1:
            reallines = reallines[1:]
       
        for realline,testline in zip(reallines,testlines):
            if find(testline, '-0.000')> -1 or find(realline, '-0.000')>-1:
                #allow 0.000 to match -0.000
                realS = strip(realline)
                testS = strip(testline)
                #compare up to zero
                self.assertEqual(realS[:70], testS[:70])
                #THIS SKIPS space 70 which has '-' sign on linux
                #compare unsigned zeros to the end
                self.assertEqual(realS[71:], testS[71:])
            else:
                self.assertEqual(strip(realline), strip(testline))
    

    def test_gp_add_solvation_paramters_replace_values(self):
        
        """checks by changing sol_par 
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        self.mv.ADgpf_chooseMacromolecule("hsg1")
        c={}
        c['sol_par']=0.5
        apply(mv.ADgpf_setGpo, (), c)
        self.assertEqual(self.mv.gpo['sol_par']['value'],0.5)
    
    
    def test_gp_set_map_types_directly_widget(self):
        
        """checks for widget of set map types
        """
        self.mv.ADgpf_readFormattedLigand("ind.out.pdbq")

        self.mv.ADgpf_setMapTypes("CAONSH",1,1,0)
        v=self.mv.ADgpf_initLigand
        self.assertEqual(v.form.root.winfo_ismapped(),1)
        v.Close_cb()
                

    def test_gp_set_map_types_directly_change_values(self):
        """checks by changing parmaters values
        """
                
        self.mv.ADgpf_setMapTypes("CAONSH",1,1,0)
        self.assertEqual(self.mv.gpo['types']['value'],"CAONSH")
        self.assertEqual(self.mv.gpo['NHB']['value'],1)
        self.assertEqual(self.mv.gpo['OHB']['value'],1)
        self.assertEqual(self.mv.gpo['SHB']['value'],0)
    

    def test_gp_set_map_types_values_in_widget(self):
        
        """checks values in widget after setting them using command
        """
        c=self.mv.ADgpf_setMapTypes
        c.gtypes.set("CAONSH")
        c.OHBonds.set(1)
        self.assertEqual(self.mv.gpo['OHB']['value'],1)
        self.assertEqual(self.mv.gpo['types']['value'],"CAONSH")
        c.OHBonds.set(0) 


    def test_gp_set_map_types_values_in_widget_update_hbonds(self):
        
        """checks by updating H bonds in atom type and checks when Hbonds are
        absent o,s,n 
        """
        
        c=self.mv.ADgpf_setMapTypes
        c.buildForm()
        self.assertEqual(c.OHBonds.get(),0) 
        self.assertEqual(c.NHBonds.get(),0)
        self.assertEqual(c.SHBonds.get(),0)
        c.updateHbonds()
        c.OHBonds.set(1)
        c.NHBonds.set(1)
        c.SHBonds.set(1)
        self.assertEqual(self.mv.gpo['types']['value'],"CAONSH")
        self.assertEqual(self.mv.gpo['OHB']['value'],1) 
        self.assertEqual(self.mv.gpo['NHB']['value'],1)
        self.assertEqual(self.mv.gpo['SHB']['value'],1)
        c.Accept_cb()
    
    
    def test_gp_set_map_types_directly_invalid_values(self):
        
        """checks by changing parameters values to invalid
        """
        self.mv.ADgpf_setMapTypes("CAONSH",1,'a',0)
        self.assertRaises(TypeError,self.mv.ADgpf_setMapTypes.OHBonds.get(),'a')


    def test_gp_set_map_types_directly_atoms_in_list(self):
        """checks the entry of map atom types
        """
        var='bc'
        types=['C','A','N','O','S','P','H','n','f','F','c','b','I','M']
        for i in range(0,len(var)):
             for l in types:      
                    if find(l,var[i])==0:
                        break;
                    else:        
                        self.mv.ADgpf_setMapTypes(var,1,1,0)
                        self.assertEqual(mv.gpo['types']['value'],'bc')
                
    
    def test_gp_set_map_types_directly_invalid_map_atom_types(self):
        """checks by changing parmaters values to invalid
        """
        c = self.mv.ADgpf_setMapTypes        
        self.assertEqual(c.gtypes.get()=='a', False)

 
    def test_gp_choose_ligand_name(self):
        """checks ligand name
        """
        self.mv.ADgpf_readFormattedLigand('ind.out.pdbq')
        self.mv.ADgpf_chooseFormattedLigand("ind_out")
        self.assertEqual(self.mv.gpo.ligand.name,self.mv.Mols[0].name)


    def test__gp_choose_ligand_name_widget(self):
        """checks widget for choose ligand name
        """
        self.mv.ADgpf_readFormattedLigand('ind.out.pdbq')
        self.mv.ADgpf_chooseFormattedLigand("ind_out")
        c=mv.ADgpf_initLigand
        c.form
        self.assertEqual(c.form.root.winfo_ismapped(),1)
        c.Close_cb()
        
    
    def test_gpf_length_of_file(self):
        """checks no. of lines in gpf file
        """
        fptr=open("hsg1.gpf")
        alllines=fptr.readlines()
        self.assertEqual(len(alllines),60)


    def test_gpf_replace_random(self):
        """ checks by setting some values in place of random
        """
        self.mv.ADgpf_readGPF("hsg1.gpf")               
        self.mv.ADgpf_readFormattedLigand("ind.out.pdbq")
        teststring = "1.0"
        self.mv.ADgpf_setGpo(sol_par=teststring)
        testfile = "hsg2.gpf"
        self.mv.ADgpf_writeGPF(testfile)
        fptr=open(testfile)
        alllines=fptr.readlines()
        for l in alllines:
            if find(l, 'sol_par')>0:
                 ll=split(alllines[find(l, 'sol_par')])
            else:
                break
            self.assertEqual(ll[i], teststring)
   

    def test_gpf_read_formatted_ligand(self):
        """checks read ligand name
        """
        self.mv.ADgpf_readFormattedLigand('ind.out.pdbq')
        self.assertEqual(self.mv.gpo.ligand_filename,'ind.out.pdbq')


    def test_gpf_read_empty_formatted_ligand(self):
        """checks empty ligand name
        """
        self.assertRaises(IOError,self.mv.ADgpf_readFormattedLigand," ")
        
    
    def test_gpf_read_non_existent_formatted_ligand(self):
        """checks non_existent ligand name
        """
        self.assertRaises(IOError,self.mv.ADgpf_readFormattedLigand,"hjgf.pdbq ")

       
    def test_gpf_set_new_parms_for_new_atom_type(self):
        """ checks values of set new params in widget after assigning
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        s=self.mv.ADgpf_defineAtomParameters
        s.elem.set("C")
        oldRij = s.Rij.get()
        s.Rij.set(1.0)
        oldepsij = s.epsij.get()
        s.epsij.set(2.0)
        self.assertEqual(s.elem.get(),"C")
        self.assertEqual(s.Rij.get(),'1.0')
        self.assertEqual(s.epsij.get(),'2.0')
        s.Rij.set(oldRij)
        s.epsij.set(oldepsij)
    
    
    def test_gpf_set_new_parms_for_new_atom_type_invalid(self):
        """checks by entering invalid parameters in setnew params
        """
        self.assertEqual(self.mv.ADgpf_defineAtomParameters.elem.get()=='W', False)
    

    def test_gpf_set_new_parms_for_new_atom_type_widget(self):
        """checks set new params for new atom type widget
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        s=self.mv.ADgpf_defineAtomParameters
        s.elem.set("C")
        oldRij = s.Rij.get()
        s.Rij.set(1.0)
        oldepsij = s.epsij.get()
        s.epsij.set(2.0)
        c=mv.ADgpf_initLigand
        self.assertEqual(c.form.root.winfo_ismapped(),1)
        s.Rij.set(oldRij)
        s.epsij.set(oldepsij)
        c.Close_cb()
         
    
    def test_gpf_set_up_covalent_map_widget(self):
        """checks set up covalent map widget
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        self.mv.ADgpf_setUpCovalentMap('-22.9870 -0.0600 4.4430',-1.0,'5.0','1000.0','CAONSHZ')
        c=mv.ADgpf_initLigand
        self.assertEqual(c.form.root.winfo_ismapped(),1)
        c.Close_cb()
            
    
    def test_gpf_set_up_covalent_map_coords(self):
        """ checks covalent coords  in set up covalent map
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        self.mv.ADgpf_readFormattedLigand('ind.out.pdbq')
        #c={}
        #c['covalent_coords']='-3.884'
        #apply(self.mv.ADgpf_setUpCovalentMap,(),c)
        self.mv.ADgpf_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_gpf_get_values_from_gpf_read(self):
        """checks gpf filename
        """
        self.mv.ADgpf_readGPF('hsg1.gpf')
        c=self.mv.gpo
        self.assertEqual(c.gpf_filename,"hsg1.gpf")
    

    def test_gpf_get_values_from_gpf_read_empty(self):
        """checks empty gpf file name
        """
        self.assertRaises(IOError,self.mv.ADgpf_readGPF," ")
    

    def test_gpf_get_values_from_gpf_read_non_existent(self):
        """checks non existent gpf file name
        """
        self.assertRaises(IOError,self.mv.ADgpf_readGPF,"hsfh.gpf")

    
    def test_gpf_write_gpf(self):
        """checks ref_hsg1.gpf test_hsg1.gpf
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        self.mv.ADgpf_readFormattedLigand('ind.out.pdbq')
        self.mv.ADgpf_writeGPF("test_hsg1.gpf")
        testptr = open("test_hsg1.gpf")
        testlines = testptr.readlines()
        refptr = open("ref_hsg1.gpf")
        reflines = refptr.readlines()
        for refline,testline in zip(reflines,testlines):
            self.assertEqual(refline,testline)
        c=mv.ADgpf_initLigand
        c.Close_cb()


    def test_set_up_grid_widget(self):
        """ checks for set up grid widget
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        self.mv.ADgpf_setGrid.buildForm()
        self.mv.ADgpf_setGrid((0.369,-0.215,-4.987),(1,0,0),0.5)
        c=self.mv.ADgpf_setGrid
        self.assertEqual(c.form.root.winfo_ismapped(),1)
        c.Close_cb()
        self.mv.ADgpf_setGrid.Accept_cb()
        

    def test_set_up_grid_coords(self):
        """ checks center coords in grid options widget and Mol
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        #have to build the form in order to create the Tkinter variables
        self.mv.ADgpf_setGrid.buildForm()
        self.mv.ADgpf_setGrid((2.428,0.812,-7.490),(1,0,0),0.5)
        #X center
        self.assertEqual(str(round(self.mv.Mols[-1].center[0],3)),'2.428')
        #Y center
        self.assertEqual(str(round(self.mv.Mols[-1].center[1],3)),'0.812')
        #Z center
        self.assertEqual(str(round(self.mv.Mols[-1].center[2],3)),'-7.49')
        #close the form opened by call to buildForm 
        self.mv.ADgpf_setGrid.Accept_cb()


    def test_set_up_grid_selected_atom_coords(self):
        """checks coords of grid after selecting an atom
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        
        self.mv.ADgpf_setGrid.buildForm()
        self.mv.ADgpf_selectCenter("hsg1:A:LYS70:CE")
        c=self.mv.ADgpf_setGrid
        d={}
        d['spacing']=0.5
        d['gridcenter']=(2.428,0.812,-7.490)
        d['npts']=(2,4,6)
        apply(c, (),d)
        #X center
        self.assertEqual(str(c.xCen.get()),'2.428')
        #Y center
        self.assertEqual(str(c.yCen.get()),'0.812')
        #Z center
        self.assertEqual(str(c.zCen.get()),'-7.49')
        self.mv.ADgpf_setGrid.Accept_cb()


    def test_set_up_grid__value_of_spacewheel(self):
        """checks space wheel value in set up grid
        """
        self.mv.ADgpf_readMacromolecule('hsg1.pdbqs')
        self.mv.ADgpf_setGrid.buildForm()
        c=self.mv.ADgpf_setGrid
        d={}
        d['spacing']=0.375
        d['gridcenter']=(2.428,0.812,-7.490)
        d['npts']=(2,4,6)
        apply(c, (),d)
        self.assertEqual(str(c.spacewheel.get()),'0.375')
        self.mv.ADgpf_setGrid.Accept_cb()


    def test_set_up_other_options_widget(self):
        """checks widget of set other options
        """
        self.mv.ADgpf_readMacromolecule("hsg1.pdbqs")
        self.mv.ADgpf_readFormattedLigand("ind.out.pdbq")
        self.mv.ADgpf_setOtherOptions(0.6,1,0.146)
        c=mv.ADgpf_initLigand
        self.assertEqual(c.form.root.winfo_ismapped(),1)
        c.Close_cb()
                
    
    def test_set_up_other_options_values(self):
        """checks smoothfactor,dielectric in set other options
        """
        self.mv.ADgpf_readMacromolecule("hsg1.pdbqs")
        c=self.mv.ADgpf_setOtherOptions
        g={}
        g['smooth']=0.6
        g['dielectric']=0.146
        apply(c, (),g)
        self.assertEqual(round(mv.gpo['smooth']['value'],3),0.60)
        self.assertEqual(round(mv.gpo['dielectric']['value'],3),0.146)
    

    def test_gpf_merge_non_polar_hydrogens(self):
        """checks number of atoms
        """
        m=self.mv.readMolecule('hsg1.pdb')[0]
        self.mv.setIcomLevel(Atom, KlassSet=None)
        self.mv.selectFromString(mols='',chains='',res= 'HOH*',atoms='*')
        self.mv.deleteAtomSet(self.mv.getSelection())
        self.mv.add_hGC("hsg1:::", polarOnly=0, renumber=1)
        self.mv.addKollmanCharges(m)
        orig_num_atoms=len(m.allAtoms)
        #self.mv.selectSet('nphs')
        nphs=m.allAtoms.get(lambda x:x.element=='H' and (x.bonds[0].atom1.element=='C' or x.bonds[0].atom2.element=='C'))
        len_nphs=len(nphs)
        c_nhps=m.allAtoms.get(lambda x:x.element=='C' and x.findHydrogens() > 0)
        #print "len(c_nphs)=", len(c_nphs)
        self.mv.ADgpf_mergeNonPolarHydrogens(m, outfile='xxxx.pdbqs')
        self.assertEqual(len(m.allAtoms),(orig_num_atoms)-(len_nphs))


    def test_gpf_merge_non_polar_hydrogens_charges(self):
        """checks atom charges after merging
        """
        m=self.mv.readMolecule('hsg1.pdb')[0]
        self.mv.setIcomLevel(Atom, KlassSet=None)
        self.mv.selectFromString(mols='',chains='',res= 'HOH*',atoms='*')
        self.mv.deleteAtomSet(self.mv.getSelection())
        self.mv.add_hGC("hsg1:::", polarOnly=0, renumber=1)
        self.mv.addKollmanCharges(m)
        #make a record of the charges on each nonpolarCarbon + its hydrogens
        c_nhps=m.allAtoms.get(lambda x:x.element=='C' and x.findHydrogens() > 0)
        #this dict keys are the nonpolarCarbons 
        #its values are their united charge [orig charge + charge on each nph]
        c_charges={}
        for c_atom in c_nhps:                                       
            hs=c_atom.findHydrogens()
            h_tot_charge=0
            for h in hs:
                h_tot_charge+=h.charge
            c_charges[c_atom]=c_atom.charge+h_tot_charge
        self.mv.ADgpf_mergeNonPolarHydrogens(m, outfile='xxxx.pdbqs')
        for c_atom in c_nhps:
            self.assertEqual(c_atom.charge, c_charges[c_atom])


    def test_gpf_init_ligand_name(self):
        """ checks ligand name read by init ligand
        """
        self.mv.ADgpf_readFormattedLigand("ind.out.pdbq")
        self.assertEqual(self.mv.gpo.ligand.name,'ind_out')



    def test_gpf_init_ligand_with_halogen(self):
        """ checks ligand with halogen sets types correctly
        """
        self.mv.ADgpf_readFormattedLigand("nutlin.pdbq")
        self.assertEqual(self.mv.gpo['types']['value'], 'CANOb')

                

    def test_gpf_init_Macro_name(self):
        """checks macro mol name read by init macro
        """
        self.mv.ADgpf_readMacromolecule("hsg1.pdbqs")
        self.assertEqual(self.mv.gpo.receptor_stem,'hsg1')
        

    def test_gpf_check_macro_types(self):
        """checks macro types
        """
        self.mv.ADgpf_readMacromolecule("hsg1.pdbqs")
        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 self.mv.gpo['mset']['value']:
            self.assertEqual(ch in types, True)



if __name__ == '__main__':
    unittest.main()


