import astropy.units as u
import matplotlib.pyplot as pl
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, 30.]*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)

fig,ax = pl.subplots(1,1,figsize=(6,6))

ax.plot(stream.pos[0], stream.pos[1], ls='none', marker='.', alpha=0.25)

x = prog_orbit[-1].pos.value
ax.set_xlim(x[0]-1., x[0]+1.)
ax.set_ylim(x[1]-1., x[1]+1.)

ax.set_xlabel("x [kpc]")
ax.set_ylabel("y [kpc]")