import matplotlib.pyplot as pl
import numpy as np
import gala.dynamics as gd
import gala.potential as gp

class HenonHeilesPotential(gp.PotentialBase):

    def __init__(self, A, units):
        pars = dict(A=A)
        super(HenonHeilesPotential, self).__init__(units=units,
                                                   parameters=pars)

    def _value(self, q, t):
        A = self.parameters['A']
        x,y = q
        return 0.5*(x**2 + y**2) + A*(x**2*y - y**3/3)

    def _gradient(self, q, t):
        A = self.parameters['A']
        x,y = q
        print(x)
        grad = np.zeros_like(q)
        grad[0] = x + 2*A*x*y
        grad[1] = y + A*(x**2 - y**2)
        return grad

pot = HenonHeilesPotential(A=1., units=None)
w0 = gd.CartesianPhaseSpacePosition(pos=[0.,0.3],
                                    vel=[0.38,0.])
orbit = pot.integrate_orbit(w0, dt=0.05, n_steps=10000)
fig = orbit.plot()