# Third-party
import astropy.coordinates as coord
import astropy.units as u
import numpy as np
import pytest
from astropy.coordinates import (
    CartesianDifferential,
    CartesianRepresentation,
    Galactic,
    SphericalCosLatDifferential,
    SphericalRepresentation,
)

from gala.tests.optional_deps import HAS_H5PY

from ...potential import Hamiltonian, HernquistPotential
from ...potential.frame import ConstantRotatingFrame, StaticFrame
from ...units import galactic, solarsystem

# Project
from ..core import PhaseSpacePosition


def test_initialize():
    with pytest.raises(ValueError):
        x = np.random.random(size=(3, 10))
        v = np.random.random(size=(3, 8))
        PhaseSpacePosition(pos=x, vel=v)

    x = np.random.random(size=(3, 10))
    v = np.random.random(size=(3, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    assert o.shape == (10,)

    x = np.random.random(size=(3, 10)) * u.kpc
    v = np.random.random(size=(3, 10)) * u.km / u.s
    o = PhaseSpacePosition(pos=x, vel=v)
    assert o.xyz.unit == u.kpc
    assert o.v_x.unit == u.km / u.s
    o.data
    assert "s" in o.data.differentials

    # Not 3D
    x = np.random.random(size=(2, 10))
    v = np.random.random(size=(2, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    assert o.ndim == 2

    o = PhaseSpacePosition(pos=x, vel=v, frame=StaticFrame(galactic))
    assert o.ndim == 2
    assert o.frame is not None
    assert isinstance(o.frame, StaticFrame)

    x = np.random.random(size=(4, 10))
    v = np.random.random(size=(4, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    assert o.ndim == 4

    # back to 3D
    pos = CartesianRepresentation(np.random.random(size=(3, 10)) * u.one)
    vel = CartesianDifferential(np.random.random(size=(3, 10)) * u.one)
    o = PhaseSpacePosition(pos=pos, vel=vel)
    assert hasattr(o, "x")
    assert hasattr(o, "y")
    assert hasattr(o, "z")
    assert hasattr(o, "v_x")
    assert hasattr(o, "v_y")
    assert hasattr(o, "v_z")

    # passing a representation with a differential attached
    pos = CartesianRepresentation(np.random.random(size=(3, 10)) * u.kpc)
    vel = CartesianDifferential(np.random.random(size=(3, 10)) * u.km / u.s)
    o = PhaseSpacePosition(pos.with_differentials({"s": vel}))
    assert hasattr(o, "x")
    assert hasattr(o, "y")
    assert hasattr(o, "z")
    assert hasattr(o, "v_x")
    assert hasattr(o, "v_y")
    assert hasattr(o, "v_z")

    o = o.represent_as(SphericalRepresentation)
    assert hasattr(o, "distance")
    assert hasattr(o, "lat")
    assert hasattr(o, "lon")
    assert hasattr(o, "radial_velocity")
    assert hasattr(o, "pm_lon")
    assert hasattr(o, "pm_lat")

    with pytest.raises(TypeError):
        o = PhaseSpacePosition(pos=x, vel=v, frame="blah blah blah")


def test_from_w():
    w = np.random.random(size=(6, 10))
    o = PhaseSpacePosition.from_w(w, galactic)
    assert o.x.unit == u.kpc
    assert o.v_x.unit == u.kpc / u.Myr
    assert o.shape == (10,)


def test_slice():
    # simple
    x = np.random.random(size=(3, 10))
    v = np.random.random(size=(3, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    new_o = o[:5]
    assert new_o.shape == (5,)

    # 1d slice on 3d
    x = np.random.random(size=(3, 10, 8))
    v = np.random.random(size=(3, 10, 8))
    o = PhaseSpacePosition(pos=x, vel=v)
    new_o = o[:5]
    assert new_o.shape == (5, 8)

    # 3d slice on 3d
    o = PhaseSpacePosition(pos=x, vel=v)
    new_o = o[:5, :4]
    assert new_o.shape == (5, 4)

    # boolean array
    x = np.random.random(size=(3, 10))
    v = np.random.random(size=(3, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    ix = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1]).astype(bool)
    new_o = o[ix]
    assert new_o.shape == (sum(ix),)

    # integer array
    x = np.random.random(size=(3, 10))
    v = np.random.random(size=(3, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    ix = np.array([0, 3, 5])
    new_o = o[ix]
    assert new_o.shape == (len(ix),)


def test_reshape():
    x = np.random.random(size=(3, 10))
    v = np.random.random(size=(3, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    new_o = o.reshape((10, 1))
    assert new_o.shape == (10, 1)
    assert new_o.x.shape == (10, 1)


def test_represent_as():
    # simple / unitless
    x = np.random.random(size=(3, 10))
    v = np.random.random(size=(3, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    new_o = o.represent_as(SphericalRepresentation)
    o.spherical
    o.cylindrical
    o.cartesian

    assert new_o.pos.distance.unit == u.one
    assert new_o.vel.d_distance.unit == u.one

    # simple / with units
    x = np.random.random(size=(3, 10)) * u.kpc
    v = np.random.normal(0.0, 100.0, size=(3, 10)) * u.km / u.s
    o = PhaseSpacePosition(pos=x, vel=v)
    sph = o.represent_as(SphericalRepresentation)
    assert sph.pos.distance.unit == u.kpc

    sph2 = o.represent_as("spherical")
    for c in sph.pos.components:
        assert u.allclose(getattr(sph.pos, c), getattr(sph2.pos, c), rtol=1e-12)

    # doesn't work for 2D
    x = np.random.random(size=(2, 10))
    v = np.random.random(size=(2, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    with pytest.raises(ValueError):
        o.represent_as(SphericalRepresentation)


def test_represent_as_expected_attributes():
    x = np.random.random(size=(3, 10))
    v = np.random.random(size=(3, 10))
    o = PhaseSpacePosition(pos=x, vel=v)

    new_o = o.spherical
    assert hasattr(new_o, "distance")
    assert hasattr(new_o, "lat")
    assert hasattr(new_o, "lon")
    assert hasattr(new_o, "radial_velocity")
    assert hasattr(new_o, "pm_lat")
    assert hasattr(new_o, "pm_lon")

    new_o = o.represent_as(SphericalRepresentation, SphericalCosLatDifferential)
    assert hasattr(new_o, "distance")
    assert hasattr(new_o, "lat")
    assert hasattr(new_o, "lon")
    assert hasattr(new_o, "radial_velocity")
    assert hasattr(new_o, "pm_lat")
    assert hasattr(new_o, "pm_lon_coslat")

    new_o = o.physicsspherical
    assert hasattr(new_o, "r")
    assert hasattr(new_o, "phi")
    assert hasattr(new_o, "theta")
    assert hasattr(new_o, "radial_velocity")
    assert hasattr(new_o, "pm_theta")
    assert hasattr(new_o, "pm_phi")

    new_o = o.cylindrical
    assert hasattr(new_o, "rho")
    assert hasattr(new_o, "phi")
    assert hasattr(new_o, "z")
    assert hasattr(new_o, "v_rho")
    assert hasattr(new_o, "pm_phi")
    assert hasattr(new_o, "v_z")

    new_o = new_o.cartesian
    assert hasattr(new_o, "x")
    assert hasattr(new_o, "y")
    assert hasattr(new_o, "z")
    assert hasattr(new_o, "xyz")
    assert hasattr(new_o, "v_x")
    assert hasattr(new_o, "v_y")
    assert hasattr(new_o, "v_z")
    assert hasattr(new_o, "v_xyz")

    # Check that this works with the NDCartesian classes too
    x = np.random.random(size=(2, 10)) * u.kpc
    v = np.random.normal(0.0, 100.0, size=(2, 10)) * u.km / u.s
    new_o = PhaseSpacePosition(pos=x, vel=v)

    assert hasattr(new_o, "x1")
    assert hasattr(new_o, "x2")
    assert hasattr(new_o, "xyz")
    assert hasattr(new_o, "v_x1")
    assert hasattr(new_o, "v_x2")
    assert hasattr(new_o, "v_xyz")


def test_to_coord_frame():
    # simple / unitless
    x = np.random.random(size=(3, 10))
    v = np.random.random(size=(3, 10))
    o = PhaseSpacePosition(pos=x, vel=v)

    with coord.galactocentric_frame_defaults.set("v4.0"):
        with pytest.raises(u.UnitConversionError):
            o.to_coord_frame(Galactic())

    # simple / with units
    x = np.random.random(size=(3, 10)) * u.kpc
    v = np.random.normal(0.0, 100.0, size=(3, 10)) * u.km / u.s
    o = PhaseSpacePosition(pos=x, vel=v)
    with coord.galactocentric_frame_defaults.set("v4.0"):
        coo = o.to_coord_frame(Galactic())
    assert coo.name == "galactic"

    # doesn't work for 2D
    x = np.random.random(size=(2, 10)) * u.kpc
    v = np.random.normal(0.0, 100.0, size=(2, 10)) * u.km / u.s
    o = PhaseSpacePosition(pos=x, vel=v)
    with coord.galactocentric_frame_defaults.set("v4.0"):
        with pytest.raises(ValueError):
            o.to_coord_frame(Galactic())


def test_w():
    # simple / unitless
    x = np.random.random(size=(3, 10))
    v = np.random.random(size=(3, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    w = o.w()
    assert w.shape == (6, 10)

    x = np.random.random(size=3)
    v = np.random.random(size=3)
    o = PhaseSpacePosition(pos=x, vel=v)
    w = o.w()
    assert w.shape == (6, 1)

    # simple / unitless, 2D
    x = np.random.random(size=(2, 10))
    v = np.random.random(size=(2, 10))
    o = PhaseSpacePosition(pos=x, vel=v)
    w = o.w()
    assert w.shape == (4, 10)

    x = np.random.random(size=2)
    v = np.random.random(size=2)
    o = PhaseSpacePosition(pos=x, vel=v)
    w = o.w()
    assert w.shape == (4, 1)

    # simple / with units
    x = np.random.random(size=(3, 10)) * u.kpc
    v = np.random.normal(0.0, 100.0, size=(3, 10)) * u.km / u.s
    o = PhaseSpacePosition(pos=x, vel=v)
    with pytest.raises(ValueError):
        o.w()
    w = o.w(units=galactic)
    assert np.allclose(x.value, w[:3])
    assert np.allclose(v.value, (w[3:] * u.kpc / u.Myr).to(u.km / u.s).value)

    # simple / with units and potential
    p = HernquistPotential(units=galactic, m=1e11, c=0.25)
    x = np.random.random(size=(3, 10)) * u.kpc
    v = np.random.normal(0.0, 100.0, size=(3, 10)) * u.km / u.s
    o = PhaseSpacePosition(pos=x, vel=v)
    w = o.w(p.units)
    assert np.allclose(x.value, w[:3])
    assert np.allclose(v.value, (w[3:] * u.kpc / u.Myr).to(u.km / u.s).value)

    w = o.w(units=solarsystem)
    assert np.allclose(x.value, (w[:3] * u.au).to(u.kpc).value)
    assert np.allclose(v.value, (w[3:] * u.au / u.yr).to(u.km / u.s).value)


# ------------------------------------------------------------------------
# Computed dynamical quantities
# ------------------------------------------------------------------------
def test_energy():  # noqa
    # with units
    x = np.random.random(size=(3, 10)) * u.kpc
    v = np.random.normal(0.0, 100.0, size=(3, 10)) * u.km / u.s
    o = PhaseSpacePosition(pos=x, vel=v)
    KE = o.kinetic_energy()
    assert KE.unit == (o.v_x.unit) ** 2
    assert KE.shape == o.shape

    # with units and potential
    p = HernquistPotential(units=galactic, m=1e11, c=0.25)
    H = Hamiltonian(p)
    x = np.random.random(size=(3, 10)) * u.kpc
    v = np.random.normal(0.0, 100.0, size=(3, 10)) * u.km / u.s
    o = PhaseSpacePosition(pos=x, vel=v)
    PE = o.potential_energy(p)  # noqa
    E = o.energy(H)  # noqa


def test_angular_momentum():
    w = PhaseSpacePosition([1.0, 0.0, 0.0], [0.0, 0.0, 1.0])
    assert u.allclose(np.squeeze(w.angular_momentum()), [0.0, -1, 0] * u.one)

    w = PhaseSpacePosition([1.0, 0.0, 0.0], [0.0, 1.0, 0.0])
    assert u.allclose(np.squeeze(w.angular_momentum()), [0.0, 0, 1] * u.one)

    w = PhaseSpacePosition([0.0, 1.0, 0.0], [0.0, 0.0, 1.0])
    assert u.allclose(np.squeeze(w.angular_momentum()), [1.0, 0, 0] * u.one)

    w = PhaseSpacePosition([1.0, 0, 0] * u.kpc, [0.0, 200.0, 0] * u.pc / u.Myr)
    assert u.allclose(
        np.squeeze(w.angular_momentum()), [0, 0, 0.2] * u.kpc**2 / u.Myr
    )

    # multiple - known
    q = np.array([[1.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0, 1.0, 0.0]]).T
    p = np.array([[0, 0, 1.0], [0, 1.0, 0.0], [0, 0, 1]]).T
    L = PhaseSpacePosition(q, p).angular_momentum()
    true_L = np.array([[0.0, -1, 0], [0.0, 0, 1], [1.0, 0, 0]]).T * u.one
    assert L.shape == (3, 3)
    assert u.allclose(L, true_L)

    # multiple - random
    q = np.random.uniform(size=(3, 128))
    p = np.random.uniform(size=(3, 128))
    L = PhaseSpacePosition(q, p).angular_momentum()
    assert L.shape == (3, 128)


def test_guiding_radius():  # noqa
    rng = np.random.default_rng(42)

    p = HernquistPotential(units=galactic, m=1e11, c=10.0)

    Rs = rng.uniform(4, 10, 128) * u.kpc
    xyz = Rs[None] * np.array([1.0, 0, 0])[:, None]

    vc = p.circular_velocity(xyz)
    vxyz = np.zeros((3, Rs.size)) * u.km / u.s
    vxyz[1] = rng.normal(vc.to_value(u.km / u.s), 15.0) * u.km / u.s

    w0 = PhaseSpacePosition(xyz, vxyz)
    Rgs = w0.guiding_radius(potential=p)
    assert np.all(Rgs > 0) and np.all(Rgs < 25 * u.kpc)


def test_frame_transform():
    static = StaticFrame(galactic)
    rotating = ConstantRotatingFrame(
        Omega=[0.53, 1.241, 0.9394] * u.rad / u.Myr, units=galactic
    )

    x = np.array([[10.0, -0.2, 0.3], [-0.232, 8.1, 0.1934]]).T * u.kpc
    v = np.array([[0.0034, 0.2, 0.0014], [0.0001, 0.002532, -0.2]]).T * u.kpc / u.Myr

    # no frame specified at init
    psp = PhaseSpacePosition(pos=x, vel=v)
    with pytest.raises(ValueError):
        psp.to_frame(rotating)

    psp.to_frame(rotating, current_frame=static, t=0.4 * u.Myr)

    # frame specified at init
    psp = PhaseSpacePosition(pos=x, vel=v, frame=static)
    psp.to_frame(rotating, t=0.4 * u.Myr)


@pytest.mark.parametrize(
    "obj",
    [
        PhaseSpacePosition([1, 2, 3.0] * u.kpc, [1, 2, 3.0] * u.km / u.s),
        PhaseSpacePosition(
            [1, 2, 3.0] * u.kpc, [1, 2, 3.0] * u.km / u.s, StaticFrame(units=galactic)
        ),
        PhaseSpacePosition(
            [1, 2, 3.0] * u.kpc,
            [1, 2, 3.0] * u.km / u.s,
            ConstantRotatingFrame(Omega=[1.0, 0, 0] * u.rad / u.Myr, units=galactic),
        ),
    ],
)
@pytest.mark.skipif(not HAS_H5PY, reason="h5py required for this test")
def test_io(tmpdir, obj):
    import h5py

    filename = str(tmpdir.join("thing.hdf5"))
    with h5py.File(filename, "w") as f:
        obj.to_hdf5(f)

    obj2 = PhaseSpacePosition.from_hdf5(filename)
    assert u.allclose(obj.xyz, obj2.xyz)
    assert u.allclose(obj.v_xyz, obj2.v_xyz)
    assert obj.frame == obj2.frame
