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

pot = gp.NFWPotential.from_circular_velocity(v_c=220*u.km/u.s,
                                             r_s=15*u.kpc,
                                             units=galactic)
H = gp.Hamiltonian(pot)
prog_w0 = gd.PhaseSpacePosition(pos=[10, 0, 0.] * u.kpc,
                                vel=[0, 170, 0.] * u.km/u.s)

df = ms.FardalStreamDF()
prog_mass = 2.5E4 * u.Msun

gen = ms.MockStreamGenerator(df, H)

stream, prog = gen.run(prog_w0, prog_mass,
                       dt=1 * u.Myr, n_steps=1000)

stream.plot(['x', 'y'], marker='o', s=4, color='k', alpha=0.1, linewidth=0)