import numpy as np
from amuse.lab import *
import matplotlib.pyplot as plt
from amuse.plot import plot
import sys

N_part = 1e4
scale_radius = 4.43 | units.kpc
converter_gal = nbody_system.nbody_to_si(1e10 | units.MSun, scale_radius) 
np.random.seed(123)
galaxy = new_halogen_model(N_part, converter_gal, alpha=1, beta=3, gamma=1, 
                           cutoff_radius=10*scale_radius, scale_radius=scale_radius)

epsilon=443 | units.pc
dt = 1.000000000000000000 | units.Myr
# gadget options
gadget_options = {'number_of_workers' : 20, 'epsilon_squared' : (epsilon)**2, 'begin_time': 0.0 | units.Myr,
                       'max_size_timestep':dt,'time_max':dt*2.**14., 'time_limit_cpu': 0.1 | units.yr,
                       'timestep_accuracy_parameter':0.01, 'opening_angle':0.5}
converter_gadget = nbody_system.nbody_to_si(dt, 1e10 | units.MSun)
galaxy.velocity*=(-2.*galaxy.kinetic_energy()/galaxy.potential_energy(smoothing_length_squared=gadget_options['epsilon_squared']))**-.5

gravity = Gadget2(converter_gadget, **gadget_options, redirection='file',redirect_file='output_short.txt')
gravity.parameters.time_max = gadget_options['time_max']
gravity.parameters.epsilon_squared=gadget_options['epsilon_squared']
gravity.parameters.begin_time=gadget_options['begin_time']
gravity.parameters.max_size_timestep=gadget_options['max_size_timestep']
gravity.parameters.time_limit_cpu=gadget_options['time_limit_cpu']
gravity.parameters.timestep_accuracy_parameter=gadget_options['timestep_accuracy_parameter']
gravity.parameters.opening_angle=gadget_options['opening_angle']
gravity.commit_parameters()
gravity.particles.add_particles(galaxy)
gravity.commit_particles()

# evolve the model
time = 0. | units.myr
t_end = dt * 50.

Menc = [] | units.MSun
lagrange_radii = [] | units.kpc
t = [] | units.Myr

while time < t_end:
    time +=dt
    print(time)
    gravity.evolve_model(time)
    sys.stdout.flush()
    t.append(time)
    print('actual timestep',gravity.get_time_step().in_(units.Myr))
    selection = (gravity.particles.position-gravity.particles.center_of_mass()).lengths()<4.43 | units.kpc
    Menc.append(gravity.particles.mass[selection].sum())
gravity.stop()

######## Now evolve with bridge ###########

gravity = Gadget2(converter_gadget, **gadget_options, redirection='file',redirect_file='output_copy_short.txt')
gravity.parameters.time_max = gadget_options['time_max']
gravity.parameters.epsilon_squared=gadget_options['epsilon_squared']
gravity.parameters.begin_time=gadget_options['begin_time']
gravity.parameters.max_size_timestep=gadget_options['max_size_timestep']
gravity.parameters.time_limit_cpu=gadget_options['time_limit_cpu']
gravity.parameters.timestep_accuracy_parameter=gadget_options['timestep_accuracy_parameter']
gravity.parameters.opening_angle=gadget_options['opening_angle']
gravity.commit_parameters()
gravity.particles.add_particles(galaxy)
gravity.commit_particles()
time = 0. | units.myr

Menc_bridge = [] | units.MSun
t_bridge = [] | units.Myr
required_attributes=['mass', 'x', 'y', 'z', 'vx', 'vy', 'vz']
att=lambda p, x : x in required_attributes
while time < t_end:
    time +=dt 
    print('evolve_model called to', time)
    print('timestep used', gravity.get_time_step().in_(units.Myr))
    gravity.evolve_model(time)
    print('the model time after evolve_model', gravity.model_time.in_(units.Myr))
    sys.stdout.flush()
    t_bridge.append(time)
    # try the channel copying
    test_copy = gravity.particles.copy(filter_attributes =att)
    channel = test_copy.new_channel_to(gravity.particles)
    channel.copy_attributes(["vx","vy","vz"]) # <------ problem line
    # gravity.particles.velocity = test_copy.velocity # <----- also breaks in the same way
    # gravity.recommit_particles()
    selection = (gravity.particles.position-gravity.particles.center_of_mass()).lengths()<4.43 | units.kpc
    Menc_bridge.append(gravity.particles.mass[selection].sum())
gravity.stop()

plt.figure()
plot(t, Menc,label='evolve model', c='tab:blue')
plot(t_bridge, Menc_bridge,label='copy', c='tab:orange')
plt.legend()
plt.title('Mass enclosed by 4.43 kpc')
plt.savefig('mass_enclosed_short.png')