import astropy.units as u
import numpy as np
import gala.potential as gp
import gala.dynamics as gd
from gala.units import galactic
from gala.dynamics.mockstream import mock_stream

pot = gp.SphericalNFWPotential(v_c=175*u.km/u.s, r_s=10*u.kpc,
                               units=galactic)
prog_mass = 1E4*u.Msun
prog_w0 = gd.CartesianPhaseSpacePosition(pos=[15, 0, 0.]*u.kpc,
                                         vel=[75, 150, 0.]*u.km/u.s)
prog_orbit = pot.integrate_orbit(prog_w0, dt=0.5, n_steps=4000)
k_mean = [2., 0, 0, 0, 0.3, 0]
k_disp = [0.5, 0, 0.5, 0, 0.5, 0.5]
stream = mock_stream(pot, prog_orbit, prog_mass,
                     k_mean=k_mean, k_disp=k_disp, release_every=1)
stream.plot()