import astropy.units as u
import numpy as np
import gala.integrate as gi
import gala.dynamics as gd
import gala.potential as gp
from gala.units import galactic

np.random.seed(42)

pot = gp.SphericalNFWPotential(v_c=200*u.km/u.s, r_s=10.*u.kpc, units=galactic)

ics = gd.CartesianPhaseSpacePosition(pos=[10,0,0.]*u.kpc,
                                     vel=[0,175,0]*u.km/u.s)
orbit = pot.integrate_orbit(ics, dt=2., n_steps=2000)

norbits = 1024
new_pos = np.random.normal(ics.pos.to(u.pc).value, 100., size=(3,norbits))*u.pc
new_vel = np.random.normal(ics.vel.to(u.km/u.s).value, 1., size=(3,norbits))*u.km/u.s
new_ics = gd.CartesianPhaseSpacePosition(pos=new_pos, vel=new_vel)
orbits = pot.integrate_orbit(new_ics, dt=2., n_steps=2000)

grid = np.linspace(-15,15,64)
fig = pot.plot_contours(grid=(grid,grid,0), cmap='Greys')
ax = fig.axes[0] # grab the first plot axes
ax.plot(orbits.pos[0,-1], orbits.pos[1,-1], # this is x and y, the last timestep
        marker='.', linestyle='none', alpha=0.75, color='#cc0000')