import numpy
import math

from amuse.support.testing.amusetest import TestWithMPI
from amuse.units import nbody_system, units, constants
from amuse.datamodel import Particles
from amuse_mi6.interface import MI6Interface, MI6

from amuse.ic.plummer import new_plummer_model


default_options = dict()


class TestMI6Interface(TestWithMPI):

    def test1(self):
        print("Test MI6Interface initialization")
        instance = MI6Interface(**default_options)
        self.assertEqual(0, instance.initialize_code())
        self.assertEqual(0, instance.commit_parameters())
        self.assertEqual(0, instance.cleanup_code())
        instance.stop()

    def test2(self):
        print("Test MI6Interface new_particle / get_state")
        instance = MI6Interface(**default_options)
        self.assertEqual(0, instance.initialize_code())
        self.assertEqual(0, instance.commit_parameters())

        id, error = instance.new_particle(mass=11.0, radius=2.0, x=0.0, y=0.0, z=0.0, vx=0.0, vy=0.0, vz=0.0)
        self.assertEqual(0, error)
        self.assertEqual(1, id)
        id, error = instance.new_particle(mass=21.0, radius=5.0, x=10.0, y=0.0, z=0.0, vx=10.0, vy=0.0, vz=0.0)
        self.assertEqual(0, error)
        self.assertEqual(2, id)
        self.assertEqual(0, instance.commit_particles())

        retrieved_state1 = instance.get_state(1)
        retrieved_state2 = instance.get_state(2)
        self.assertEqual(0,  retrieved_state1['__result'])
        self.assertEqual(0,  retrieved_state2['__result'])
        self.assertEqual(11.0,  retrieved_state1['mass'])
        self.assertEqual(21.0,  retrieved_state2['mass'])
        self.assertEqual(0.0,  retrieved_state1['x'])
        self.assertEqual(10.0,  retrieved_state2['x'])

        self.assertEqual(0, instance.cleanup_code())
        instance.stop()

    def xtest3(self):
        print("Test MI6Interface get_index_of_first_particle, get_index_of_next_particle")
        instance = MI6Interface(**default_options)
        instance.initialize_code()
        instance.commit_parameters()
        for i in [1, 2, 3]:
            result = instance.new_particle(mass=i, radius=1.0, x=0.0, y=0.0, z=0.0, vx=0.0, vy=0.0, vz=0.0)
            self.assertEqual(i, result['index_of_the_particle'])

        instance.commit_particles()
        self.assertEqual(1, instance.get_index_of_first_particle()['index_of_the_particle'])
        self.assertEqual(2, instance.get_index_of_next_particle(1)['index_of_the_next_particle'])
        self.assertEqual(3, instance.get_index_of_next_particle(2)['index_of_the_next_particle'])

        instance.delete_particle(1)

        self.assertEqual(2, instance.get_number_of_particles()['number_of_particles'])

        # the deletion does a swap, so 3 is copied to 1, (overwriting old 1 and treesize -> treesize-1
        self.assertEqual(3, instance.get_index_of_first_particle()['index_of_the_particle'])

        self.assertEqual(1, instance.get_index_of_next_particle(2)['__result'])

        instance.cleanup_code()
        instance.stop()

    def test4(self):
        print("Test MI6Interface particle property getters/setters")
        instance = MI6Interface(**default_options)
        self.assertEqual(0, instance.initialize_code())
        self.assertEqual(0, instance.commit_parameters())
        self.assertEqual([1, 0], list(instance.new_particle(0.01,  1, 0, 0,  0, 1, 0, 0.1).values()))
        self.assertEqual([2, 0], list(instance.new_particle(0.02, -1, 0, 0,  0, -1, 0, 0.1).values()))
        self.assertEqual(-3, instance.get_mass(1)['__result'])  # Have to commit first
        self.assertEqual(0, instance.commit_particles())

        # getters
        mass, result = instance.get_mass(1)
        self.assertAlmostEqual(0.01, mass)
        self.assertEqual(0, result)
        radius, result = instance.get_radius(2)
        self.assertAlmostEqual(0.1, radius)
        self.assertEqual(0, result)
        self.assertEqual(-3, instance.get_mass(3)['__result'])  # Particle not found
        self.assertEqual([1, 0, 0,  0], list(instance.get_position(1).values()))
        self.assertEqual([-1, 0, 0,  0], list(instance.get_position(2).values()))
        self.assertEqual([0, 1, 0,  0], list(instance.get_velocity(1).values()))
        self.assertEqual([0, -1, 0,  0], list(instance.get_velocity(2).values()))

        # setters
        self.assertEqual(0, instance.set_state(1, 0.01, 1, 2, 3, 4, 5, 6, 0.1))
        self.assertEqual([0.01, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.1, 0], list(instance.get_state(1).values()))
        self.assertEqual(0, instance.set_mass(1, 0.02))
        self.assertEqual([0.02, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.1, 0], list(instance.get_state(1).values()))
        self.assertEqual(0, instance.set_radius(1, 0.2))
        self.assertEqual([0.02, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.2, 0], list(instance.get_state(1).values()))
        self.assertEqual(0, instance.set_position(1, 10, 20, 30))
        self.assertEqual([0.02, 10.0, 20.0, 30.0, 4.0, 5.0, 6.0, 0.2, 0], list(instance.get_state(1).values()))
        self.assertEqual(0, instance.set_velocity(1, 40, 50, 60))
        self.assertEqual([0.02, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 0.2, 0], list(instance.get_state(1).values()))

        self.assertEqual(0, instance.cleanup_code())
        instance.stop()

    def test5(self):
        print("Test MI6Interface parameters")
        instance = MI6Interface(**default_options)
        self.assertEqual(0, instance.initialize_code())

        # MI6 has separate epsilon_squared parameters for different interactions!
        self.assertEqual([0, 0], list(instance.get_eps2_fs_fs().values()))
        self.assertEqual([0, 0], list(instance.get_eps2_fs_bh().values()))
        self.assertEqual([0, 0], list(instance.get_eps2_bh_bh().values()))
        self.assertEqual(-2, instance.get_eps2()['__result'])  # Not implemented (would be ambiguous)

        self.assertEqual(0, instance.set_eps2_fs_fs(0.2))
        self.assertEqual([0.2, 0], list(instance.get_eps2_fs_fs().values()))
        self.assertEqual(0, instance.set_eps2_fs_bh(0.3))
        self.assertEqual([0.3, 0], list(instance.get_eps2_fs_bh().values()))
        self.assertEqual(0, instance.set_eps2_bh_bh(0.4))
        self.assertEqual([0.4, 0], list(instance.get_eps2_bh_bh().values()))
        self.assertEqual(-2, instance.set_eps2(0.1))  # Not implemented (would be ambiguous)

        self.assertEqual([1.0e-4, 0], list(instance.get_eta_s().values()))
        self.assertEqual([0.1, 0], list(instance.get_eta_fs().values()))
        self.assertEqual([0.4, 0], list(instance.get_eta_smbh().values()))
        self.assertEqual([0.4, 0], list(instance.get_eta_imbh().values()))

        self.assertEqual(0, instance.set_eta_s(0.01))
        self.assertEqual([0.01, 0], list(instance.get_eta_s().values()))
        self.assertEqual(0, instance.set_eta_fs(0.02))
        self.assertEqual([0.02, 0], list(instance.get_eta_fs().values()))
        self.assertEqual(0, instance.set_eta_smbh(0.03))
        self.assertEqual([0.03, 0], list(instance.get_eta_smbh().values()))
        self.assertEqual(0, instance.set_eta_imbh(0.04))
        self.assertEqual([0.04, 0], list(instance.get_eta_imbh().values()))

        self.assertEqual(0, instance.set_max_relative_energy_error(1.0e-6))
        self.assertEqual([1.0e-6, 0], list(instance.get_max_relative_energy_error().values()))
        self.assertEqual(0, instance.set_maximum_timestep(1.0e-6))
        self.assertEqual([1.0e-6, 0], list(instance.get_maximum_timestep().values()))

        self.assertEqual(0, instance.set_smbh_mass(0.1))
        self.assertEqual([0.1, 0], list(instance.get_smbh_mass().values()))

        self.assertEqual([0, 0], list(instance.get_include_smbh_flag().values()))
        self.assertEqual(0, instance.set_include_smbh_flag(1))
        self.assertEqual([1, 0], list(instance.get_include_smbh_flag().values()))
        self.assertEqual([1, 0], list(instance.get_calculate_postnewtonian().values()))
        self.assertEqual(0, instance.set_calculate_postnewtonian(0))
        self.assertEqual([0, 0], list(instance.get_calculate_postnewtonian().values()))

        self.assertEqual(0, instance.commit_parameters())
        self.assertEqual(0, instance.cleanup_code())
        instance.stop()

    def test6(self):
        print("Test MI6Interface evolve_model, single particle (+SMBH)")
        instance = MI6Interface(**default_options)
        self.assertEqual(0, instance.initialize_code())
        self.assertEqual(0, instance.set_include_smbh_flag(1))
        self.assertEqual(0, instance.set_calculate_postnewtonian(0))
        self.assertEqual(0, instance.commit_parameters())

        # Set up a light particle on a circular orbit around the SMBH:
        self.assertEqual([1, 0], list(instance.new_particle(1.0e-4,  1.0, 0, 0,  0, 1.0, 0, 0.001).values()))
        self.assertEqual(0, instance.commit_particles())

        self.assertEqual(0, instance.evolve_model(math.pi))  # half an orbit
        for result, expected in zip(instance.get_position(1).values(), [-1.0, 0.0, 0.0, 0]):
            self.assertAlmostEqual(result, expected, 4)

        self.assertEqual(0, instance.evolve_model(2 * math.pi))  # full orbit
        for result, expected in zip(instance.get_position(1).values(), [1.0, 0.0, 0.0, 0]):
            self.assertAlmostEqual(result, expected, 4)

        self.assertEqual(0, instance.cleanup_code())
        instance.stop()

    def test7(self):
        print("Test MI6Interface evolve_model, 2 particles orbiting the SMBH")
        instance = MI6Interface(**default_options)
        self.assertEqual(0, instance.initialize_code())
        self.assertEqual(0, instance.set_include_smbh_flag(1))
        self.assertEqual(0, instance.set_calculate_postnewtonian(0))
        self.assertEqual(0, instance.commit_parameters())
        # Set up a light binary on circular orbits around the SMBH:
        mass = 1.0e-4
        dv = 0.5 * (math.sqrt(1.0 + 0.5*mass) - 1.0)
        self.assertEqual([1, 0], list(instance.new_particle(mass,  1.0, 0, 0,  0, 1.0+dv, 0, 0.001).values()))
        self.assertEqual([2, 0], list(instance.new_particle(mass, -1.0, 0, 0,  0, -1.0-dv, 0, 0.001).values()))
        self.assertEqual(0, instance.commit_particles())

        P = 2 * math.pi / (1 + dv)
        self.assertEqual(0, instance.evolve_model(P / 2))  # half an orbit
        for result, expected in zip(instance.get_position(1).values(), [-1.0, 0.0, 0.0, 0]):
            self.assertAlmostEqual(result, expected, 4)

        self.assertEqual(0, instance.evolve_model(P))  # full orbit
        for result, expected in zip(instance.get_position(1).values(), [1.0, 0.0, 0.0, 0]):
            self.assertAlmostEqual(result, expected, 4)

        self.assertEqual(0, instance.cleanup_code())
        instance.stop()

    def test8(self):
        print("Test MI6Interface evolve_model, binary (+SMBH) --> accretion!?")
        instance = MI6Interface(**default_options)
        self.assertEqual(0, instance.initialize_code())
        self.assertEqual(0, instance.set_include_smbh_flag(1))
        self.assertEqual(0, instance.set_calculate_postnewtonian(0))
        self.assertEqual(0, instance.commit_parameters())
        # Set up a light binary on circular orbits around the SMBH:
        mass = 1.0e-6
        self.assertEqual([1, 0], list(instance.new_particle(mass,  1.0, 0, 0,  0, 1.0, 0, 0.001).values()))
        self.assertEqual([2, 0], list(instance.new_particle(mass, -1.0, 0, 0,  0, -1.0, 0, 0.001).values()))
        self.assertEqual(0, instance.commit_particles())

        self.assertEqual(0, instance.evolve_model(0.00001))
        result1 = list(instance.get_position(1).values())
        result2 = list(instance.get_position(2).values())
        print(result1)
        print(result2)

        self.assertEqual(0, instance.cleanup_code())
        instance.stop()


