##################
# Toy Potentials #
##################
import numpy as np
from scipy.optimize import leastsq

# in units kpc, km/s and 10^11 M_solar
# Grav = 430091.5694
Grav = 430091.7270069976  # This was a bug in Sanders' code!

# Triaxial harmonic


def H_ho(x, omega):
    """Simple harmonic oscillator Hamiltonian = 0.5 * omega**2 * x**2"""
    return 0.5 * np.sum(x[3:] ** 2 + (omega * x[:3]) ** 2)


def angact_ho(x, omega):
    """Calculate angle and action variable in sho potential with
    parameter omega"""
    action = (x[3:] ** 2 + (omega * x[:3]) ** 2) / (2.0 * omega)
    angle = np.array(
        [
            np.arctan(-x[3 + i] / omega[i] / x[i])
            if x[i] != 0.0
            else -np.sign(x[3 + i]) * np.pi / 2.0
            for i in range(3)
        ]
    )
    for i in range(3):
        if x[i] < 0:
            angle[i] += np.pi
    return np.concatenate((action, angle % (2.0 * np.pi)))


def deltaH_ho(omega, xsamples):
    if np.any(omega < 1e-5):
        return np.nan
    H = 0.5 * np.sum(xsamples.T[3:] ** 2, axis=0) + 0.5 * np.sum(
        (omega[:3] * xsamples.T[:3].T) ** 2, axis=1
    )
    return H - np.mean(H)


def Jac_deltaH_ho(omega, xsamples):
    dHdparams = omega[:3] * xsamples.T[:3].T ** 2
    return dHdparams - np.mean(dHdparams, axis=0)


def findbestparams_ho(xsamples):
    """Minimize sum of square differences of H_sho-<H_sho> for timesamples"""
    return np.abs(
        leastsq(
            deltaH_ho,
            np.array([10.0, 10.0, 10.0]),
            Dfun=Jac_deltaH_ho,
            args=(xsamples,),
        )[0]
    )[:3]


# Isochrone


def cart2spol(X):
    """Performs coordinate transformation from cartesian
    to spherical polar coordinates with (r, phi, theta) having
    usual meanings."""
    x, y, z, vx, vy, vz = X
    r = np.sqrt(x * x + y * y + z * z)
    p = np.arctan2(y, x)
    t = np.arccos(z / r)
    vr = (vx * np.cos(p) + vy * np.sin(p)) * np.sin(t) + np.cos(t) * vz
    vp = -vx * np.sin(p) + vy * np.cos(p)
    vt = (vx * np.cos(p) + vy * np.sin(p)) * np.cos(t) - np.sin(t) * vz
    return np.array([r, p, t, vr, vp, vt])


def H_iso(x, params):
    """Isochrone Hamiltonian = -GM/(b+sqrt(b**2+(r-r0)**2))"""
    # r = (np.sqrt(np.sum(x[:3]**2))-params[2])**2
    r = np.sum(x[:3] ** 2)
    return 0.5 * np.sum(x[3:] ** 2) - Grav * params[0] / (
        params[1] + np.sqrt(params[1] ** 2 + r)
    )


def angact_iso(x, params):
    """Calculate angle and action variable in isochrone potential with
    parameters params = (M, b)"""
    GM = Grav * params[0]
    E = H_iso(x, params)
    r, p, t, vr, vphi, vt = cart2spol(x)
    st = np.sin(t)
    Lz = r * vphi * st
    L = np.sqrt(r * r * vt * vt + Lz * Lz / st / st)
    if E > 0.0:  # Unbound
        return (np.nan, np.nan, np.nan, np.nan, np.nan, np.nan)
    Jr = GM / np.sqrt(-2 * E) - 0.5 * (L + np.sqrt(L * L + 4 * GM * params[1]))
    action = np.array([Jr, Lz, L - abs(Lz)])

    c = GM / (-2 * E) - params[1]
    e = np.sqrt(1 - L * L * (1 + params[1] / c) / GM / c)
    eta = np.arctan2(
        r * vr / np.sqrt(-2.0 * E), params[1] + c - np.sqrt(params[1] ** 2 + r * r)
    )
    OmR = np.power(-2 * E, 1.5) / GM
    Omp = 0.5 * OmR * (1 + L / np.sqrt(L * L + 4 * GM * params[1]))
    thetar = eta - e * c * np.sin(eta) / (c + params[1])

    if abs(vt) > 1e-10:
        psi = np.arctan2(np.cos(t), -np.sin(t) * r * vt / L)
    else:
        psi = np.pi / 2.0
    a = np.sqrt((1 + e) / (1 - e))
    ap = np.sqrt((1 + e + 2 * params[1] / c) / (1 - e + 2 * params[1] / c))

    def F(x, y):
        return (
            np.pi / 2.0 - np.arctan(np.tan(np.pi / 2.0 - 0.5 * y) / x)
            if y > np.pi / 2.0
            else -np.pi / 2.0 + np.arctan(np.tan(np.pi / 2.0 + 0.5 * y) / x)
            if y < -np.pi / 2.0
            else np.arctan(x * np.tan(0.5 * y))
        )

    thetaz = (
        psi
        + Omp * thetar / OmR
        - F(a, eta)
        - F(ap, eta) / np.sqrt(1 + 4 * GM * params[1] / L / L)
    )

    LR = Lz / L
    sinu = LR / np.sqrt(1.0 - LR**2) / np.tan(t)
    u = 0
    if sinu > 1.0:
        u = np.pi / 2.0
    elif sinu < -1.0:
        u = -np.pi / 2.0
    else:
        u = np.arcsin(sinu)
    if vt > 0.0:
        u = np.pi - u
    thetap = p - u + np.sign(Lz) * thetaz
    angle = np.array([thetar, thetap, thetaz])
    return np.concatenate((action, angle % (2.0 * np.pi)))


def deltaH_iso(params, p, r):
    deltaH = p - Grav * params[0] / (params[1] + np.sqrt(params[1] ** 2 + r))
    if params[0] < 0.0 or params[1] < 0.0 or np.any(deltaH > 0.0):
        return np.nan
    return deltaH - np.mean(deltaH)
    # return JR-np.mean(JR)


def Jac_deltaH_iso(params, p, r):
    H_o = -Grav / (params[1] + np.sqrt(params[1] ** 2 + r))
    H_1 = (
        Grav
        * params[0]
        * (1.0 + params[1] / np.sqrt(params[1] ** 2 + r))
        / (params[1] + np.sqrt(params[1] ** 2 + r)) ** 2
    )
    return np.array([(H_o - np.mean(H_o)), (H_1 - np.mean(H_1))])


def findbestparams_iso(xsamples):
    """Minimize sum of square differences of H_iso-<H_iso> for timesamples"""
    p = 0.5 * np.sum(xsamples.T[3:] ** 2, axis=0)
    r = np.sum(xsamples.T[:3] ** 2, axis=0)
    return np.abs(
        leastsq(
            deltaH_iso,
            np.array([10.0, 10.0]),
            Dfun=None,
            col_deriv=1,
            args=(
                p,
                r,
            ),
        )[0]
    )  # Jac_deltaH_iso
