def test_simplestride():
    if haveStride:
        from MolKit import Read
        mol = Read("./Data/1crn.pdb")[0]
        mol.secondaryStructureFromStride()
        assert hasattr(mol.chains[0], 'secondarystructureset') and \
               len(mol.chains[0].secondarystructureset)==11

def test_getsecondarystructureChain():
    if haveStride:
        from MolKit import Read
        mol = Read("./Data/fxnohtatm.pdb")[0]
        from MolKit.getsecondarystructure import GetSecondaryStructureFromStride
        ssb = GetSecondaryStructureFromStride(mol)
        mol.secondaryStructure(ssb)

 def test_Write_6(self):
     
     if not haveStride: return
     from MolKit import Read
     self.mol = Read("Data/1crn.pdb")[0]
     self.mol.secondaryStructureFromStride()
     from MolKit.pdbWriter import PdbWriter
     writer = PdbWriter()
     writer.write('Data/1crn_writer.pdb', self.mol,
                  records=['ATOM', 'HETATM', 'CONECT',
                           'TURN','HELIX', 'SHEET'],
                  bondOrigin='all', ssOrigin='Stride')
     # Make sure that the ss information has been written out
     # properly.
     nmol = Read("Data/1crn_writer.pdb")[0]
     nmol.secondaryStructureFromFile()
     nsset = nmol.chains[0].secondarystructureset
     osset = self.mol.chains[0].secondarystructureset
     self.assertEqual(len(nsset), len(osset))
     for nss, oss in map(None, nsset, osset):
         self.assertEqual(nss.name, oss.name)
         self.assertEqual(nss.start.name,oss.start.name)
         self.assertEqual(nss.end.name, oss.end.name)
         self.assertEqual(len(nss.residues), len(oss.residues))
     os.system("rm Data/1crn_writer.pdb")

 def set_types_from_directory(self, directory):
     if self.verbose: 
         print "reading directory ", directory
     filelist = glob.glob(directory + "/*.pdb*")
     if self.verbose: 
         print "len(filelist)=", len(filelist)
     ad4_typer = AutoDock4_AtomTyper()
     type_dict = {}
     for f in filelist:
         m = Read(f)[0]
         m.buildBondsByDistance()
         ad4_typer.setAutoDockElements(m)
         for a in m.allAtoms:
             type_dict[a.autodock_element] = 1
         self.getSideLengths(m) #sets ligand.center
         npts = m.npts
         #only make the box bigger, do NOT make it smaller
         for ix, val in enumerate(self.gpo['npts']['value']):
             if npts[ix]>val:
                 self.gpo['npts']['value'][ix] =  npts[ix]
                 if self.verbose: 
                     print m.name, " increased grid dimension ", ix, " to ", npts[ix]
     d_types = type_dict.keys()
     if self.verbose: 
         print "found ", d_types, " atom types in directory ", directory
     self.gpo['ligand_types']['value'] =  string.join(d_types)
     if self.verbose: 
         print "now ligand_types is ", self.gpo['ligand_types']['value']

def test_buildBondsByDistance_1():
    """
    Test the buildBondsByDistance for 1crn.pdb with one chain
    
    """
    from MolKit import Read
    mol = Read('Data/1crn.pdb')[0]
    mol.buildBondsByDistance()
    bonds, nobonds = mol.allAtoms.bonds[0], mol.allAtoms.bonds[1]
    assert len(bonds) == 337 and len(nobonds) == 0

def test_merge():
    # merge mol2 into mol1.
    from MolKit import Read
    mol1 = Read("./Data/protease.pdb")[0]
    mol2 = Read("./Data/indinavir.pdb")[0]
    from MolKit.protein import ProteinSet
    pset = ProteinSet([mol1,mol2])
    chains = pset.chains[:]
    mol1.merge(mol2)
    mol1.chains.sort()
    chains.sort()
    assert mol1.chains == chains

