# NOTE: This file is generated from modules/core/tools/container_tempates/test.py
# do not edit.

import unittest
import IMP
import IMP.test
import IMP.core
import IMP.container
import math

def compare(a,b):
    return cmp(a.get_name(), b.get_name())

# This file is generated by the make-container script

class ParticleContainerTest(IMP.test.TestCase):
    """Tests for SingletonContainer related objects"""

    def create_particle(self,m):
        p= IMP.Particle(m)
        d=IMP.core.XYZ.setup_particle(p)
        d.set_coordinates(IMP.algebra.get_random_vector_in(IMP.algebra.get_unit_bounding_box_3d()))
        p.add_attribute(IMP.FloatKey("thekey"), d.get_x())
        return p

    def create_particle_pair(self,m):
        p0= IMP.Particle(m)
        p1= IMP.Particle(m)
        d0= IMP.core.XYZ.setup_particle(p0)
        d1= IMP.core.XYZ.setup_particle(p1)
        d0.set_coordinates(IMP.algebra.get_random_vector_in(IMP.algebra.get_unit_bounding_box_3d()))
        d1.set_coordinates(IMP.algebra.get_random_vector_in(IMP.algebra.get_unit_bounding_box_3d()))
        return IMP.ParticlePair(p0,p1)

    def create_singleton_score(self):
        uf= IMP.core.Linear(0,1)
        return IMP.core.AttributeSingletonScore(uf,IMP.FloatKey("thekey"))

    def create_pair_score(self):
        uf= IMP.core.Linear(0,1)
        return IMP.core.DistancePairScore(uf)

    def create_singleton_restraint(self, ps,p):
        return IMP.core.SingletonRestraint(ps, p)

    def create_pair_restraint(self, ps, p):
        return IMP.core.PairRestraint(ps, IMP.ParticlePair(p[0], p[1]))


    def test_restraint(self):
        """Test the SingletonsRestraint"""
        m= IMP.Model()
        gs=self.create_singleton_score()
        c= IMP.container.ListSingletonContainer()
        f=0
        for i in range(0,10):
            p=self.create_particle(m)
            f= f+ gs.evaluate(p, None)
            c.add_particle(p)
            r= IMP.container.SingletonsRestraint(gs, c)
            r.set_was_used(True)
        m.add_restraint(r)
        self.assertInTolerance(m.evaluate(False), f, .1*f)
        p=self.create_particle(m)
        f= f+ gs.evaluate(p,None)
        c.add_particle(p)
        self.assertInTolerance(m.evaluate(False), f, .1*f)


    def test_irestraint(self):
        """Test the incremental evaluation of the SingletonsRestraint"""
        m= IMP.Model()
        m.set_log_level(IMP.TERSE)
        gs=self.create_singleton_score()
        c= IMP.container.ListSingletonContainer()
        ps=IMP.Particles()
        ps2= IMP.Particles()
        f=0
        for i in range(0,20):
            p=self.create_particle(m)
            ps.append(p)
            f=f+gs.evaluate(p,None)
        for i in range(0,10):
            p=self.create_particle(m)
            ps2.append(p)
            #f=f+evaluate_singleton_score(gs, p)
        c.set_particles(ps)
        r= IMP.container.SingletonsRestraint(gs, c)
        m.add_restraint(r)
        m.set_is_incremental(True)
        self.assertInTolerance(m.evaluate(False), f, .1*f)
        self.assertInTolerance(m.evaluate(False), f, .1*f)
        ps= ps+ps2
        f=0
        for p in ps:
            f=f+ gs.evaluate(p,None)
        c.set_particles(ps)
        self.assertInTolerance(m.evaluate(False), f, .1*f)


    def test_srestraint(self):
        """Test the SingletonRestraint"""
        m= IMP.Model()
        gs=self.create_singleton_score()
        p=self.create_particle(m)
        f= gs.evaluate(p, None)
        r= self.create_singleton_restraint(gs, p)
        r.set_was_used(True)
        m.add_restraint(r)
        self.assertInTolerance(m.evaluate(False), f, .1*f)

    def test_min_restraint(self):
        """Test the MinimumSingletonRestraint"""
        m= IMP.Model()
        c= IMP.container.ListSingletonContainer()
        self.assertEqual(c.get_ref_count(), 1)
        for i in range(0,10):
            c.add_particle(self.create_particle(m))
        print c.get_number_of_particles()
        d= self.create_singleton_score()
        self.assertEqual(d.get_ref_count(), 1)
        r= IMP.container.MinimumSingletonRestraint(d, c)
        self.assertEqual(d.get_ref_count(), 2)
        self.assertEqual(c.get_ref_count(), 2)
        r.set_n(4)
        m.add_restraint(r)
        f= m.evaluate(False)
        print f
        ms= []
        print c.get_number_of_particles()
        for i in range(0,10):
            ps= c.get_particle(i)
            cm= d.evaluate(ps, None)
            ms.append(cm)
        print ms
        ms.sort()
        mt=0
        for i in range(0, 4):
            mt = mt+ ms[i]
        print mt
        self.assertInTolerance(mt, f, .1*f)

    def test_max_restraint(self):
        """Test the MaximumSingletonRestraint"""
        m= IMP.Model()
        c= IMP.container.ListSingletonContainer()
        self.assertEqual(c.get_ref_count(), 1)
        for i in range(0,10):
            c.add_particle(self.create_particle(m))
        print c.get_number_of_particles()
        d= self.create_singleton_score()
        self.assertEqual(d.get_ref_count(), 1)
        r= IMP.container.MaximumSingletonRestraint(d, c)
        self.assertEqual(c.get_ref_count(), 2)
        self.assertEqual(d.get_ref_count(), 2)
        r.set_n(4)
        m.add_restraint(r)
        f= m.evaluate(False)
        print f
        ms= []
        print c.get_number_of_particles()
        for i in range(0,10):
            ps= c.get_particle(i)
            cm= d.evaluate(ps, None)
            ms.append(cm)
        ms.sort()
        print ms
        mt=0
        for i in range(0, 4):
            mt = mt+ ms[-i-1]
        print mt
        self.assertInTolerance(mt, f, .1*f)
    def test_max_score(self):
        """Test the MaximumSingletonScore"""
        m= IMP.Model()
        s= IMP.SingletonScoresTemp()
        for i in range(0,5):
            s.append(IMP.test.ConstSingletonScore(i))
        ps= IMP.container.MaximumSingletonScore(s, 2)
        p= self.create_particle(m)
        ps.set_was_used(True)
        v= ps.evaluate(p, None)
        self.assertEqual(v, 7)
    def test_min_score(self):
        """Test the MinimumSingletonScore"""
        m= IMP.Model()
        s= IMP.SingletonScoresTemp()
        for i in range(0,5):
            s.append(IMP.test.ConstSingletonScore(i))
        ps= IMP.container.MinimumSingletonScore(s, 3)
        p= self.create_particle(m)
        ps.set_was_used(True)
        v= ps.evaluate(p, None)
        self.assertEqual(v, 3)

    def test_set(self):
        """Testing SingletonContainerSet"""
        IMP.set_log_level(IMP.VERBOSE)
        m= IMP.Model()
        print "hi"
        c= IMP.container.SingletonContainerSet()
        ls=[]
        cs=[]
        for i in range(0,3):
            l= IMP.container.ListSingletonContainer()
            c.add_singleton_container(l)
            for j in range(0,3):
                t=self.create_particle(m)
                l.add_particle(t)
                cs.append(t)
        for p in cs:
            self.assert_(c.get_contains_particle(p))
        ret=[]
        for i in range(0, len(cs)):
            ret.append(c.get_particle(i))
        ret.sort(compare)
        #print ret
        cs.sort(compare)
        #rint cs
        for i in range(0, len(ret)):
            self.assertEqual(compare(ret[i], cs[i]), 0)
        self.assertEqual(c.get_number_of_particles(), len(cs))

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