class TestMI6(TestWithMPI):

    default_converter = nbody_system.nbody_to_si(1.0e4 | units.MSun, 1.0 | units.AU)

    def new_sun_earth_system(self):
        particles = Particles(2)
        particles.mass = [1.0, 3.0037e-6] | units.MSun
        particles.radius = 1.0 | units.RSun
        particles.position = [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]] | units.AU
        particles.velocity = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] | units.km / units.s
        particles[1].vy = (constants.G * particles.total_mass() / (1.0 | units.AU)).sqrt()
        return particles

    def test1(self):
        print("Testing MI6 initialization")
        instance = MI6(self.default_converter, **default_options)
        instance.initialize_code()
        instance.commit_parameters()
        instance.cleanup_code()
        instance.stop()

    def test2(self):
        print("Testing MI6 parameters")
        instance = MI6(self.default_converter, **default_options)
        instance.initialize_code()

        self.assertEqual(instance.parameters.epsilon_squared,
            instance.unit_converter.to_si(0.0 | nbody_system.length**2))
        self.assertEqual(instance.parameters.timestep_parameter, 0.1)

        for par, value in [('epsilon_squared_star_star', 0.0 | nbody_system.length**2),
                ('epsilon_squared_star_blackhole', 0.0 | nbody_system.length**2),
                ('epsilon_squared_blackhole_blackhole', 0.0 | nbody_system.length**2),
                ('initial_timestep_parameter', 1.0e-4),
                ('timestep_parameter_stars', 0.1),
                ('timestep_parameter_supermassive_black_holes', 0.4),
                ('timestep_parameter_intermediate_mass_black_holes', 0.4),
                ('max_relative_energy_error', 5.0e-5),
                ('maximum_timestep', 1.0/1024.0 | nbody_system.time),
                ('smbh_mass', 1.0 | nbody_system.mass),
                ('lightspeed', 1.0 | nbody_system.speed)]:
            self.assertEqual(instance.unit_converter.to_si(value),
                getattr(instance.parameters, par))

            if hasattr(value, 'unit'):
                new_value = 3.0 | value.unit
            else:
                new_value = 3.0

            setattr(instance.parameters, par, new_value)
            self.assertEqual(instance.unit_converter.to_si(new_value),
                getattr(instance.parameters, par))

        # epsilon_squared is an alias for epsilon_squared_star_star, so epsilon_squared also has become 3:
        self.assertEqual(instance.parameters.epsilon_squared,
            instance.unit_converter.to_si(3.0 | nbody_system.length**2))
        instance.parameters.epsilon_squared = 0.1 | nbody_system.length**2
        self.assertEqual(instance.parameters.epsilon_squared,
            instance.unit_converter.to_si(0.1 | nbody_system.length**2))
        # timestep_parameter is an alias for timestep_parameter_stars, so timestep_parameter also has become 3:
        self.assertEqual(instance.parameters.timestep_parameter, 3.0)
        instance.parameters.timestep_parameter = 0.01
        self.assertEqual(instance.parameters.timestep_parameter, 0.01)

        self.assertEqual(instance.parameters.include_smbh, False)
        instance.parameters.include_smbh = True
        self.assertEqual(instance.parameters.include_smbh, True)
        self.assertEqual(instance.parameters.calculate_postnewtonian, True)
        instance.parameters.calculate_postnewtonian = False
        self.assertEqual(instance.parameters.calculate_postnewtonian, False)

        self.assertEqual(instance.parameters.drink, "Vodka martini. Shaken, not stirred.")

        instance.stop()

    def test3(self):
        print("Testing MI6 particles")
        instance = MI6(self.default_converter, **default_options)
        instance.initialize_code()
        instance.commit_parameters()
        instance.particles.add_particles(self.new_sun_earth_system())
        instance.commit_particles()

        self.assertAlmostEqual(instance.particles.mass, [1.0, 3.0037e-6] | units.MSun)
        self.assertAlmostEqual(instance.particles.radius, 1.0 | units.RSun)
        self.assertAlmostEqual(instance.particles.position,
            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]] | units.AU)
        self.assertAlmostEqual(instance.particles.velocity,
            [[0.0, 0.0, 0.0], [0.0, 29.7885, 0.0]] | units.km / units.s, 3)

        instance.cleanup_code()
        instance.stop()

    def test4(self):
        print("Testing MI6 evolve_model, 2 particles orbiting the SMBH")
        particles = Particles(2)
        particles.mass = 1.0 | units.MSun
        particles.radius = 1.0 | units.RSun
        particles.position = [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]] | units.AU
        particles.velocity = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] | units.km / units.s
        particles[1].vy = ((constants.G * (10001.0 | units.MSun) / (1.0 | units.AU)).sqrt() +
                           (constants.G * (10000.0 | units.MSun) / (1.0 | units.AU)).sqrt())
        particles.move_to_center()
        print(particles)

        instance = MI6(self.default_converter, **default_options)
        instance.initialize_code()
        instance.parameters.include_smbh = True
        instance.parameters.lightspeed = constants.c
        instance.commit_parameters()
        instance.particles.add_particles(particles)
        instance.commit_particles()
        primary = instance.particles[0]

        P = 2 * math.pi * primary.x / primary.vy
        P_corrected = (P) * (2.0 / (1.0 + math.sqrt(1.0001)))

        position_at_start = primary.position.x
        instance.evolve_model(P_corrected / 4.0)
        self.assertAlmostRelativeEqual(position_at_start, primary.position.y, 3)

        instance.evolve_model(P_corrected / 2.0)
        self.assertAlmostRelativeEqual(position_at_start, -primary.position.x, 3)

        instance.evolve_model(P_corrected)
        self.assertAlmostRelativeEqual(position_at_start, primary.position.x, 3)

        instance.cleanup_code()
        instance.stop()

    def test5(self):
        print("Testing MI6 evolve_model, 2 particles, no SMBH")
        particles = Particles(2)
        particles.mass = 1.0 | units.MSun
        particles.radius = 1.0 | units.RSun
        particles.position = [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]] | units.AU
        particles.velocity = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] | units.km / units.s
        particles[1].vy = (constants.G * (2.0 | units.MSun) / (2.0 | units.AU)).sqrt()
        particles.move_to_center()
        print(particles)

        converter = nbody_system.nbody_to_si(1.0 | units.MSun, 1.0 | units.AU)
        instance = MI6(converter, **default_options)
        instance.initialize_code()
        instance.parameters.smbh_mass = 0.0 | units.MSun
        instance.commit_parameters()
        instance.particles.add_particles(particles)
        instance.commit_particles()
        primary = instance.particles[0]

        P = 2 * math.pi * primary.x / primary.vy

        position_at_start = primary.position.x
        instance.evolve_model(P / 4.0)
        self.assertAlmostRelativeEqual(position_at_start, primary.position.y, 3)

        instance.evolve_model(P / 2.0)
        self.assertAlmostRelativeEqual(position_at_start, -primary.position.x, 3)

        instance.evolve_model(P)
        self.assertAlmostRelativeEqual(position_at_start, primary.position.x, 3)

        instance.cleanup_code()
        instance.stop()

    def test6(self):
        print("Testing MI6 evolve_model, earth-sun system, no SMBH")
        converter = nbody_system.nbody_to_si(1.0 | units.MSun, 1.0 | units.AU)
        instance = MI6(converter, **default_options)
        instance.initialize_code()
        instance.parameters.smbh_mass = 0.0 | units.MSun
        instance.commit_parameters()
        instance.particles.add_particles(self.new_sun_earth_system())
        instance.commit_particles()
        earth = instance.particles[1]

        position_at_start = earth.position.x
        instance.evolve_model(0.25 | units.yr)
        self.assertAlmostRelativeEqual(position_at_start, earth.position.y, 3)

        instance.evolve_model(0.5 | units.yr)
        self.assertAlmostRelativeEqual(position_at_start, -earth.position.x, 3)

        instance.evolve_model(1.0 | units.yr)
        self.assertAlmostRelativeEqual(position_at_start, earth.position.x, 3)

        instance.cleanup_code()
        instance.stop()

    def test7(self):
        print("Testing effect of MI6 parameter epsilon_squared")
        converter = nbody_system.nbody_to_si(1.0 | units.MSun, 1.0 | units.AU)
        particles = self.new_sun_earth_system()
        particles.rotate(0.0, 0.0, -math.pi/4)
        particles.move_to_center()

        tan_initial_direction = particles[1].vy/particles[1].vx
        self.assertAlmostEqual(tan_initial_direction, math.tan(math.pi/4))
        tan_final_direction = []
        for log_eps2 in range(-9, 10, 2):
            instance = MI6(converter, **default_options)
            instance.initialize_code()
            instance.parameters.epsilon_squared = 10.0**log_eps2 | units.AU ** 2
            instance.parameters.smbh_mass = 0.0 | units.MSun
            instance.commit_parameters()
            instance.particles.add_particles(particles)
            instance.commit_particles()
            instance.evolve_model(0.25 | units.yr)
            tan_final_direction.append(instance.particles[1].velocity[1] /
                instance.particles[1].velocity[0])
            instance.cleanup_code()
            instance.stop()
        # Small values of epsilon_squared should result in normal earth-sun dynamics: rotation of 90 degrees
        self.assertAlmostEqual(tan_final_direction[0], math.tan(3 * math.pi / 4.0), 2)
        # Large values of epsilon_squared should result in ~ no interaction
        self.assertAlmostEqual(tan_final_direction[-1], tan_initial_direction, 2)
        # Outcome is most sensitive to epsilon_squared when epsilon_squared = d(earth, sun)^2
        delta = [abs(tan_final_direction[i+1]-tan_final_direction[i]) for i in range(len(tan_final_direction)-1)]
        self.assertEqual(delta[len(tan_final_direction)//2 - 1], max(delta))

    def test8(self):
        print("Testing MI6 get_gravity_at_point and get_potential_at_point")
        instance = MI6(**default_options)
        instance.initialize_code()
        instance.parameters.epsilon_squared = 0.0 | nbody_system.length**2
        instance.parameters.smbh_mass = 0.0 | nbody_system.mass

        particles = Particles(2)
        particles.mass = 1.0 | nbody_system.mass
        particles.radius = 0.0 | nbody_system.length
        particles.position = [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]] | nbody_system.length
        particles.velocity = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] | nbody_system.speed
        instance.particles.add_particles(particles)

        zero = 0.0 | nbody_system.length
        fx, fy, fz = instance.get_gravity_at_point(zero, 1.0 | nbody_system.length, zero, zero)
        self.assertAlmostEqual(fx, 0.0 | nbody_system.acceleration)
        self.assertAlmostEqual(fy, 0.0 | nbody_system.acceleration)
        self.assertAlmostEqual(fz, 0.0 | nbody_system.acceleration)

        for x in (0.25, 0.5, 0.75):
            x0 = x | nbody_system.length
            x1 = (2.0 - x) | nbody_system.length
            potential0 = instance.get_potential_at_point(zero, x0, zero, zero)
            potential1 = instance.get_potential_at_point(zero, x1, zero, zero)
            fx0, fy0, fz0 = instance.get_gravity_at_point(zero, x0, zero, zero)
            fx1, fy1, fz1 = instance.get_gravity_at_point(zero, x1, zero, zero)

            self.assertAlmostEqual(fy0, 0.0 | nbody_system.acceleration)
            self.assertAlmostEqual(fz0, 0.0 | nbody_system.acceleration)
            self.assertAlmostEqual(fy1, 0.0 | nbody_system.acceleration)
            self.assertAlmostEqual(fz1, 0.0 | nbody_system.acceleration)

            self.assertAlmostEqual(fx0, -1.0 * fx1)
            fx = (-1.0 / (x0**2) + 1.0 / (x1**2)) * (1.0 | nbody_system.length ** 3 / nbody_system.time ** 2)
            self.assertAlmostEqual(fx, fx0)
            self.assertAlmostEqual(potential0, potential1)

        instance.stop()

    def test9(self):
        print("Testing MI6 evolve_model and getters energy, plummer sphere, no SMBH")
        converter = nbody_system.nbody_to_si(1.0e2 | units.MSun, 1.0 | units.parsec)
        instance = MI6(converter, **default_options)
        instance.initialize_code()
        instance.parameters.smbh_mass = 0.0 | units.MSun
        instance.commit_parameters()
        numpy.random.seed(987654321)
        instance.particles.add_particles(new_plummer_model(100, convert_nbody=converter))
        instance.commit_particles()

        kinetic_energy = instance.kinetic_energy
        potential_energy = instance.potential_energy
        self.assertAlmostRelativeEqual(kinetic_energy, 2.12292810174e+37 | units.J, 10)
        self.assertAlmostRelativeEqual(potential_energy, -4.33511391248e+37 | units.J, 10)

        initial_total_energy = kinetic_energy + potential_energy
        instance.evolve_model(0.1 | nbody_system.time)
        kinetic_energy = instance.kinetic_energy
        potential_energy = instance.potential_energy
        self.assertAlmostRelativeEqual(kinetic_energy, 2.13633848369e+37 | units.J, 5)
        self.assertAlmostRelativeEqual(potential_energy, -4.34851806763e+37 | units.J, 5)

        self.assertAlmostRelativeEqual(potential_energy + kinetic_energy, initial_total_energy, 5)

        instance.cleanup_code()
        instance.stop()

    def test10(self):
        print("Testing MI6 collision_detection")
        particles = Particles(7)
        particles.mass = 0.00000001 | nbody_system.mass
        particles.radius = 0.0001 | nbody_system.length
        particles.x = [-101.0, -100.0, -0.5, 0.5, 100.0, 101.0, 104.0] | nbody_system.length
        particles.y = 0 | nbody_system.length
        particles.z = 0 | nbody_system.length
        particles.velocity = [[2, 0, 0], [-2, 0, 0]]*3 + [[-4, 0, 0]] | nbody_system.speed

        instance = MI6(**default_options)
        instance.initialize_code()
        instance.parameters.set_defaults()
        instance.particles.add_particles(particles)
        collisions = instance.stopping_conditions.collision_detection
        collisions.enable()
        instance.evolve_model(1.0 | nbody_system.time)

        self.assertTrue(collisions.is_set())
        self.assertAlmostEqual(instance.model_time, -(particles[0].x-particles[1].x)/(particles[0].vx-particles[1].vx), 3)
        self.assertEqual(len(collisions.particles(0)), 3)
        self.assertEqual(len(collisions.particles(1)), 3)
        self.assertEqual(len(particles - collisions.particles(0) - collisions.particles(1)), 1)
        self.assertEqual(abs(collisions.particles(0).x - collisions.particles(1).x) <
                (collisions.particles(0).radius + collisions.particles(1).radius),
                [True, True, True])

        sticky_merged = Particles(len(collisions.particles(0)))
        sticky_merged.mass = collisions.particles(0).mass + collisions.particles(1).mass
        sticky_merged.radius = collisions.particles(0).radius
        for p1, p2, merged in zip(collisions.particles(0), collisions.particles(1), sticky_merged):
            merged.position = (p1 + p2).center_of_mass()
            merged.velocity = (p1 + p2).center_of_mass_velocity()

        print(instance.model_time)
        print(instance.particles)
        instance.particles.remove_particles(collisions.particles(0) + collisions.particles(1))
        instance.particles.add_particles(sticky_merged)

        instance.evolve_model(1.0 | nbody_system.time)
        print()
        print(instance.model_time)
        print(instance.particles)
        self.assertTrue(collisions.is_set())
        self.assertAlmostEqual(instance.model_time, -(particles[6].x-0.5*(particles[4].x+particles[5].x))/particles[6].vx, 3)
        self.assertEqual(len(collisions.particles(0)), 1)
        self.assertEqual(len(collisions.particles(1)), 1)
        self.assertEqual(len(instance.particles - collisions.particles(0) - collisions.particles(1)), 2)
        self.assertEqual(abs(collisions.particles(0).x - collisions.particles(1).x) <
                (collisions.particles(0).radius + collisions.particles(1).radius),
                [True])
        instance.stop()

    def test11(self):
        print("Testing MI6 properties")
        numpy.random.seed(12345)
        particles = new_plummer_model(100, do_scale=True)
        particles.position += [1, 2, 3] | nbody_system.length
        cluster_velocity = [4, 5, 6] | nbody_system.speed
        particles.velocity += cluster_velocity
        external_kinetic_energy = (0.5 | nbody_system.mass) * cluster_velocity.length_squared()

        instance = MI6(**default_options)
        instance.particles.add_particles(particles)

        kinetic_energy = instance.kinetic_energy - external_kinetic_energy
        potential_energy = instance.potential_energy
        self.assertAlmostRelativeEqual(kinetic_energy, 0.25 | nbody_system.energy, 10)
        self.assertAlmostRelativeEqual(potential_energy, -0.5 | nbody_system.energy, 10)
        self.assertAlmostRelativeEqual(instance.total_mass, 1.0 | nbody_system.mass, 10)
        self.assertAlmostRelativeEqual(instance.center_of_mass_position,
            [1, 2, 3] | nbody_system.length, 10)
        self.assertAlmostRelativeEqual(instance.center_of_mass_velocity,
            [4, 5, 6] | nbody_system.speed, 10)
        initial_total_energy = kinetic_energy + potential_energy

        instance.evolve_model(0.1 | nbody_system.time)
        self.assertAlmostRelativeEqual(instance.model_time, 0.1 | nbody_system.time)
        kinetic_energy = instance.kinetic_energy - external_kinetic_energy
        potential_energy = instance.potential_energy
        self.assertAlmostRelativeEqual(kinetic_energy+potential_energy, -0.25 | nbody_system.energy, 3)
        self.assertAlmostRelativeEqual(instance.total_mass, 1.0 | nbody_system.mass, 3)
        self.assertAlmostRelativeEqual(instance.center_of_mass_position,
            [1.4, 2.5, 3.6] | nbody_system.length, 3)
        self.assertAlmostRelativeEqual(instance.center_of_mass_velocity,
            [4, 5, 6] | nbody_system.speed, 3)

        instance.cleanup_code()
        instance.stop()
