import numpy as np
from amuse.lab import *
from amuse.couple import bridge
from amuse.datamodel import Particles
import matplotlib.pyplot as plt
from amuse.plot import plot

def circular_velocity_profile(galaxy):
    radii = np.linspace(0.02, 10, 60) | units.kpc
    Vcirc = [] | units.kms
    for radius in radii:
        selection = (galaxy.position - galaxy.center_of_mass()).lengths()<radius
        Menc=galaxy[selection].mass.sum()
        Vcirc.append((constants.G * Menc/radius)**.5)
    return Vcirc, radii

# a hydro class wrapper for gadget2
class Hydro():
    def __init__(self, code, particles, gadget_options, dt= 5 | units.Myr):
        # self.particles.add_particles(particles)
        self.converter = nbody_system.nbody_to_si(dt/4, 1e10 | units.MSun) 
        self.code = code(self.converter, **gadget_options) 
        self.code.parameters.time_max = gadget_options['time_max']
        self.code.parameters.epsilon_squared=gadget_options['epsilon_squared']
        self.code.parameters.begin_time=gadget_options['begin_time']
        self.code.parameters.max_size_timestep=gadget_options['max_size_timestep']
        self.code.parameters.time_limit_cpu=gadget_options['time_limit_cpu']
        self.code.parameters.timestep_accuracy_parameter=gadget_options['timestep_accuracy_parameter']
        self.code.parameters.opening_angle=gadget_options['opening_angle']
        self.code.commit_parameters()
        self.code.particles.add_particles(particles)

    def evolve_model(self, time): 
        self.code.evolve_model(time)

    def copy_to_framework(self): 
        self.channel_to_framework.copy()

    def get_gravity_at_point(self, r, x, y, z):
        # edit this to use fastkick!
        return self.code.get_gravity_at_point(r, x, y, z)
    
    def get_potential_at_point(self, r, x, y, z):
        return self.code.get_potential_at_point(r, x, y, z)
    
    def get_timestep(self):
        return self.code.parameters.timestep
    
    @property
    def model_time(self):
        return self.code.model_time 
    
    @property
    def particles(self):
        return self.code.particles
    
    @property
    def total_energy(self):
        return self.code.kinetic_energy + self.code.potential_energy 
    
    @property
    def stop(self):
        return self.code.stop


N_part = 1e5 
scale_radius = 4.1 | 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=75 | units.kpc, scale_radius=scale_radius)

epsilon=18. | units.pc
dt = 5. | units.Myr
# gadget options
gadget_options = {'number_of_workers' : 15, 'epsilon_squared' : (epsilon)**2, 'begin_time': 0.0 | units.Myr,
                       'max_size_timestep':dt/4.,'time_max':dt*2**13, 'time_limit_cpu': 0.1 | units.yr,
                       'timestep_accuracy_parameter':0.01, 'opening_angle':0.5}
converter_gadget = nbody_system.nbody_to_si(dt/4., 1e10 | units.MSun)

vcirc_init, rcirc_init = circular_velocity_profile(galaxy)

# gravity = Gadget2(converter_gadget, **gadget_options, redirection='none')
# 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 = BHTree(converter_gadget, number_of_workers=15, redirection='none')
# gravity.parameters.timestep = 1/4 | nbody_system.time
# gravity.parameters.epsilon_squared = converter_gadget.to_nbody(epsilon**2)
# gravity.particles.add_particles(galaxy)

gravity = Hydro(code=Gadget2, particles=galaxy,gadget_options=gadget_options, dt=dt)

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

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

while time < t_end:
    time +=dt
    print(time)
    gravity.evolve_model(time)
    t.append(time)
    lr,mf=gravity.particles.LagrangianRadii(mf=[0.01,0.05,0.1,0.2])
    lagrange_radii.append(lr)
    selection = (gravity.particles.position-gravity.particles.center_of_mass()).lengths()<4.1 | units.kpc
    Menc.append(gravity.particles.mass[selection].sum())
vcirc_final, rcirc_final = circular_velocity_profile(gravity.particles.copy())
gravity.stop()

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

# gravity = Gadget2(converter_gadget, **gadget_options, redirection='none')
# 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 = BHTree(converter_gadget, number_of_workers=15, redirection='none')
# gravity.parameters.timestep = 1/4 | nbody_system.time
# gravity.parameters.epsilon_squared = converter_gadget.to_nbody(epsilon**2)
# gravity.particles.add_particles(galaxy)

gravity = Hydro(code=Gadget2, particles=galaxy,gadget_options=gadget_options, dt=dt)

p= Particles(1)
p.position = [0,0,0] | units.pc
p.velocity =[0,0,0] | units.kms
p.mass = 0.01 | units.MSun
test_particle =BHTree(converter_gadget)
test_particle.parameters.timestep = 1/4 | nbody_system.time
test_particle.particles.add_particles(p)

integrator=bridge.Bridge(verbose=True,timestep=dt,use_threading=False)
# integrator.add_code(gravity) # this works, but not add_system, with or without kickers
integrator.add_system(gravity, (test_particle,))
time = 0. | units.myr

Menc_bridge = [] | units.MSun
lagrange_radii_bridge = [] | units.kpc
t_bridge = [] | units.Myr

while time < t_end:
    time +=dt 
    print(time)
    integrator.evolve_model(time)
    t_bridge.append(time)
    lr,mf=gravity.particles.LagrangianRadii(mf=[0.01,0.05,0.1,0.2])
    lagrange_radii_bridge.append(lr)
    selection = (gravity.particles.position-gravity.particles.center_of_mass()).lengths()<4.1 | units.kpc
    Menc_bridge.append(gravity.particles.mass[selection].sum())
vcirc_bridge, rcirc_bridge = circular_velocity_profile(gravity.particles.copy())
gravity.stop()

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

plt.figure()
plot(t, lagrange_radii,label='evolve model', c='tab:blue')
plot(t_bridge, lagrange_radii_bridge,label='bridge', c='tab:orange')
plt.yscale('log')
plt.legend()
plt.title('0.01,0.05,0.1,0.2 Lagrangian radii')
plt.savefig('lagrangian_radii_class_testp.png')

plt.figure()
plot(rcirc_init, vcirc_init, label='initial')
plot(rcirc_final, vcirc_final, label='evolve model')
plot(rcirc_bridge, vcirc_bridge, label='bridge')
plt.legend()
plt.savefig('final_circular_velocity_class_testp.png')