from amuse.test import amusetest

from amuse.units import units
from amuse.units import constants
from amuse.units import nbody_system
import numpy
import time
import sys

from amuse.datamodel import Particles
from amuse.datamodel import Particle
from amuse.couple import encounters


# codes to use
from amuse.community.kepler.interface import Kepler
from amuse.community.smalln.interface import SmallN


def new_binary(
        mass1, mass2, semi_major_axis,
        eccentricity=0, keyoffset=1,
        is_at_periapsis=True,
        G=nbody_system.G
    ):
    total_mass = mass1 + mass2
    mass_fraction_particle_1 = mass1 / (total_mass)

    binary = Particles(keys=range(keyoffset, keyoffset+2))
    binary[0].mass = mass1
    binary[1].mass = mass2
    binary.child1 = None
    binary.child2 = None

    mu = G * total_mass

    if is_at_periapsis:
        velocity = numpy.sqrt(mu / semi_major_axis * ((1.0 + eccentricity)/(1.0 - eccentricity)))
        radius = semi_major_axis * (1.0 - eccentricity)
    else:
        velocity = numpy.sqrt(mu / semi_major_axis * ((1.0 - eccentricity)/(1.0 + eccentricity)))
        radius = semi_major_axis * (1.0 + eccentricity)

    binary[0].position = ((1.0 - mass_fraction_particle_1) * radius * [1.0, 0.0, 0.0])
    binary[1].position = -(mass_fraction_particle_1 * radius * [1.0, 0.0, 0.0])

    binary[0].velocity = ((1.0 - mass_fraction_particle_1) * velocity * [0.0, 1.0, 0.0])
    binary[1].velocity = -(mass_fraction_particle_1 * velocity * [0.0, 1.0, 0.0])

    return binary


class TryHandleEncounter(encounters.AbstractHandleEncounter, encounters.SelectNeighboursByDistanceMixin):

    def __init__(self,
        kepler_code,
        G=nbody_system.G
    ):
        encounters.AbstractHandleEncounter.__init__(
            self,
            kepler_code,
            G
        )
        encounters.SelectNeighboursByDistanceMixin.__init__(self)