def test_findType():
    from MolKit import Read

    mol = Read("./Data/1crn.pdb")[0]
    from MolKit.protein import Residue, Chain, Protein, ResidueSet
    from MolKit.molecule import Atom, AtomSet

    # FindType below from top of the tree
    res = mol.findType(Residue)
    res.sort()
    molres = mol.chains.residues
    molres.sort()
    assert res == molres

    # FindType above from bottom of the tree
    atms = mol.allAtoms
    chains = atms.findType(Chain)
    # Automatically does a uniq when going above
    assert len(chains) == 327

    # FindType above uniq from bottom of the tree
    chainsuniq = atms.findType(Chain, uniq=1)
    assert len(chainsuniq) == 1

    # FindType above from middle of the tree
    prot = res.findType(Protein, uniq=1)
    assert len(prot) == 1 and prot == mol.setClass([mol])

    # FindType below from middle of the tree
    atoms = chainsuniq.findType(Atom)
    assert len(atoms) == 327

    # FindType on empty set
    emptyres = ResidueSet([])
    atoms = emptyres.findType(Atom)
    assert len(atoms) == 0 and isinstance(atoms, AtomSet)

    # FindType on same type
    atoms = atms.findType(Atom)
    assert atoms == atms

    # Raise exception
    from MolKit.protein import Helix

    try:
        nodes = atms.findType(Helix)
    except RuntimeError:
        print "passed"

 def set_ligand(self, ligand_filename, center_on_ligand=False): 
     self.ligand = Read(ligand_filename)[0]
     if self.ligand==None:
         print 'ERROR reading: ', ligand_filename
         return 
     if self.verbose: 
         print "read ", self.ligand.name
     ligand_types = self.getTypes(self.ligand)
     self.gpo.set_ligand4(ligand_filename, types=ligand_types)
     #this sets ligand_types, gpo.ligand_stem and gpo.ligand_filename
     if self.verbose: 
         print "set gpo.ligand_stem to", self.gpo.ligand_stem
         print "set gpo.ligand_filename to", self.gpo.ligand_filename
         print "set gpo.ligand_types to", self.gpo['ligand_types']['value'].__class__
     #need to get npts
     if self.size_box_to_include_ligand:
         self.getSideLengths(self.ligand) #sets ligand.center
     #gridcenter IS NOT SET BY THIS!!!
     if center_on_ligand:
         cen = self.ligand.getCenter()
         self.gpo['gridcenter']['value'] =  [round(cen[0],4), round(cen[1],4),\
                                         round(cen[2],4)]
         self.gpo['gridcenterAuto']['value'] =  0
         if self.verbose: print "set gridcenter to ", self.gpo['gridcenter']['value']
     #only make the box bigger, do NOT make it smaller
     for ix, val in enumerate(self.gpo['npts']['value']):
         #npts
         if hasattr(self.ligand, 'npts'):
             npts = self.ligand.npts
             if npts[ix]>val:
                 self.gpo['npts']['value'][ix] =  npts[ix]
     if self.verbose: print "set npts to ", self.gpo['npts']['value']

class BaseTests(unittest.TestCase):
    def setUp(self):
        from MolKit import Read
        self.mol = Read('Data/1crn.pdb')[0]
        self.mol.buildBondsByDistance()
    

    def tearDown(self):
        """
        clean-up
        """
        del(self.mol)


    def test_constructor(self):
        """
        instantiate an HydrogenBuilder
        """
        h_builder = HydrogenBuilder()
        self.assertEquals(h_builder.__class__, HydrogenBuilder)


    def test_constructorOptions(self):
        """
         test possible constructor options
            options = htype, renumber, method
        """
    
        h_builder = HydrogenBuilder(htype='polarOnly')
        self.assertEquals(h_builder.__class__, HydrogenBuilder)
        h_builder = HydrogenBuilder(renumber=0)
        self.assertEquals(h_builder.__class__, HydrogenBuilder)
        h_builder = HydrogenBuilder(method='withBondOrder')
        self.assertEquals(h_builder.__class__, HydrogenBuilder)


    def test_addHydrogens(self):
        """
         test addHydrogens 
        """
        beforeLen = len(self.mol.allAtoms)
        h_builder = HydrogenBuilder()
        h_builder.addHydrogens(self.mol)
        afterLen = len(self.mol.allAtoms)
        #print "beforeLen=", beforeLen, ' afterLen=', afterLen
        self.assertEquals(beforeLen<afterLen, True)

