import astropy.coordinates as coord
import astropy.units as u
import numpy as np
import gala.coordinates as gc
import gala.dynamics as gd
import gala.potential as gp
from gala.units import galactic
from gala.dynamics import mockstream as ms

coord.galactocentric_frame_defaults.set('v4.0')

pot = gp.CCompositePotential()
pot['disk'] = gp.MiyamotoNagaiPotential(m=6E10*u.Msun,
                                        a=3.5*u.kpc, b=280*u.pc,
                                        units=galactic)
pot['halo'] = gp.NFWPotential(m=1E12, r_s=20*u.kpc, units=galactic)

c = coord.ICRS(ra=229 * u.deg, dec=-0.124 * u.deg,
               distance=22.9 * u.kpc,
               pm_ra_cosdec=-2.296 * u.mas/u.yr,
               pm_dec=-2.257 * u.mas/u.yr,
               radial_velocity=-58.7 * u.km/u.s)

c_gc = c.transform_to(coord.Galactocentric()).cartesian
pal5_w0 = gd.PhaseSpacePosition(c_gc)

pal5_mass = 2.5e4 * u.Msun
pal5_pot = gp.PlummerPotential(m=pal5_mass, b=4*u.pc, units=galactic)

df = ms.FardalStreamDF(gala_modified=True)
gen_pal5 = ms.MockStreamGenerator(df, pot, progenitor_potential=pal5_pot)
pal5_stream, _ = gen_pal5.run(pal5_w0, pal5_mass,
                             dt=-1 * u.Myr, n_steps=4000)

pal5_stream.plot(alpha=0.1)