class TestAbstractHandleEncounter(amusetest.TestWithMPI):

    def new_kepler(self):
        x = Kepler()
        x.initialize_code()
        return x

    def test1(self):
        particles_in_encounter = Particles(2)
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.5, 0] | nbody_system.speed
        particles_in_encounter.radius = 0 | nbody_system.length

        particles_in_field = Particles(5)
        particles_in_field.mass = [1, 2, 3, 4, 5] | nbody_system.mass
        particles_in_field[0].position = [2, 0, 0] | nbody_system.length
        particles_in_field[1].position = [1.5, 0, 0] | nbody_system.length
        particles_in_field[2].position = [0.5, 1, 0] | nbody_system.length
        particles_in_field[3].position = [0.5, -0.5, 0] | nbody_system.length
        particles_in_field[4].position = [0, 0, 2] | nbody_system.length
        particles_in_field.velocity = [0, 0.0, 0] | nbody_system.speed
        particles_in_field.radius = 0 | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.particles_in_field.add_particles(particles_in_field)

        x.execute()

        self.assertAlmostRelativeEqual(x.large_scale_of_particles_in_the_encounter, 1.0 | nbody_system.length)
        self.assertAlmostRelativeEqual(x.small_scale_of_particles_in_the_encounter, 3.0 | nbody_system.length)

        self.assertEqual(len(x.particles_close_to_encounter), 3)
        self.assertAlmostRelativeEqual(x.particles_close_to_encounter.mass, [2, 3, 4] | nbody_system.mass)

    def test2(self):

        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.5, 0] | nbody_system.speed
        particles_in_encounter.radius = 0 | nbody_system.length

        particles_in_field = Particles(keys=(5,))
        particles_in_field.mass = 2 | nbody_system.mass
        particles_in_field.position = [-0.5, 0, 0] | nbody_system.length
        particles_in_field.velocity = [0, 0.5, 0] | nbody_system.speed
        particles_in_field.radius = 0 | nbody_system.length

        particles_in_multiples = Particles()
        particles_in_multiples.add_particle(particles_in_encounter[0])
        multiple = particles_in_multiples[0]
        multiple.components = Particles(keys=(3, 4))
        multiple.components.mass = 0.5 | nbody_system.mass
        multiple.components[0].position = [0, 0.1, 0] | nbody_system.length
        multiple.components[1].position = [0, -0.1, 0] | nbody_system.length
        multiple.components[0].velocity = [0, 0, 0.2] | nbody_system.speed
        multiple.components[1].velocity = [0, 0.1, -0.2] | nbody_system.speed
        multiple.components.child1 = None
        multiple.components.child2 = None

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.particles_in_field.add_particles(particles_in_field)
        x.existing_multiples.add_particles(particles_in_multiples)

        x.execute()

        self.assertEqual(len(x.all_singles_in_encounter), 3)
        self.assertEqual(len(x.all_singles_in_evolve), 4)

        self.assertAlmostRelativeEqual(x.all_singles_in_evolve.mass, [0.5, 0.5, 1.0, 2.0] | nbody_system.mass)

        self.assertEqual(len(x.released_singles), 2)
        print(x.particles_after_encounter)
        child1 = x.particles_after_encounter[-2]
        child2 = x.particles_after_encounter[-1]
        print(child1.position)
        print(child2.position)
        self.assertAlmostRelativeEqual(child1.position, [1, 0.1, 0] | nbody_system.length)
        self.assertAlmostRelativeEqual(child1.velocity, [0, 0.5, 0.2] | nbody_system.speed)
        self.assertAlmostRelativeEqual(child2.position, [1, -0.1, 0] | nbody_system.length)
        self.assertAlmostRelativeEqual(child2.velocity, [0, 0.6, -0.2] | nbody_system.speed)

    def test3(self):

        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.0, 0] | nbody_system.speed
        particles_in_encounter.radius = 0 | nbody_system.length

        particles_in_field = Particles(keys=(2, 4, 5, 6, 7,))
        particles_in_field.mass = [1, 2, 3, 4, 5] | nbody_system.mass
        particles_in_field[0].position = [2, 0, 0] | nbody_system.length
        particles_in_field[1].position = [1.5, 0, 0] | nbody_system.length
        particles_in_field[2].position = [0.5, 1, 0] | nbody_system.length
        particles_in_field[3].position = [0.5, -0.5, 0] | nbody_system.length
        particles_in_field[4].position = [0, 0, 2] | nbody_system.length
        particles_in_field.velocity = [0, 0, 0] | nbody_system.speed
        particles_in_field.radius = 0 | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.particles_in_field.add_particles(particles_in_field)

        simple_binary = new_binary(
            1 | nbody_system.mass,
            1 | nbody_system.mass,
            0.2 | nbody_system.length
        )

        def evolve_singles_in_encounter_until_end_state():
            particles = x.singles_and_multiples_after_evolve
            particles.add_particles(x.all_singles_in_encounter)
            particles.child1 = None
            particles.child2 = None

            root_particle = particles.add_particle(Particle(
                key=10,
                mass=2.0 | nbody_system.mass,
                position=particles[0:2].center_of_mass(),
                velocity=particles[0:2].center_of_mass_velocity(),
            ))
            root_particle.child1 = particles[0]
            root_particle.child2 = particles[1]
            particles[0].position = simple_binary[0].position + root_particle.position
            particles[1].position = simple_binary[1].position + root_particle.position

            particles[0].velocity = simple_binary[0].velocity + root_particle.velocity
            particles[1].velocity = simple_binary[1].velocity + root_particle.velocity

        x.evolve_singles_in_encounter_until_end_state = evolve_singles_in_encounter_until_end_state
        x.determine_structure_of_the_evolved_state = lambda: 1

        x.execute()
        self.assertEqual(len(x.new_multiples), 1)
        self.assertEqual(len(x.new_binaries), 1)
        multiple = x.new_multiples[0]
        self.assertEqual(len(multiple.components), 2)

        self.assertAlmostRelativeEqual(multiple.components[0].position, simple_binary[0].position)
        self.assertAlmostRelativeEqual(multiple.components[1].position, simple_binary[1].position)

    def test4(self):

        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.0, 0] | nbody_system.speed
        particles_in_encounter.radius = 0 | nbody_system.length

        particles_in_field = Particles(keys=(3, 4, 5, 6, 7,))
        particles_in_field.mass = [1, 2, 3, 4, 5] | nbody_system.mass
        particles_in_field[0].position = [2, 0, 0] | nbody_system.length
        particles_in_field[1].position = [1.5, 0, 0] | nbody_system.length
        particles_in_field[2].position = [0.5, 1, 0] | nbody_system.length
        particles_in_field[3].position = [0.5, -0.5, 0] | nbody_system.length
        particles_in_field[4].position = [0, 0, 2] | nbody_system.length
        particles_in_field.velocity = [0, 0, 0] | nbody_system.speed
        particles_in_field.radius = 0 | nbody_system.length

        binaries = Particles(keys=(20,))
        binaries[0].child1 = particles_in_encounter[0]
        binaries[0].child2 = particles_in_encounter[1]

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.particles_in_field.add_particles(particles_in_field)
        x.existing_binaries.add_particles(binaries)

        simple_binary = new_binary(
            1 | nbody_system.mass,
            1 | nbody_system.mass,
            0.2 | nbody_system.length
        )

        def evolve_singles_in_encounter_until_end_state():
            particles = x.singles_and_multiples_after_evolve
            particles.add_particles(x.all_singles_in_encounter)
            particles.child1 = None
            particles.child2 = None

            root_particle = particles.add_particle(Particle(
                key=10,
                mass=2.0 | nbody_system.mass,
                position=particles[0:2].center_of_mass(),
                velocity=particles[0:2].center_of_mass_velocity(),
            ))
            root_particle.child1 = particles[0]
            root_particle.child2 = particles[1]
            particles[0].position = simple_binary[0].position + root_particle.position
            particles[1].position = simple_binary[1].position + root_particle.position

            particles[0].velocity = simple_binary[0].velocity + root_particle.velocity
            particles[1].velocity = simple_binary[1].velocity + root_particle.velocity

        x.evolve_singles_in_encounter_until_end_state = evolve_singles_in_encounter_until_end_state
        x.determine_structure_of_the_evolved_state = lambda: 1

        x.execute()
        self.assertEqual(len(x.new_multiples), 1)
        self.assertEqual(len(x.new_binaries), 0)
        self.assertEqual(len(x.updated_binaries), 1)
        self.assertEqual(x.updated_binaries[0], binaries[0])

    def test5(self):

        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.0, 0] | nbody_system.speed
        particles_in_encounter.radius = 0 | nbody_system.length

        particles_in_field = Particles(keys=(2, 4, 5, 6, 7,))
        particles_in_field.mass = [1, 2, 3, 4, 5] | nbody_system.mass
        particles_in_field[0].position = [2, 0, 0] | nbody_system.length
        particles_in_field[1].position = [1.5, 0, 0] | nbody_system.length
        particles_in_field[2].position = [0.5, 1, 0] | nbody_system.length
        particles_in_field[3].position = [0.5, -0.5, 0] | nbody_system.length
        particles_in_field[4].position = [0, 0, 2] | nbody_system.length
        particles_in_field.velocity = [0, 0, 0] | nbody_system.speed
        particles_in_field.radius = 0 | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.parameters.hard_binary_factor = 1
        x.small_scale_factor = 1

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.particles_in_field.add_particles(particles_in_field)

        simple_binary = new_binary(
            1 | nbody_system.mass,
            1 | nbody_system.mass,
            2 | nbody_system.length
        )

        def evolve_singles_in_encounter_until_end_state():
            particles = x.singles_and_multiples_after_evolve
            particles.add_particles(x.all_singles_in_encounter)
            particles.child1 = None
            particles.child2 = None

            root_particle = particles.add_particle(Particle(
                key=10,
                mass=2.0 | nbody_system.mass,
                position=particles[0:2].center_of_mass(),
                velocity=particles[0:2].center_of_mass_velocity(),
            ))
            root_particle.child1 = particles[0]
            root_particle.child2 = particles[1]
            particles[0].position = simple_binary[0].position + root_particle.position
            particles[1].position = simple_binary[1].position + root_particle.position

            particles[0].velocity = simple_binary[0].velocity + root_particle.velocity
            particles[1].velocity = simple_binary[1].velocity + root_particle.velocity

        x.evolve_singles_in_encounter_until_end_state = evolve_singles_in_encounter_until_end_state
        x.determine_structure_of_the_evolved_state = lambda: 1

        x.execute()

        # no multiples as the binary is larger than the
        # hard binary scale
        self.assertEqual(len(x.new_multiples), 0)
        self.assertEqual(len(x.new_binaries), 0)

    def test6(self):

        particles_in_encounter = Particles(keys=(1, 2, 3))
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter[2].position = [0, 0.5, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.0, 0] | nbody_system.speed
        particles_in_encounter.radius = 0 | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )
        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.small_scale_factor = 1.0

        simple_binary_1 = new_binary(
            1 | nbody_system.mass,
            1 | nbody_system.mass,
            0.4 | nbody_system.length
        )
        simple_binary_top = new_binary(
            2 | nbody_system.mass,
            1 | nbody_system.mass,
            2 | nbody_system.length
        )

        def evolve_singles_in_encounter_until_end_state():
            particles = x.singles_and_multiples_after_evolve
            particles.add_particles(x.all_singles_in_encounter)
            particles.child1 = None
            particles.child2 = None

            inner_binary_particle = particles.add_particle(Particle(
                key=10,
                mass=2.0 | nbody_system.mass,
                position=particles[0:2].center_of_mass(),
                velocity=particles[0:2].center_of_mass_velocity(),
            ))
            inner_binary_particle.child1 = particles[0]
            inner_binary_particle.child2 = particles[1]
            particles[0].position = simple_binary_1[0].position + inner_binary_particle.position
            particles[1].position = simple_binary_1[1].position + inner_binary_particle.position

            particles[0].velocity = simple_binary_1[0].velocity + inner_binary_particle.velocity
            particles[1].velocity = simple_binary_1[1].velocity + inner_binary_particle.velocity

            root_particle = particles.add_particle(Particle(
                key=11,
                mass=3.0 | nbody_system.mass,
                position=particles.center_of_mass(),
                velocity=particles.center_of_mass_velocity(),
            ))

            root_particle.child1 = inner_binary_particle
            root_particle.child2 = particles[2]
            inner_binary_particle.position = simple_binary_top[0].position + root_particle.position
            particles[2].position = simple_binary_top[1].position + root_particle.position

            inner_binary_particle.velocity = simple_binary_top[0].velocity + root_particle.velocity
            particles[2].velocity = simple_binary_top[1].velocity + root_particle.velocity

        x.evolve_singles_in_encounter_until_end_state = evolve_singles_in_encounter_until_end_state
        x.determine_structure_of_the_evolved_state = lambda: 1

        x.execute()

        # no multiples as the binary is larger than the
        # hard binary scale
        self.assertEqual(len(x.new_multiples), 1)
        self.assertEqual(len(x.new_binaries), 1)
        multiple = x.new_multiples[0]
        self.assertEqual(len(multiple.components), 2)
        self.assertAlmostRelativeEqual(multiple.components[0].key, particles_in_encounter[0].key)
        self.assertAlmostRelativeEqual(multiple.components[1].key, particles_in_encounter[1].key)

    def test7(self):

        particles_in_encounter = Particles(keys=(1, 2, 3))
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter[2].position = [0, 0.5, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.0, 0] | nbody_system.speed
        particles_in_encounter.radius = 0 | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )
        x.particles_in_encounter.add_particles(particles_in_encounter)

        multiples = Particles()
        particle_in_multiples = multiples.add_particle(particles_in_encounter[2])
        particle_in_multiples.components = Particles(keys=(4, 5))
        particle_in_multiples.components.mass = 1 | nbody_system.mass
        particle_in_multiples.components[0].position = [0, 0, 0.2] | nbody_system.length
        particle_in_multiples.components[1].position = [0, 0, -0.2] | nbody_system.length
        particle_in_multiples.components.velocity = [0, 0.0, 0] | nbody_system.speed
        particle_in_multiples.components.child1 = None
        particle_in_multiples.components.child2 = None
        x.existing_multiples.add_particles(multiples)

        def evolve_singles_in_encounter_until_end_state():

            particles = x.singles_and_multiples_after_evolve
            particles.add_particles(x.all_singles_in_encounter)
            print(particles)
            particles.child1 = None
            particles.child2 = None

        x.evolve_singles_in_encounter_until_end_state = evolve_singles_in_encounter_until_end_state
        x.determine_structure_of_the_evolved_state = lambda: 1

        x.execute()

        # no multiples as the binary is larger than the
        # hard binary scale
        self.assertEqual(len(x.new_multiples), 0)
        self.assertEqual(len(x.new_binaries), 0)
        self.assertEqual(len(x.dissolved_multiples), 1)
        self.assertEqual(len(x.released_singles), 2)
        self.assertTrue(particle_in_multiples.components[0] in x.released_singles)
        self.assertTrue(particle_in_multiples.components[1] in x.released_singles)

    def xtest8(self):
        particles_in_encounter = Particles(3)
        particles_in_encounter.mass = 1. | nbody_system.mass
        particles_in_encounter[0].position = [0., 0., 0.] | nbody_system.length
        particles_in_encounter[1].position = [1., 0., 0.] | nbody_system.length
        particles_in_encounter[2].position = [2., 0., 0.] | nbody_system.length
        particles_in_encounter.velocity = [0., 0.5, 0.] | nbody_system.speed
        particles_in_encounter.radius = [0.5, 1., 0.2] | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)

        x.execute()

        self.assertAlmostRelativeEqual(x.large_scale_of_particles_in_the_encounter, 2.0 | nbody_system.length)
        self.assertAlmostRelativeEqual(x.small_scale_of_particles_in_the_encounter, 1.5 | nbody_system.length)

    def test8b(self):
        particles_in_encounter = Particles(2)
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.5, 0] | nbody_system.speed
        particles_in_encounter.radius = [0.5, 0.5] | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)

        x.execute()

        self.assertAlmostRelativeEqual(x.initial_potential_in_field,  0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_energy, -0.75 | nbody_system.energy)

    def test9(self):
        particles_in_encounter = Particles(2)
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.5, 0] | nbody_system.speed
        particles_in_encounter.radius = [0.5, 0.5] | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)

        x.execute()

        self.assertAlmostRelativeEqual(x.initial_potential_in_field,  0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_energy, -0.75 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_singles_energy, -0.75 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.delta_phi_1,  0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_multiple_energy,  0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.final_energy, -0.75 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.delta_phi_2, 0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.final_multiple_energy,  0 | nbody_system.energy)

    def test10(self):
        particles_in_encounter = Particles(2)
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.5, 0] | nbody_system.speed
        particles_in_encounter.radius = [0.5, 0.5] | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.particles_in_field.add_particle(Particle(
            mass=1 | nbody_system.mass,
            position=[5, 0, 0] | nbody_system.length,
            velocity=[0, 0, 0] | nbody_system.speed,
            radius=0 | nbody_system.length
        ))

        x.execute()

        self.assertAlmostRelativeEqual(x.initial_potential_in_field,   -0.45 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_energy, -0.75 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_singles_energy, -0.75 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.delta_phi_1,  0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_multiple_energy,  0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.final_energy, -0.75 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.delta_phi_2, 0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.final_multiple_energy,  0 | nbody_system.energy)

    def test11(self):
        particles_in_encounter = Particles(2)
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.5, 0] | nbody_system.speed
        particles_in_encounter.radius = [0.5, 0.5] | nbody_system.length

        particles_in_multiples = Particles()
        particles_in_multiples.add_particle(particles_in_encounter[0])
        multiple = particles_in_multiples[0]
        multiple.components = Particles(keys=(3, 4))
        multiple.components.mass = 0.5 | nbody_system.mass
        multiple.components[0].position = [0, 0.1, 0] | nbody_system.length
        multiple.components[1].position = [0, -0.1, 0] | nbody_system.length
        multiple.components[0].velocity = [0, 0, 0.2] | nbody_system.speed
        multiple.components[1].velocity = [0, 0.1, -0.2] | nbody_system.speed
        multiple.components.child1 = None
        multiple.components.child2 = None

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.particles_in_field.add_particle(Particle(
            mass=1 | nbody_system.mass,
            position=[5, 0, 0] | nbody_system.length,
            velocity=[0, 0, 0] | nbody_system.speed,
            radius=0 | nbody_system.length
        ))
        x.existing_multiples.add_particles(particles_in_multiples)

        x.execute()

        self.assertAlmostRelativeEqual(x.initial_potential_in_field,   -0.45 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_energy, -0.75 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_singles_energy, -1.912495 | nbody_system.energy, 4)
        self.assertAlmostRelativeEqual(x.delta_phi_1,  0.06500400 | nbody_system.energy, 4)
        self.assertAlmostRelativeEqual(x.initial_multiple_energy,  -1.2275 | nbody_system.energy, 4)
        self.assertAlmostRelativeEqual(x.final_energy, -1.9124959996 | nbody_system.energy, 4)
        self.assertAlmostRelativeEqual(x.delta_phi_2, 0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.final_multiple_energy,  0 | nbody_system.energy)

    def test12(self):
        particles_in_encounter = Particles(2)
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.5, 0] | nbody_system.speed
        particles_in_encounter.radius = [0.5, 0.5] | nbody_system.length

        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.particles_in_field.add_particle(Particle(
            mass=0.1 | nbody_system.mass,
            position=[1.5, 0, 0] | nbody_system.length,
            velocity=[0, 0, 0] | nbody_system.speed,
            radius=1 | nbody_system.length
        ))

        x.execute()

        self.assertAlmostRelativeEqual(x.initial_potential_in_field,   0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_energy, -1.016666 | nbody_system.energy, 4)
        self.assertAlmostRelativeEqual(x.initial_singles_energy, -1.016666 | nbody_system.energy, 4)
        self.assertAlmostRelativeEqual(x.delta_phi_1,  0 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.initial_multiple_energy,  0 | nbody_system.energy)

    def test13(self):
        x = TryHandleEncounter(
            G=nbody_system.G,
            kepler_code=self.new_kepler()
        )

        x.parameters.hard_binary_factor = 1
        self.assertEqual(x.parameters.hard_binary_factor, 1)
        self.assertEqual(x.parameters.scatter_factor, 10)
        x.parameters.scatter_factor = 2
        self.assertEqual(x.parameters.scatter_factor, 2)
        self.assertEqual(x.get_scatter_factor(), 2)


