import pylab as plt
from mw_plot import MWPlot

from galpy.potential import MWPotential2014
from galpy.orbit import Orbit
import numpy as np
from astropy import units as u

# Orbit Integration using galpy for the Sun
op = Orbit([0., 0., 0., 0., 0., 0.], radec=True, ro=8., vo=220.)
ts = np.linspace(0, 5, 10000) * u.Gyr
op.integrate(ts, MWPotential2014)
x = op.x(ts) * u.kpc
y = op.y(ts) * u.kpc
z = op.z(ts)

# setup a mw-plot instance of bird's eye view of the disc
mw1 = MWPlot(radius=20 * u.kpc, unit=u.kpc, coord='galactocentric', annotation=True, figsize=(15, 12), r0=8)

# set up plot title
mw1.title = 'Orbit of Sun in 5Gyr'

# use mw_scatter instead of scatter because we want a colorbar
mw1.mw_scatter(x, y, c=[z, 'kpc above galactic plane'], s=1)
plt.tight_layout()