class BaseTests(unittest.TestCase):
    def setUp(self):
        from MolKit import Read

        self.mol = Read("Data/2plv_no_oxt.pdb")[0]
        self.mol.buildBondsByDistance()

    def tearDown(self):
        """
        clean-up
        """
        del (self.mol)

    def test_constructor(self):
        """
        instantiate an HydrogenBuilder
        """
        oxt_builder = OxtBuilder()
        self.assertEquals(oxt_builder.__class__, OxtBuilder)

    #    def test_constructorOptions(self):
    #        """
    #         test possible constructor options
    #            options = NONE!!
    #        """
    #
    #        oxt_builder = OxtBuilder(???=???)
    #        self.assertEquals(oxt_builder.__class__, OxtBuilder)

    def test_add_oxt(self):
        """
         test add_oxt 
        """
        beforeLen = len(self.mol.allAtoms)
        oxt_builder = OxtBuilder()
        # the last chain is all waters so skip it
        for ch in self.mol.chains[:-1]:
            catom = ch.residues[-1].atoms[2]
            new_at = oxt_builder.add_oxt(catom)
            # print "added ", new_at.full_name()
        afterLen = len(self.mol.allAtoms)
        # print "beforeLen=", beforeLen, ' afterLen=', afterLen
        self.assertEquals(beforeLen + len(self.mol.chains[:-1]), afterLen)

 def set_ligand4(self, ligand_filename, types=None):
     #this should set ligand_types
     #print "in set_ligand4: types=", types
     if types is None:
         ligand = Read(ligand_filename)[0]
         ligand.buildBondsByDistance()
         ad4_typer = AutoDock4_AtomTyper()
         ad4_typer.setAutoDockElements(ligand)
         dict = {}
         for a in ligand.allAtoms:
             dict[a.autodock_element] = 1
         d_types = dict.keys()
         d_types.sort()
         types = d_types[0]
         for t in d_types[1:]:
             types = types + " " + t
     self['ligand_types']['value'] = types
     #print "set_ligand4: self['ligand_types']['value']=", self['ligand_types']['value']
     self.ligand_filename = os.path.basename(ligand_filename)
     #print "GPO: set ligand_filename to ", self.ligand_filename
     self.ligand_stem = os.path.splitext(self.ligand_filename)[0]

class PdbqtWriterAtomLinesTest(PdbWriterTest):

    def setUp(self):
        from MolKit import Read
        self.mol = Read('Data/hsg1.pdbqt')[0]
        self.mol.buildBondsByDistance()
    

    def tearDown(self):
        del(self.mol)


    def test_write(self):
        """
        test writing a pdbqs file
        """
        writer = PdbqtWriter()
        writer.write('test_pdbqtWriter.pdbqt', self.mol, bondOrigin=('File',))
        ans, errors = self.compare('Data/hsg1.pdbqt', 'test_pdbqtWriter.pdbqt') 
        self.assertEquals(errors, None)
        self.assertEquals(ans, True)

def get_best_energy_info(d, rms_tolerance, build_hydrogen_bonds=False, \
                                report_energy_breakdown=False,
                                report_unbound_energy=False,
                                receptor_filename=None):
    ostr = "" 
    clust0 = d.clusterer.clustering_dict[rms_tolerance][0]
    c = clust0[0]
    #find the filename of the best result
    dlo_filename = get_filename(d, c)
    if verbose: print "set dlo_filename to ", dlo_filename
    ostr = dlo_filename + ", "
    ostr += "%3d,%3d,%3d,% 8.4f,% 8.4f," %(len(d.clusterer.data), num_clusters, len(clust0),  c.binding_energy, c.getRMSD())
    d.ch.set_conformation(c)
    receptor = None
    if build_hydrogen_bonds:
        if not receptor_filename:
            print "receptor_filename must be specified in order to build_hydrogen_bonds"
            return 
        receptor = Read(receptor_filename)[0] 
        receptor.buildBondsByDistance()
        ostr += construct_hydrogen_bonds(receptor.allAtoms, d.ligMol.allAtoms)
    if report_energy_breakdown:
        if not receptor_filename:
            print "receptor_filename must be specified in order to build_hydrogen_bonds"
        if not receptor:
            receptor = Read(receptor_filename)[0] 
            receptor.buildBondsByDistance()
        ostr += get_energy_breakdown(receptor.allAtoms, d.ligMol.allAtoms, c)
    if report_unbound_energy:
        if c.unbound_energy is not None: 
            ostr += "% 8.4f," %c.unbound_energy
        else:
            if verbose: print "conformation's unbound energy is None!"
            ostr += "%    0.0," 
    return ostr