class TestKeplerOrbits(amusetest.TestWithMPI):

    def new_kepler(self):
        x = Kepler()
        x.initialize_code()
        return x

    def test1(self):
        x = encounters.KeplerOrbits(self.new_kepler())
        binary = new_binary(
            1 | nbody_system.mass,
            0.5 | nbody_system.mass,
            1.2 | nbody_system.length,
        )
        semimajor_axis, eccentricity = x.get_semimajor_axis_and_eccentricity_for_binary_components(
            binary[0],
            binary[1]
        )
        self.assertAlmostRelativeEquals(semimajor_axis,  1.2 | nbody_system.length)
        self.assertAlmostRelativeEquals(eccentricity,  0)

    def test2(self):
        x = encounters.KeplerOrbits(self.new_kepler())
        binary = new_binary(
            1 | nbody_system.mass,
            0.5 | nbody_system.mass,
            2.0 | nbody_system.length,
            0.5,
            is_at_periapsis=False

        )
        binary.position += [0.1, 0.2, 0.3] | nbody_system.length
        binary.velocity += [0.4, 0.5, 0.6] | nbody_system.speed

        dpos, dvel = x.compress_binary(
            binary,
            1.2 | nbody_system.length
        )

        center_of_mass_before = binary.center_of_mass()
        center_of_mass_velocity_before = binary.center_of_mass_velocity()

        binary.position += dpos
        binary.velocity += dvel

        center_of_mass_after = binary.center_of_mass()
        center_of_mass_velocity_after = binary.center_of_mass_velocity()

        self.assertAlmostRelativeEquals(center_of_mass_before, center_of_mass_after)
        self.assertAlmostRelativeEquals(center_of_mass_velocity_before, center_of_mass_velocity_after)
        separation = (binary[0].position - binary[1].position).length()
        self.assertAlmostRelativeEquals(separation, 1.2 | nbody_system.length)

    def test3(self):
        x = encounters.KeplerOrbits(self.new_kepler())
        binary = new_binary(
            1 | nbody_system.mass,
            0.5 | nbody_system.mass,
            2.0 | nbody_system.length,
            0.5,
            is_at_periapsis=True

        )
        binary.position += [0.1, 0.2, 0.3] | nbody_system.length
        binary.velocity += [0.4, 0.5, 0.6] | nbody_system.speed

        dpos, dvel = x.expand_binary(
            binary,
            2.5 | nbody_system.length
        )

        center_of_mass_before = binary.center_of_mass()
        center_of_mass_velocity_before = binary.center_of_mass_velocity()

        binary.position += dpos
        binary.velocity += dvel

        center_of_mass_after = binary.center_of_mass()
        center_of_mass_velocity_after = binary.center_of_mass_velocity()

        self.assertAlmostRelativeEquals(center_of_mass_before, center_of_mass_after, 8)
        self.assertAlmostRelativeEquals(center_of_mass_velocity_before, center_of_mass_velocity_after, 8)
        separation = (binary[0].position - binary[1].position).length()
        self.assertAlmostRelativeEquals(separation, 2.5 | nbody_system.length)

    def test4(self):
        converter = nbody_system.nbody_to_si(1 | units.MSun, 1 | units.AU)
        kepler = Kepler(converter)
        kepler.initialize_code()
        x = encounters.KeplerOrbits(kepler)
        binary = new_binary(
            1 | units.MSun,
            0.5 | units.MSun,
            1.2 | units.AU,
            G=constants.G
        )
        semimajor_axis, eccentricity = x.get_semimajor_axis_and_eccentricity_for_binary_components(
            binary[0],
            binary[1]
        )
        self.assertAlmostRelativeEquals(semimajor_axis,  1.2 | units.AU, 8)
        self.assertAlmostRelativeEquals(eccentricity,  0)


