import gala.integrate as gi
import scipy.optimize as so

pot = gp.CCompositePotential()
pot['bar'] = gp.LongMuraliBarPotential(m=2E10*u.Msun, a=4*u.kpc,
                                       b=0.5*u.kpc, c=0.5*u.kpc,
                                       alpha=25*u.degree,
                                       units=galactic)
pot['disk'] = gp.MiyamotoNagaiPotential(m=5E10*u.Msun, a=3.*u.kpc,
                                        b=280.*u.pc, units=galactic)
pot['halo'] = gp.NFWPotential(m=6E11*u.Msun, r_s=20.*u.kpc,
                              units=galactic)

Om_bar = 42. * u.km/u.s/u.kpc
frame = gp.ConstantRotatingFrame(Omega=[0,0,Om_bar.value]*Om_bar.unit,
                                 units=galactic)
H = gp.Hamiltonian(potential=pot, frame=frame)

def func(r):
    Om = pot.circular_velocity([r[0], 0, 0]*u.kpc)[0] / (r[0]*u.kpc)
    return (Om - Om_bar).to(Om_bar.unit).value**2

res = so.minimize(func, x0=10., method='powell')
r_corot = res.x[0] * u.kpc
v_circ = Om_bar * r_corot

w0 = gd.PhaseSpacePosition(pos=[r_corot.value, 0, 0] * r_corot.unit,
                           vel=[0,v_circ.value, 0.] * v_circ.unit)

orbit = H.integrate_orbit(w0, dt=0.1, n_steps=40000,
                          Integrator=gi.DOPRI853Integrator)

fig = orbit.plot(marker=',', linestyle='none', alpha=0.5) # doctest: +SKIP
for ax in fig.axes:
    ax.set_xlim(-15,15)
    ax.set_ylim(-15,15)