def autodock_scoring(receptor, ligand):
    receptorfilename =  receptor
    ligandfilename =  ligand
    write_file_mode = False
    parameter_library_filename = None
    exclude_torsFreeEnergy = False
    verbose = None
    ad_scorer = AutoDock41Scorer(exclude_torsFreeEnergy=exclude_torsFreeEnergy)
    supported_types = ad_scorer.supported_types
    receptor = Read(receptorfilename)[0]
    receptor.buildBondsByDistance()
    ligand = Read(ligandfilename)[0]
    ligand.buildBondsByDistance()

    ms = MolecularSystem()
    ms.add_entities(receptor.allAtoms)
    ms.add_entities(ligand.allAtoms)
    ad_scorer.set_molecular_system(ms)
    #get the scores, score per term:
    [estat, hb, vdw ,dsolv] = ad_scorer.get_score_per_term()
    torsEnrg = ligand.TORSDOF * ad_scorer.tors_weight
    score = estat +hb +vdw +dsolv +torsEnrg
    output_score = {'score':score, 'estat':estat, 'hb':hb, 'vdw':vdw, 'dsolv,':dsolv, 'torsEnrg':torsEnrg}
    
    return output_score

 def set_receptor4(self, receptor_filename, types=None):
     #this should set receptor_types
     if types is None:
         receptor = Read(receptor_filename)[0]
         receptor.buildBondsByDistance()
         ad4_typer = AutoDock4_AtomTyper()
         ad4_typer.setAutoDockElements(receptor)
         dict = {}
         for a in receptor.allAtoms:
             dict[a.autodock_element] = 1
         d_types = dict.keys()
         d_types.sort()
         types = d_types[0]
         for t in d_types[1:]:
             types = types + " " + t
     self['receptor_types']['value'] = types
     basename = os.path.basename(receptor_filename)
     self.receptor_filename = basename
     self.receptor_stem = os.path.splitext(basename)[0]
     if receptor_filename!='':
         self['receptor']['value'] = basename
         self['gridfld']['value'] = self.receptor_stem + '.maps.fld'
         self['elecmap']['value'] = self.receptor_stem + '.e.map'
         self['dsolvmap']['value'] = self.receptor_stem + '.d.map'

 def test_Write_2(self):
     """
     Test writing out CONECT records for the bonds described in
     the pdb file and built by distance.
     """
     from MolKit import Read
     self.mol = Read("Data/1crn.pdb")[0]
     self.mol.buildBondsByDistance()
     # instanciate a PdbWriter and call the write method with the
     # default arguments
     from MolKit.pdbWriter import PdbWriter
     writer = PdbWriter()
     writer.write('Data/1crn_writer.pdb', self.mol)
     from MolKit.pdbParser import PdbParser
     nmol = Read("Data/1crn_writer.pdb")[0]
     os.system("rm Data/1crn_writer.pdb")

    def test_Write_1(self):
        """
        Test the default option of the write method of a PdbWriter.
        Just write the ATOM, HETATM and CONECT records and bondOrigin
        is File.
        """
        # read a molecule
        from MolKit import Read
        self.mol = Read("Data/1crn.pdb")[0]

        # instanciate a PdbWriter and call the write method with the
        # default arguments
        from MolKit.pdbWriter import PdbWriter
        writer = PdbWriter()
        writer.write('Data/1crn_writer.pdb', self.mol)
        # This should write only the ATOM and CONECT records
        # The TER records are automatically created with the ATOM records

        # 1- Make sure that the file has been created
        import os
        self.failUnless(os.path.exists('Data/1crn_writer.pdb'))
        # 2- Make sure that the file created the proper records
        # Get the default records from the new pdb files
        from MolKit.pdbParser import PdbParser
        nmol = Read("Data/1crn_writer.pdb")[0]

        # COMPARE 
        ncoords = nmol.allAtoms.coords
        ocoords = self.mol.allAtoms.coords
        self.assertEqual(ncoords, ocoords)

        nname = nmol.allAtoms.name
        oname = nmol.allAtoms.name
        self.assertEqual(nname, oname)

        nbonds = nmol.allAtoms.bonds[0]
        obonds = self.mol.allAtoms.bonds[0]
        self.assertEqual(len(nbonds), len(obonds))
        nbatms = nbonds.atom1+nbonds.atom2
        nbatms = nbatms.uniq()
        nbatms.sort()
        
        nbatms = nbonds.atom1+nbonds.atom2
        nbatms = nbatms.uniq()
        nbatms.sort()
        os.system("rm Data/1crn_writer.pdb")

 def set_ligand(self, ligand_filename): 
     self.ligand_filename = os.path.basename(ligand_filename)
     if verbose:
         print "set ligand_filename to", self.ligand_filename
     self.dpo.set_ligand(ligand_filename)
     #expect a filename like ind.out.pdbq: get 'ind' from it
     self.ligand_stem = string.split(self.ligand_filename,'.')[0]
     if verbose: print "set ligand_stem to", self.ligand_stem
     self.ligand = Read(ligand_filename)[0]
     if self.ligand==None:
         print 'ERROR reading: ', ligand_filename
         return 
     if verbose: 
         print "read ", self.ligand.name
     #set dpo:
     #move
     self.dpo['move']['value'] = self.ligand_filename
     if verbose: print "set move to ", self.dpo['move']['value']
     #ndihe
     #assumes ligand has torTree
     self.dpo['ndihe']['value'] = self.ligand.parser.keys.count("BRANCH")
     #self.dpo['ndihe']['value'] = len(self.ligand.torTree.torsionMap)
     if verbose: print "set ndihe to ", self.dpo['ndihe']['value']
     #torsdof
     #caution dpo['torsdof']['value'] is a list [ndihe, 0.3113]
     self.dpo['torsdof']['value'][0] = self.ligand.TORSDOF
     if verbose: print "set torsdof to ", self.dpo['torsdof']['value']
     #types
     d = {}
     for a in self.ligand.allAtoms:
         d[a.autodock_element] = 1
     sortKeyList =  ['C','A','N','O','S','H','P','n','f','F','c','b','I','M']
     lig_types = ""
     for t in sortKeyList:
         if t in d.keys():
             lig_types = lig_types + t
     self.ligand.types = lig_types
     self.dpo['types']['value'] = self.ligand.types
     if verbose: print "set types to ", self.dpo['types']['value']
     #about
     self.ligand.getCenter()
     cen = self.ligand.center
     self.dpo['about']['value'] =  [round(cen[0],4), round(cen[1],4),\
                                     round(cen[2],4)]
     if verbose: print "set about to ", self.dpo['about']['value']

