import numpy

from amuse.test import amusetest
from amuse.units import nbody_system
from amuse.units import units
from amuse.ic.plummer import new_plummer_model, MakePlummerModel

class TestData(amusetest.TestCase):
    pass
   
class TestPlummer(TestData):
    def test1(self):
        numpy.random.seed(0)
        #print numpy.random.get_state()
        m = MakePlummerModel(2)
        m1, p, v = m.new_model()
        self.assertEquals(m1[0,0], 0.5)
        self.assertEquals(m1[1,0], 0.5)
        self.assertAlmostEqual(p[0,0], -0.729636617171, 5)
        self.assertAlmostEqual(p[1,0], -0.713272921751 , 5)
        self.assertAlmostEqual(p[0,1],  0.379570256435, 5)
        self.assertAlmostEqual(p[1,1],  -0.930290757081, 5)
        
    def test2(self):
        convert_nbody = nbody_system.nbody_to_si(6|units.kg, 7 | units.m) 
        stars =  new_plummer_model(2, convert_nbody)
        self.assertEquals(stars[0].mass.value_in(units.kg), 3.0)
        self.assertEquals(stars[1].mass.value_in(units.kg), 3.0)
        
    def test3(self):
        stars =  new_plummer_model(2, None)
        self.assertEquals(stars[0].mass.value_in(nbody_system.mass), 0.5)
        self.assertEquals(stars[1].mass.value_in(nbody_system.mass), 0.5)
        
    def test4(self):
        stars = new_plummer_model(2, do_scale = True)
        self.assertAlmostEqual(stars.kinetic_energy(),             0.25 | nbody_system.energy)
        self.assertAlmostEqual(stars.potential_energy(G=nbody_system.G), -0.50 | nbody_system.energy)
        self.assertAlmostEqual(stars.center_of_mass(),          [0,0,0] | nbody_system.length)
        self.assertAlmostEqual(stars.center_of_mass_velocity(), [0,0,0] | nbody_system.speed)
        self.assertAlmostEqual(stars.mass.sum(),                   1.00 | nbody_system.mass)
        self.assertAlmostEqual(stars.virial_radius(),              1.00 | nbody_system.length)
        
