import astropy.units as u
import numpy as np
import pytest

from gala.dynamics import (
    Orbit,
    PhaseSpacePosition,
    combine,
    estimate_dt_n_steps,
    peak_to_peak_period,
)
from gala.potential import Hamiltonian, NFWPotential, StaticFrame
from gala.units import galactic


def test_peak_to_peak_period():
    ntimes = 16384

    # trivial test
    for true_T in [1.0, 2.0, 4.123]:
        t = np.linspace(0, 10.0, ntimes)
        f = np.sin(2 * np.pi / true_T * t)
        T = peak_to_peak_period(t, f)
        assert np.allclose(T, true_T, atol=1e-3)

    # modulated trivial test
    true_T = 2.0
    t = np.linspace(0, 10.0, ntimes)
    f = np.sin(2 * np.pi / true_T * t) + 0.1 * np.cos(2 * np.pi / (10 * true_T) * t)
    T = peak_to_peak_period(t, f)
    assert np.allclose(T, true_T, atol=1e-3)


def test_estimate_dt_n_steps():
    nperiods = 128
    pot = NFWPotential.from_circular_velocity(v_c=1.0, r_s=10.0, units=galactic)
    w0 = [10.0, 0.0, 0.0, 0.0, 0.9, 0.0]

    H = Hamiltonian(pot)
    dt, n_steps = estimate_dt_n_steps(
        w0, H, n_periods=nperiods, n_steps_per_period=256, func=np.nanmin
    )

    orbit = H.integrate_orbit(w0, dt=dt, n_steps=n_steps)
    T = orbit.physicsspherical.estimate_period()["r"]
    assert int(np.round((orbit.t.max() / T).decompose().value)) == nperiods


class TestCombine:
    def setup_method(self):
        x = np.random.random(size=(3,))
        v = np.random.random(size=(3,))
        p1 = PhaseSpacePosition(pos=x, vel=v)
        p2 = PhaseSpacePosition(pos=x, vel=v, frame=StaticFrame(galactic))
        x = np.random.random(size=(3, 5))
        v = np.random.random(size=(3, 5))
        p3 = PhaseSpacePosition(pos=x, vel=v)
        p4 = PhaseSpacePosition(pos=x * u.kpc, vel=v * u.km / u.s)
        x = np.random.random(size=(2, 5))
        v = np.random.random(size=(2, 5))
        p5 = PhaseSpacePosition(pos=x, vel=v)
        self.psps = [p1, p2, p3, p4, p5]

        x = np.random.random(size=(3, 8))
        v = np.random.random(size=(3, 8))
        o1 = Orbit(pos=x, vel=v)
        o2 = Orbit(pos=x, vel=v, t=np.arange(8))

        pot = NFWPotential.from_circular_velocity(v_c=1.0, r_s=10.0, units=galactic)
        o3 = Orbit(
            pos=x * u.kpc,
            vel=v * u.km / u.s,
            t=np.arange(8) * u.Myr,
            potential=pot,
            frame=StaticFrame(galactic),
        )

        x = np.random.random(size=(2, 8))
        v = np.random.random(size=(2, 8))
        o4 = Orbit(pos=x, vel=v, t=np.arange(8))
        self.orbs = [o1, o2, o3, o4]

    def test_combine_fail(self):
        with pytest.raises(ValueError):
            combine([])

        with pytest.raises(ValueError):
            combine(self.psps[0])

        with pytest.raises(TypeError):
            combine([self.psps[0], self.orbs[0]])

        with pytest.raises(TypeError):
            combine([5, 5, 5])

        with pytest.raises(ValueError):
            combine(self.psps)

        with pytest.raises(ValueError):
            combine(self.orbs)

    def test_combine_psp(self):
        for psp in self.psps:
            psps = [psp] * 3
            new_psp = combine(psps)
            assert new_psp.ndim == psp.ndim

            shp = psp.pos.shape if psp.pos.shape else (1,)

            assert new_psp.pos.shape == (3 * shp[0],)
            assert new_psp.frame == psp.frame

    def test_combine_orb(self):
        for orb in self.orbs:
            orbs = [orb] * 4
            new_orb = combine(orbs)
            assert new_orb.ndim == orb.ndim

            shp = orb.shape
            shp = (*shp, 4) if len(shp) < 2 else (*shp[:-1], 4 * shp[-1])

            assert new_orb.pos.shape == shp
            assert new_orb.frame == orb.frame
            assert new_orb.potential == orb.potential