class TestScaleSystem(amusetest.TestWithMPI):

    def new_kepler(self, converter=None):
        x = Kepler(converter)
        x.initialize_code()
        return x

    def test1(self):
        kepler = encounters.KeplerOrbits(self.new_kepler())
        binary = new_binary(
            1 | nbody_system.mass,
            0.5 | nbody_system.mass,
            1.2 | nbody_system.length,
            0.5,
            is_at_periapsis=False
        )
        binary.radius = 0 | nbody_system.length

        x = encounters.ScaleSystem(kepler)
        self.assertTrue((binary[0].position - binary[1].position).length() > 1.5 | nbody_system.length)
        x.scale_particles_to_sphere(binary, 0.75 | nbody_system.length)
        self.assertTrue((binary[0].position - binary[1].position).length() <= 1.6 | nbody_system.length)
        print((binary[0].position - binary[1].position).length())
        self.assertTrue((binary[0].position - binary[1].position).length() >= (1.5 - 1e-6) | nbody_system.length)

    def test2(self):
        kepler = encounters.KeplerOrbits(self.new_kepler())
        binary = new_binary(
            1 | nbody_system.mass,
            0.5 | nbody_system.mass,
            1.2 | nbody_system.length,
            0.5,
            is_at_periapsis=True
        )
        binary.radius = 0 | nbody_system.length

        x = encounters.ScaleSystem(kepler)
        self.assertTrue((binary[0].position - binary[1].position).length() < 1.5 | nbody_system.length)
        x.scale_particles_to_sphere(binary, 0.75 | nbody_system.length)
        self.assertTrue((binary[0].position - binary[1].position).length() <= 1.6 | nbody_system.length)
        print((binary[0].position - binary[1].position).length())
        self.assertTrue((binary[0].position - binary[1].position).length() >= (1.5 - 1e-6) | nbody_system.length)

    def test3(self):
        kepler = encounters.KeplerOrbits(self.new_kepler())

        particles = Particles(keys=(1, 2))
        particles.mass = 1 | nbody_system.mass
        particles.child2 = None
        particles.child1 = None
        particles[0].position = [1, 0, 0] | nbody_system.length
        particles[0].velocity = [1, 0.0, 0] | nbody_system.speed

        particles[1].position = [-1, 0, 0] | nbody_system.length
        particles[1].velocity = [-1, 0.0, 0] | nbody_system.speed

        particles.radius = 0 | nbody_system.length

        x = encounters.ScaleSystem(kepler)

        print((particles[0].position - particles[1].position).length())
        self.assertTrue((particles[0].position - particles[1].position).length() > 1.5 | nbody_system.length)

        x.scale_particles_to_sphere(particles, 0.75 | nbody_system.length)

        self.assertTrue((particles[0].position - particles[1].position).length() <= 1.6 | nbody_system.length)
        print(particles)
        self.assertTrue((particles[0].position - particles[1].position).length() >= (1.5 - 1e-6) | nbody_system.length)
        self.assertAlmostRelativeEquals(particles[0].position,  [0.75, 0, 0] | nbody_system.length)
        self.assertAlmostRelativeEquals(particles[1].position,  [-0.75, 0, 0] | nbody_system.length)

    def test4(self):
        kepler = encounters.KeplerOrbits(self.new_kepler())

        particles = Particles(keys=(1, 2, 3, 4, 5, 6))
        particles.child2 = None
        particles.child1 = None
        particles.mass = 1 | nbody_system.mass
        for i in range(3):
            position = [0, 0, 0] | nbody_system.length
            position[i] = 1 | nbody_system.length
            particles[i].position = position
            position[i] = -1 | nbody_system.length
            particles[i+3].position = position

            velocity = [0, 0, 0] | nbody_system.speed
            velocity[i] = 1 | nbody_system.speed
            particles[i].velocity = velocity
            velocity[i] = -1 | nbody_system.speed
            particles[i+3].velocity = velocity

        particles.radius = 0 | nbody_system.length

        x = encounters.ScaleSystem(kepler)
        potential_energy0 = particles.potential_energy(G=nbody_system.G)
        kinetic_energy0 = particles.kinetic_energy()

        x.scale_particles_to_sphere(particles, 0.5 | nbody_system.length)

        self.assertAlmostRelativeEquals(particles[0].position,  [1/numpy.sqrt(2), 0, 0] | nbody_system.length)
        self.assertAlmostRelativeEquals(particles[3].position,  [-1/numpy.sqrt(2), 0, 0] | nbody_system.length)
        self.assertAlmostRelativeEquals(particles[0].velocity,  [1.542, 0, 0] | nbody_system.speed, 3)
        self.assertAlmostRelativeEquals(particles[3].velocity,  [-1.542, 0, 0] | nbody_system.speed, 3)

        potential_energy1 = particles.potential_energy(G=nbody_system.G)
        kinetic_energy1 = particles.kinetic_energy()

        self.assertAlmostRelativeEquals(potential_energy0 + kinetic_energy0, potential_energy1 + kinetic_energy1)

    def test5(self):
        converter = nbody_system.nbody_to_si(1 | units.MSun, 1 | units.AU)
        kepler = encounters.KeplerOrbits(self.new_kepler(converter))

        particles = Particles(keys=(1, 2, 3))
        particles.child2 = None
        particles.child1 = None
        particles.position = [
            [-1.28230200e-05,  -3.69457095e-05,  -2.02383488e-05],
            [-2.91749746e-05,  -1.21387289e-05,   1.56377986e-07],
            [2.92123436e-05,   1.22463965e-05,  -9.73992061e-08]
        ] | units.parsec
        particles.velocity = [
            [-8685.98414414, -413260.30543051, -135268.19175611],
            [19639.55455979,  -30251.55372943,     972.69648982],
            [-19614.24178601,   31455.88066005,    -578.49669843]
        ] | (units.m / units.s)
        particles.mass = [0.14571045,  50., 50.] | units.MSun
        particles.radius = [8, 0., 0.0] | units.AU

        x = encounters.ScaleSystem(kepler, G=constants.G)

        potential_energy0 = particles.potential_energy()
        kinetic_energy0 = particles.kinetic_energy()

        x.scale_particles_to_sphere(particles, 4 | units.AU)

        potential_energy1 = particles.potential_energy()
        kinetic_energy1 = particles.kinetic_energy()

        self.assertAlmostRelativeEquals(potential_energy0 + kinetic_energy0, potential_energy1 + kinetic_energy1)