def get_largest_cluster_info(d, rms_tolerance, build_hydrogen_bonds=False, \
                                    report_energy_breakdown=False,
                                    report_unbound_energy=False,
                                    receptor_filename=None, refCoords=None,
                                    subtract_internal_energy=False):
    global be_lc
    largest = d.clusterer.clustering_dict[rms_tolerance][0]
    if verbose: print "set largest to ", len(largest)
    for clust in d.clusterer.clustering_dict[rms_tolerance]:
        if verbose: print "current largest cluster len= ", len(clust)
        if len(clust)>len(largest): 
            if verbose: print "resetting largest clust: now len=", len(clust)
            largest = clust
    c = largest[0]   #print info about the lowest energy member of this cluster
    ostr = " "
    dlo_filename = get_filename(d, c)
    ostr += dlo_filename + ","
    be_lc = c.binding_energy
    if subtract_internal_energy:
        be_lc = c.binding_energy - c.total_internal
    if refCoords:
        ostr += "%3d,%3d,%3d,% 8.4f,% 8.4f," %(len(d.clusterer.data), num_clusters, len(largest),  be_lc, c.getRMSD(refCoords=refCoords))
    else:
        ostr += "%3d,%3d,%3d,% 8.4f,% 8.4f," %(len(d.clusterer.data), num_clusters, len(largest),  be_lc, c.getRMSD())
    if verbose: print "set dlo_filename to ", dlo_filename
    #update the coords only if you need to?
    d.ch.set_conformation(c)
    receptor = None
    if build_hydrogen_bonds:
        if not receptor_filename:
            print "receptor_filename must be specified in order to build_hydrogen_bonds"
            return 
        receptor = Read(receptor_filename)[0] 
        receptor.buildBondsByDistance()
        ostr += construct_hydrogen_bonds(receptor.allAtoms, d.ligMol.allAtoms, verbose=verbose)
    if report_energy_breakdown:
        if not receptor_filename:
            print "receptor_filename must be specified in order to build_hydrogen_bonds"
            return 
        if not receptor:
            receptor = Read(receptor_filename)[0] 
            receptor.buildBondsByDistance()
        #ostr += get_energy_breakdown(receptor.allAtoms, d.ligMol.allAtoms, c)
        ostr += get_energy_breakdown(receptor.allAtoms, d.ligMol.allAtoms, c.getCoords()[:])
    if report_unbound_energy: 
        if c.unbound_energy is not None:
            ostr += ",% 8.4f" %c.unbound_energy
        else:
            if verbose: print "conformation's unbound energy is None!"
            ostr += ", 0.0" 
    return ostr

    def test_Write_4(self):
        """
        Test writing out the ATOM records and the CONECT records
        without the HETATM
        """
        from MolKit import Read
        self.mol = Read("Data/1bsr.pdb")[0]
        # instanciate a PdbWriter and call the write method with the
        # default arguments
        from MolKit.pdbWriter import PdbWriter
        writer = PdbWriter()
        writer.write('Data/1bsr_writer.pdb', self.mol,
                     records=['ATOM', 'CONECT'], bondOrigin=('File',))
        nmol = Read("Data/1bsr_writer.pdb")[0]

        nhoh = nmol.chains.residues.get(lambda x: x.type=="HOH")
        ohoh = self.mol.chains.residues.get(lambda x: x.type=="HOH")
        self.failUnless(len(nhoh)==0)
        os.system("rm Data/1bsr_writer.pdb")