import IMP
import IMP.test
import IMP.core
import IMP.atom

rk = IMP.FloatKey("radius")

class Test(IMP.test.TestCase):
    """Tests for bond refiner"""


    def test_rops(self):
        """Checking rigid distance pair score"""
        IMP.set_log_level(IMP.VERBOSE)
        m= IMP.Model()
        p0= IMP.atom.read_pdb(self.get_input_file_name("input.pdb"), m)
        p1= IMP.atom.read_pdb(self.get_input_file_name("input.pdb"), m)
        IMP.atom.add_radii(p0)
        IMP.atom.add_radii(p1)
        rb0= IMP.atom.create_rigid_body(p0)
        rb1= IMP.atom.create_rigid_body(p1)
        randt=IMP.algebra.Transformation3D(IMP.algebra.get_random_rotation_3d(), IMP.algebra.get_random_vector_in(IMP.algebra.BoundingBox3D(IMP.algebra.Vector3D(0,0,0), IMP.algebra.Vector3D(100,100,100))))
        IMP.core.RigidBody(rb0).set_transformation(randt)
        sdps= IMP.core.SphereDistancePairScore(IMP.core.Linear(0,1))
        rdps= IMP.core.RigidBodyDistancePairScore(sdps, IMP.core.LeavesRefiner(IMP.atom.Hierarchy.get_traits()))
        v= rdps.evaluate((p0.get_particle(), p1.get_particle()), None)
        dm= 1000000
        bp=None
        for l0 in IMP.core.get_leaves(p0):
            for l1 in IMP.core.get_leaves(p1):
                d= sdps.evaluate(IMP.ParticlePair(l0.get_particle(), l1.get_particle()), None)
                if d< dm:
                    print "found ", l0.get_name(), l1.get_name(), d
                    dm=d
        self.assertAlmostEqual(v, dm, delta=.1)


    def test_rops_against_one(self):
        """Checking rigid distance pair score against one"""
        IMP.set_log_level(IMP.VERBOSE)
        m= IMP.Model()
        p0= IMP.atom.read_pdb(self.get_input_file_name("input.pdb"), m)
        p1= IMP.Particle(m)
        randt=IMP.algebra.get_random_vector_in(IMP.algebra.BoundingBox3D(IMP.algebra.Vector3D(0,0,0), IMP.algebra.Vector3D(100,100,100)))
        IMP.core.XYZR.setup_particle(p1, IMP.algebra.Sphere3D(randt, 3))
        IMP.atom.add_radii(p0)
        rb=IMP.atom.create_rigid_body(p0)
        sdps= IMP.core.SphereDistancePairScore(IMP.core.Linear(0,1))
        rdps= IMP.core.RigidBodyDistancePairScore(sdps, IMP.core.LeavesRefiner(IMP.atom.Hierarchy.get_traits()))
        v= rdps.evaluate((p0, p1), None)
        dm= 1000000
        bp=None
        for l0 in IMP.core.get_leaves(p0):
            d= sdps.evaluate(IMP.ParticlePair(l0.get_particle(), p1), None)
            if d< dm:
                dm=d
        self.assertAlmostEqual(v, dm, delta=.1)


if __name__ == '__main__':
    IMP.test.main()