class TestHandleEncounter(amusetest.TestWithMPI):

    def new_kepler(self):
        x = Kepler()
        x.initialize_code()
        return x

    def test1(self):
        particles_in_encounter = Particles(keys=(1, 2, 3))
        particles_in_encounter.mass = 1. | nbody_system.mass
        particles_in_encounter[0].position = [1., 0., 0.] | nbody_system.length
        particles_in_encounter[1].position = [0., 0., 0.] | nbody_system.length
        particles_in_encounter[2].position = [0., 0.5, 0.] | nbody_system.length
        particles_in_encounter.velocity = [0., 0.0, 0.] | nbody_system.speed
        particles_in_encounter[2].velocity = [0., 0.0, 3.] | nbody_system.speed
        particles_in_encounter.radius = 0. | nbody_system.length

        particles_in_field = Particles()

        x = encounters.HandleEncounter(
            kepler_code=self.new_kepler(),
            resolve_collision_code=SmallN(),
            interaction_over_code=None,
            G=nbody_system.G
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)

        x.execute()
        self.assertEqual(len(x.new_multiples), 1)
        self.assertEqual(len(x.new_binaries), 1)
        multiple = x.new_multiples[0]
        self.assertEqual(len(multiple.components), 2)
        self.assertAlmostRelativeEqual(multiple.components[0].key, particles_in_encounter[0].key)
        self.assertAlmostRelativeEqual(multiple.components[1].key, particles_in_encounter[1].key)
        self.assertEqual(len(x.captured_singles), 2)
        self.assertEqual(x.captured_singles.key, [1, 2])

    def test2(self):
        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 2. | nbody_system.mass
        particles_in_encounter[0].position = [1., 0., 0.] | nbody_system.length
        particles_in_encounter[1].position = [0., 0., 0.] | nbody_system.length
        particles_in_encounter.velocity = [0., 0.0, 0.] | nbody_system.speed
        particles_in_encounter.radius = 0. | nbody_system.length

        binary1 = new_binary(
            1. | nbody_system.mass,
            1. | nbody_system.mass,
            0.01 | nbody_system.length,
            keyoffset=30
        )
        binary2 = new_binary(
            1. | nbody_system.mass,
            1. | nbody_system.mass,
            0.01 | nbody_system.length,
            keyoffset=40
        )
        binaries = Particles(keys=(20, 21))
        binaries[0].child1 = binary1[0]
        binaries[0].child2 = binary1[1]
        binaries[1].child1 = binary2[0]
        binaries[1].child2 = binary2[1]
        binary1.child1 = None
        binary1.child2 = None
        binary2.child1 = None
        binary2.child2 = None

        multiples = Particles()
        multiple = particles_in_encounter[0]
        multiple.components = binary1
        multiples.add_particle(multiple)
        multiple = particles_in_encounter[1]
        multiple.components = binary2
        multiples.add_particle(multiple)

        x = encounters.HandleEncounter(
            kepler_code=self.new_kepler(),
            resolve_collision_code=SmallN(),
            interaction_over_code=None,
            G=nbody_system.G
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.existing_binaries.add_particles(binaries)
        x.existing_multiples.add_particles(multiples)

        x.execute()
        self.assertEqual(len(x.new_multiples), 2)
        self.assertEqual(len(x.new_binaries), 2)
        self.assertEqual(len(x.captured_singles), 0)
        self.assertEqual(len(x.released_singles), 0)

        multiple = x.new_multiples[0]
        self.assertEqual(len(multiple.components), 2)

        self.assertAlmostRelativeEqual(multiple.components[0].key, binaries[0].child1.key)
        self.assertAlmostRelativeEqual(multiple.components[1].key, binaries[1].child1.key)
        multiple = x.new_multiples[1]
        self.assertEqual(len(multiple.components), 2)
        self.assertAlmostRelativeEqual(multiple.components[0].key, binaries[0].child2.key)
        self.assertAlmostRelativeEqual(multiple.components[1].key, binaries[1].child2.key)

    def test3(self):
        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 2. | nbody_system.mass
        particles_in_encounter[0].position = [1., 0., 0.] | nbody_system.length
        particles_in_encounter[1].position = [0., 0., 0.] | nbody_system.length
        particles_in_encounter.velocity = [0., 0.0, 0.] | nbody_system.speed
        particles_in_encounter.radius = 0. | nbody_system.length

        binary1 = new_binary(
            1. | nbody_system.mass,
            1. | nbody_system.mass,
            0.01 | nbody_system.length,
            keyoffset=30
        )
        binaries = Particles(keys=(20,))
        binaries[0].child1 = binary1[0]
        binaries[0].child2 = binary1[1]
        binary1.child1 = None
        binary1.child2 = None

        multiples = Particles()
        multiple = particles_in_encounter[0]
        multiple.components = binary1
        multiples.add_particle(multiple)

        x = encounters.HandleEncounter(
            kepler_code=self.new_kepler(),
            resolve_collision_code=SmallN(),
            interaction_over_code=None,
            G=nbody_system.G
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.existing_binaries.add_particles(binaries)
        x.existing_multiples.add_particles(multiples)

        x.execute()

        self.assertEqual(len(x.new_multiples), 1)
        self.assertEqual(len(x.dissolved_multiples), 1)
        self.assertEqual(len(x.new_binaries), 1)
        self.assertEqual(len(x.captured_singles), 1)
        self.assertEqual(len(x.released_singles), 1)

        multiple = x.new_multiples[0]
        print(multiple.child1)
        self.assertEqual(len(multiple.components), 2)
        print(multiple.components)

    def test4(self):
        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 1.0 / 20.0 | nbody_system.mass

        particles_in_encounter[0].position = [0.303976184547589401, 0.273137803329168094, 0.0] | nbody_system.length
        particles_in_encounter[1].position = [0.290167020486133631, 0.273139253307546515, 0.0] | nbody_system.length
        particles_in_encounter[0].velocity = [-2.544712989335638387, -1.224759650260411004, 0.0] | nbody_system.speed
        particles_in_encounter[1].velocity = [0.898326624897966997, 0.870611747842778838, 0.0] | nbody_system.speed
        particles_in_encounter.radius = 0.007 | nbody_system.length

        x = encounters.HandleEncounter(
            kepler_code=self.new_kepler(),
            resolve_collision_code=SmallN(),
            interaction_over_code=None,
            G=nbody_system.G
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)

        x.execute()

        self.assertEqual(len(x.new_multiples), 0)
        self.assertEqual(len(x.dissolved_multiples), 0)
        self.assertEqual(len(x.new_binaries), 0)
        self.assertEqual(len(x.captured_singles), 0)
        self.assertEqual(len(x.released_singles), 0)
        r_before = (particles_in_encounter[0].position - particles_in_encounter[1].position).length()
        r_after = (x.particles_after_encounter[0].position - x.particles_after_encounter[1].position).length()
        print(r_before, r_after)
        self.assertFalse(r_after > (10 * r_before))

    def test5(self):
        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 1.0 / 20.0 | nbody_system.mass

        particles_in_encounter[0].position = [0.0779377282404, -0.559210143918, 0.0] | nbody_system.length
        particles_in_encounter[1].position = [0.07802860386, -0.561207706614, 0.0] | nbody_system.length
        particles_in_encounter[0].velocity = [-1.79086847491, -4.88551917358, 0.0] | nbody_system.speed
        particles_in_encounter[1].velocity = [2.13208424698, 4.31403500143, 0.0] | nbody_system.speed
        particles_in_encounter.radius = 0.001 | nbody_system.length

        x = encounters.HandleEncounter(
            kepler_code=self.new_kepler(),
            resolve_collision_code=SmallN(),
            interaction_over_code=None,
            G=nbody_system.G
        )
        x.particles_in_encounter.add_particles(particles_in_encounter)

        x.execute()

        self.assertEqual(len(x.new_multiples), 0)
        self.assertEqual(len(x.dissolved_multiples), 0)
        self.assertEqual(len(x.new_binaries), 0)
        self.assertEqual(len(x.captured_singles), 0)
        self.assertEqual(len(x.released_singles), 0)
        r_before = (particles_in_encounter[0].position - particles_in_encounter[1].position).length()
        r_after = (x.particles_after_encounter[0].position - x.particles_after_encounter[1].position).length()
        print(r_before, r_after)
        self.assertFalse(r_after > (10 * r_before))

    def test6(self):
        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 2 | nbody_system.mass
        particles_in_encounter[0].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.0, 0] | nbody_system.speed
        particles_in_encounter.radius = 0 | nbody_system.length

        particles_in_field = Particles(10)
        particles_in_field.mass = 1 | nbody_system.mass
        for i in range(len(particles_in_field)):
            if i == 5:
                j = len(particles_in_field)
            else:
                j = i
            particles_in_field[i].position = [-10 + (2*j), -10 + (2*j), 0] | nbody_system.length
        particles_in_field.velocity = [0, 0.0, 0] | nbody_system.speed
        particles_in_field.radius = 0 | nbody_system.length

        x = encounters.HandleEncounter(
            kepler_code=self.new_kepler(),
            resolve_collision_code=SmallN(),
            interaction_over_code=None,
            G=nbody_system.G
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.particles_in_field.add_particles(particles_in_field)

        x.execute()

        self.assertEqual(len(x.new_multiples), 1)
        self.assertEqual(len(x.dissolved_multiples), 0)
        self.assertEqual(len(x.new_binaries), 1)
        self.assertEqual(len(x.captured_singles), 2)
        self.assertEqual(len(x.released_singles), 0)

    def test7(self):
        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 2. | nbody_system.mass
        particles_in_encounter[0].position = [1., 0., 0.] | nbody_system.length
        particles_in_encounter[1].position = [0., 0., 0.] | nbody_system.length
        particles_in_encounter.velocity = [0., 0.0, 0.] | nbody_system.speed
        particles_in_encounter.radius = 0. | nbody_system.length

        binary1 = new_binary(
            1. | nbody_system.mass,
            1. | nbody_system.mass,
            0.01 | nbody_system.length,
            keyoffset=30
        )
        binary_energy = binary1.kinetic_energy() + binary1.potential_energy(G=nbody_system.G)
        self.assertAlmostRelativeEqual(binary_energy, -50 | nbody_system.energy)

        binaries = Particles(keys=(20,))
        binaries[0].child1 = binary1[0]
        binaries[0].child2 = binary1[1]
        binary1.child1 = None
        binary1.child2 = None

        multiples = Particles()
        multiple = particles_in_encounter[0]
        multiple.components = binary1
        multiples.add_particle(multiple)

        x = encounters.HandleEncounter(
            kepler_code=self.new_kepler(),
            resolve_collision_code=SmallN(),
            interaction_over_code=None,
            G=nbody_system.G
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.existing_binaries.add_particles(binaries)
        x.existing_multiples.add_particles(multiples)

        x.execute()
        self.assertEqual(len(x.new_multiples), 1)
        self.assertEqual(len(x.new_binaries), 1)
        self.assertEqual(len(x.captured_singles), 1)
        self.assertEqual(len(x.released_singles), 1)

        multiple = x.new_multiples[0]
        self.assertEqual(len(multiple.components), 2)

        self.assertAlmostRelativeEqual(multiple.components[0].key, binaries[0].child1.key)
        self.assertAlmostRelativeEqual(multiple.components[1].key, particles_in_encounter[1].key)
        self.assertAlmostRelativeEqual(x.initial_multiple_energy, -50 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.final_multiple_energy, -85.651753404 | nbody_system.energy, 4)

    def test8(self):
        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter[0].mass = 2. | nbody_system.mass
        particles_in_encounter[1].mass = 0.1 | nbody_system.mass
        particles_in_encounter[0].position = [1., 0., 0.] | nbody_system.length
        particles_in_encounter[1].position = [0., 0., 0.] | nbody_system.length
        particles_in_encounter.velocity = [0., 0.0, 0.] | nbody_system.speed
        particles_in_encounter.radius = 0. | nbody_system.length

        binary1 = new_binary(
            1. | nbody_system.mass,
            1. | nbody_system.mass,
            0.01 | nbody_system.length,
            keyoffset=30
        )

        binary_energy = binary1.kinetic_energy() + binary1.potential_energy(G=nbody_system.G)
        self.assertAlmostRelativeEqual(binary_energy, -50 | nbody_system.energy)

        binaries = Particles(keys=(20,))
        binaries[0].child1 = binary1[0]
        binaries[0].child2 = binary1[1]
        binary1.child1 = None
        binary1.child2 = None

        multiples = Particles()
        multiple = particles_in_encounter[0]
        multiple.components = binary1
        multiples.add_particle(multiple)

        x = encounters.HandleEncounter(
            kepler_code=self.new_kepler(),
            resolve_collision_code=SmallN(),
            interaction_over_code=None,
            G=nbody_system.G
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)
        x.existing_binaries.add_particles(binaries)
        x.existing_multiples.add_particles(multiples)

        x.execute()
        self.assertEqual(len(x.new_multiples), 0)
        self.assertEqual(len(x.updated_multiples), 1)
        self.assertEqual(len(x.new_binaries), 0)
        self.assertEqual(len(x.captured_singles), 0)
        self.assertEqual(len(x.released_singles), 0)

        multiple = x.updated_multiples[0]
        self.assertEqual(len(multiple.components), 2)
        self.assertEqual(multiple.key, multiples[0].key)

        self.assertEqual(multiple.components[0].key, binaries[0].child1.key)
        self.assertEqual(multiple.components[1].key, binaries[0].child2.key)
        self.assertAlmostRelativeEqual(x.initial_multiple_energy, -50 | nbody_system.energy)
        self.assertAlmostRelativeEqual(x.final_multiple_energy,  -56.2906032319 | nbody_system.energy, 4)


class TestStickyHandleEncounter(amusetest.TestWithMPI):

    def test1(self):
        particles_in_encounter = Particles(keys=(1, 2))
        particles_in_encounter.mass = 1 | nbody_system.mass
        particles_in_encounter[0].position = [1, 0, 0] | nbody_system.length
        particles_in_encounter[1].position = [0, 0, 0] | nbody_system.length
        particles_in_encounter.velocity = [0, 0.0, 0] | nbody_system.speed
        particles_in_encounter[0].velocity = [-1, 0.0, 0] | nbody_system.speed
        particles_in_encounter.radius = 0 | nbody_system.length

        particles_in_field = Particles()

        x = encounters.StickyHandleEncounter(
            G=nbody_system.G
        )

        x.particles_in_encounter.add_particles(particles_in_encounter)

        x.execute()
        self.assertEqual(len(x.new_multiples), 1)
        self.assertEqual(len(x.new_binaries), 1)
        multiple = x.new_multiples[0]
        self.assertEqual(len(multiple.components), 2)
        self.assertAlmostRelativeEqual(multiple.components[0].key, particles_in_encounter[0].key)
        self.assertAlmostRelativeEqual(multiple.components[1].key, particles_in_encounter[1].key)
        self.assertEqual(len(x.captured_singles), 2)
        self.assertEqual(x.captured_singles.key, [1, 2])
        self.assertAlmostRelativeEqual(multiple.velocity, [-0.5, 0.0, 0] | nbody_system.speed)
        self.assertAlmostRelativeEqual(multiple.mass, 2 | nbody_system.mass)
        self.assertAlmostRelativeEqual(multiple.radius, 0.5 | nbody_system.length)
        self.assertAlmostRelativeEqual(multiple.position, [0.5, 0, 0] | nbody_system.length)
