import copy

import numpy
import pytest

from galpy.actionAngle import actionAngleIsochroneApprox
from galpy.df import chen24spraydf, fardal15spraydf, streamdf, streamspraydf
from galpy.orbit import Orbit
from galpy.potential import (
    ChandrasekharDynamicalFrictionForce,
    HernquistPotential,
    LogarithmicHaloPotential,
    MovingObjectPotential,
    MWPotential2014,
    PlummerPotential,
    TriaxialNFWPotential,
)
from galpy.util import conversion  # for unit conversions
from galpy.util import coords

################################ Tests against streamdf ######################


def test_streamspraydf_deprecation():
    # Check if the deprecating class raises the correct warning
    lp = LogarithmicHaloPotential(normalize=1.0, q=0.9)
    obs = Orbit(
        [1.56148083, 0.35081535, -1.15481504, 0.88719443, -0.47713334, 0.12019596]
    )
    ro, vo = 8.0, 220.0
    with pytest.warns(DeprecationWarning):
        spdf = streamspraydf(
            2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
            progenitor=obs,
            pot=lp,
            tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
        )


# Setup both DFs
@pytest.fixture(scope="module")
def setup_testStreamsprayAgainstStreamdf():
    lp = LogarithmicHaloPotential(normalize=1.0, q=0.9)
    aAI = actionAngleIsochroneApprox(pot=lp, b=0.8)
    obs = Orbit(
        [1.56148083, 0.35081535, -1.15481504, 0.88719443, -0.47713334, 0.12019596]
    )
    ro, vo = 8.0, 220.0
    # Set up streamdf
    sigv = 0.365  # km/s
    sdf_bovy14 = streamdf(
        sigv / 220.0,
        progenitor=obs,
        pot=lp,
        aA=aAI,
        leading=True,
        nTrackChunks=11,
        tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
    )
    # Set up streamspraydf
    f15spdf_bovy14 = fardal15spraydf(
        2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
        progenitor=obs,
        pot=lp,
        tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
    )
    c24spdf_bovy14 = chen24spraydf(
        2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
        progenitor=obs,
        pot=lp,
        tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
    )
    return sdf_bovy14, [f15spdf_bovy14, c24spdf_bovy14]


def test_sample_bovy14(setup_testStreamsprayAgainstStreamdf):
    # Load objects that were setup above
    sdf_bovy14, spdfs_bovy14 = setup_testStreamsprayAgainstStreamdf
    for spdf_bovy14 in spdfs_bovy14:
        numpy.random.seed(1)
        RvR_sdf = sdf_bovy14.sample(n=1000)
        RvR_spdf = spdf_bovy14.sample(n=1000, integrate=True, return_orbit=False)
        # Sanity checks
        # Range in Z
        indx = (RvR_sdf[3] > 4.0 / 8.0) * (RvR_sdf[3] < 5.0 / 8.0)
        # mean
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[0][indx]) - numpy.mean(RvR_spdf[0][indx]))
            < 6e-2
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[1][indx]) - numpy.mean(RvR_spdf[1][indx]))
            < 5e-2
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[2][indx]) - numpy.mean(RvR_spdf[2][indx]))
            < 5e-2
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[4][indx]) - numpy.mean(RvR_spdf[4][indx]))
            < 5e-2
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[5][indx]) - numpy.mean(RvR_spdf[5][indx]))
            < 1e-1
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
        # Another range in Z
        indx = (RvR_sdf[3] > 5.0 / 8.0) * (RvR_sdf[3] < 6.0 / 8.0)
        # mean
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[0][indx]) - numpy.mean(RvR_spdf[0][indx]))
            < 1e-1
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[1][indx]) - numpy.mean(RvR_spdf[1][indx]))
            < 3e-2
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[2][indx]) - numpy.mean(RvR_spdf[2][indx]))
            < 4e-2
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[4][indx]) - numpy.mean(RvR_spdf[4][indx]))
            < 3e-2
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
        assert (
            numpy.fabs(numpy.mean(RvR_sdf[5][indx]) - numpy.mean(RvR_spdf[5][indx]))
            < 1e-1
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean)"
        )
    return None


def test_bovy14_sampleorbit(setup_testStreamsprayAgainstStreamdf):
    # Load objects that were setup above
    sdf_bovy14, spdfs_bovy14 = setup_testStreamsprayAgainstStreamdf
    for spdf_bovy14 in spdfs_bovy14:
        numpy.random.seed(1)
        XvX_sdf = sdf_bovy14.sample(n=1000, xy=True)
        XvX_spdf = spdf_bovy14.sample(
            n=1000
        )  # returns Orbit, from which we can get anything we want
        # Sanity checks
        # Range in Z
        indx = (XvX_sdf[2] > 4.0 / 8.0) * (XvX_sdf[2] < 5.0 / 8.0)
        # mean
        assert (
            numpy.fabs(numpy.mean(XvX_sdf[0][indx]) - numpy.mean(XvX_spdf.x()[indx]))
            < 6e-2
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean, xy)"
        )
        assert (
            numpy.fabs(numpy.mean(XvX_sdf[1][indx]) - numpy.mean(XvX_spdf.y()[indx]))
            < 2e-1
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean, xy)"
        )
        assert (
            numpy.fabs(numpy.mean(XvX_sdf[4][indx]) - numpy.mean(XvX_spdf.vy()[indx]))
            < 3e-2
        ), (
            "streamdf and streamspraydf do not generate similar samples for the Bovy (2014) stream (mean, xy)"
        )
    return None


def test_integrate(setup_testStreamsprayAgainstStreamdf):
    # Test that sampling at stripping + integrate == sampling at the end
    # Load objects that were setup above
    _, spdfs_bovy14 = setup_testStreamsprayAgainstStreamdf
    for spdf_bovy14 in spdfs_bovy14:
        # Sample at at stripping
        numpy.random.seed(4)
        RvR_noint, dt_noint = spdf_bovy14.sample(
            n=100, return_orbit=False, returndt=True, integrate=False
        )
        # and integrate
        for ii in range(len(dt_noint)):
            to = Orbit(RvR_noint[:, ii])
            to.integrate(numpy.linspace(-dt_noint[ii], 0.0, 1001), spdf_bovy14._pot)
            RvR_noint[:, ii] = [
                to.R(0.0),
                to.vR(0.0),
                to.vT(0.0),
                to.z(0.0),
                to.vz(0.0),
                to.phi(0.0),
            ]
        # Sample today
        numpy.random.seed(4)
        RvR, dt = spdf_bovy14.sample(
            n=100, return_orbit=False, returndt=True, integrate=True
        )
        # Should agree
        assert numpy.amax(numpy.fabs(dt - dt_noint)) < 1e-10, (
            "Times not the same when sampling with and without integrating"
        )
        assert numpy.amax(numpy.fabs(RvR - RvR_noint)) < 1e-7, (
            "Phase-space points not the same when sampling with and without integrating"
        )
    return None


def test_integrate_rtnonarray():
    # Test that sampling at stripping + integrate == sampling at the end
    # For a potential that doesn't support array inputs
    nfp = TriaxialNFWPotential(normalize=1.0, b=0.9, c=0.8)
    obs = Orbit(
        [1.56148083, 0.35081535, -1.15481504, 0.88719443, -0.47713334, 0.12019596]
    )
    ro, vo = 8.0, 220.0
    for streamspraydf in [fardal15spraydf, chen24spraydf]:
        # Set up streamspraydf
        spdf_bovy14 = streamspraydf(
            2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
            progenitor=obs,
            pot=nfp,
            tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
        )
        # Sample at at stripping
        numpy.random.seed(4)
        RvR_noint, dt_noint = spdf_bovy14.sample(
            n=100, return_orbit=False, returndt=True, integrate=False
        )
        # and integrate
        for ii in range(len(dt_noint)):
            to = Orbit(RvR_noint[:, ii])
            to.integrate(numpy.linspace(-dt_noint[ii], 0.0, 1001), spdf_bovy14._pot)
            RvR_noint[:, ii] = [
                to.R(0.0),
                to.vR(0.0),
                to.vT(0.0),
                to.z(0.0),
                to.vz(0.0),
                to.phi(0.0),
            ]
        # Sample today
        numpy.random.seed(4)
        RvR, dt = spdf_bovy14.sample(
            n=100, return_orbit=False, returndt=True, integrate=True
        )
        # Should agree
        assert numpy.amax(numpy.fabs(dt - dt_noint)) < 1e-10, (
            "Times not the same when sampling with and without integrating"
        )
        assert numpy.amax(numpy.fabs(RvR - RvR_noint)) < 1e-7, (
            "Phase-space points not the same when sampling with and without integrating"
        )
    return None


def test_center():
    # Test that a stream around a different center is generated
    # when using center
    # In this example, we'll generate a stream in the LMC orbiting the MW
    # LMC and its orbit
    ro, vo = 8.0, 220.0
    o = Orbit(
        [5.13200034, 1.08033051, 0.2332339, -3.48068653, 0.94950884, -1.54626091]
    )  # Result from from_name('LMC')
    tMWPotential2014 = copy.deepcopy(MWPotential2014)
    tMWPotential2014[2] *= 1.5
    cdf = ChandrasekharDynamicalFrictionForce(
        GMs=10 / conversion.mass_in_1010msol(vo, ro),
        rhm=5.0 / ro,
        dens=tMWPotential2014,
    )
    ts = numpy.linspace(0.0, -10.0, 1001) / conversion.time_in_Gyr(vo, ro)
    o.integrate(ts, tMWPotential2014 + cdf)
    lmcpot = HernquistPotential(
        amp=2 * 10 / conversion.mass_in_1010msol(vo, ro),
        a=5.0 / ro / (1.0 + numpy.sqrt(2.0)),
    )  # rhm = (1+sqrt(2)) a
    moving_lmcpot = MovingObjectPotential(o, pot=lmcpot)
    # Now generate a stream within the LMC, progenitor at 8x kpc on circular orbit
    of = o(ts[-1])  # LMC at final point, earliest time, for convenience
    # Following pos in kpc, vel in km/s
    R_in_lmc = 1.0
    prog_phasespace = (
        of.x(use_physical=False) + R_in_lmc,
        of.y(use_physical=False),
        of.z(use_physical=False),
        of.vx(use_physical=False),
        of.vy(use_physical=False) + lmcpot.vcirc(R_in_lmc, use_physical=False),
        of.vz(use_physical=False),
    )
    prog_pos = coords.rect_to_cyl(
        prog_phasespace[0], prog_phasespace[1], prog_phasespace[2]
    )
    prog_vel = coords.rect_to_cyl_vec(
        prog_phasespace[3],
        prog_phasespace[4],
        prog_phasespace[5],
        None,
        prog_pos[1],
        None,
        cyl=True,
    )
    prog = Orbit(
        [prog_pos[0], prog_vel[0], prog_vel[1], prog_pos[2], prog_vel[2], prog_pos[1]],
        ro=8.0,
        vo=220.0,
    )
    # Integrate prog forward
    prog.integrate(ts[::-1], tMWPotential2014 + moving_lmcpot)
    for streamspraydf in [fardal15spraydf, chen24spraydf]:
        # Then set up streamspraydf
        spdf = streamspraydf(
            2e4 / conversion.mass_in_msol(vo, ro),
            progenitor=prog(0.0),
            pot=tMWPotential2014 + moving_lmcpot,
            rtpot=lmcpot,
            tdisrupt=10.0 / conversion.time_in_Gyr(vo, ro),
            center=o,
            centerpot=tMWPotential2014 + cdf,
        )
        # Generate stream
        numpy.random.seed(1)
        stream_RvR = spdf.sample(n=300, return_orbit=False, integrate=True)
        stream_pos = coords.cyl_to_rect(stream_RvR[0], stream_RvR[5], stream_RvR[3])
        # Stream should lie on a circle with radius R_in_lmc
        stream_R_wrt_LMC = numpy.sqrt(
            (stream_pos[0] - o.x(use_physical=False)) ** 2.0
            + (stream_pos[1] - o.y(use_physical=False)) ** 2.0
        )
        assert numpy.fabs(numpy.mean(stream_R_wrt_LMC) - R_in_lmc) < 0.1, (
            "Stream generated in the LMC does not appear to be on a circle within the LMC"
        )
        assert numpy.fabs(numpy.std(stream_R_wrt_LMC)) < 0.15, (
            "Stream generated in the LMC does not appear to be on a circle within the LMC"
        )
    return None


def test_sample_orbit_rovoetc():
    # Test that the sample orbit output has the same ro/vo/etc. as the
    # input progenitor
    lp = LogarithmicHaloPotential(normalize=1.0, q=0.9)
    ro, vo = 9.0, 230.0
    zo, solarmotion = 0.03, [-20.0, 30.0, 40.0]
    obs = Orbit(
        [1.56148083, 0.35081535, -1.15481504, 0.88719443, -0.47713334, 0.12019596],
        ro=ro,
        vo=vo,
        zo=zo,
        solarmotion=solarmotion,
    )
    for streamspraydf in [fardal15spraydf, chen24spraydf]:
        # Set up streamspraydf
        spdf_bovy14 = streamspraydf(
            2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
            progenitor=obs,
            pot=lp,
            tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
        )
        sam = spdf_bovy14.sample(n=10)
        assert obs._roSet is sam._roSet, (
            "Sampled streamspraydf orbits do not have the same roSet as the progenitor orbit"
        )
        assert obs._voSet is sam._voSet, (
            "Sampled streamspraydf orbits do not have the same voSet as the progenitor orbit"
        )
        assert numpy.fabs(obs._ro - sam._ro) < 1e-10, (
            "Sampled streamspraydf orbits do not have the same ro as the progenitor orbit"
        )
        assert numpy.fabs(obs._vo - sam._vo) < 1e-10, (
            "Sampled streamspraydf orbits do not have the same vo as the progenitor orbit"
        )
        assert numpy.fabs(obs._zo - sam._zo) < 1e-10, (
            "Sampled streamspraydf orbits do not have the same zo as the progenitor orbit"
        )
        assert numpy.all(numpy.fabs(obs._solarmotion - sam._solarmotion) < 1e-10), (
            "Sampled streamspraydf orbits do not have the same solarmotion as the progenitor orbit"
        )
    # Another one
    ro = 9.0
    zo, solarmotion = 0.03, [-20.0, 30.0, 40.0]
    obs = Orbit(
        [1.56148083, 0.35081535, -1.15481504, 0.88719443, -0.47713334, 0.12019596],
        ro=ro,
        zo=zo,
        solarmotion=solarmotion,
    )
    for streamspraydf in [fardal15spraydf, chen24spraydf]:
        # Set up streamspraydf
        spdf_bovy14 = streamspraydf(
            2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
            progenitor=obs,
            pot=lp,
            tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
        )
        sam = spdf_bovy14.sample(n=10)
        assert obs._roSet, (
            "Test requires that ro be set for the progenitor orbit, but it appears not to have been set"
        )
        assert not obs._voSet, (
            "Test requires that vo not be set for the progenitor orbit, but it appears to have been set"
        )
        assert obs._roSet is sam._roSet, (
            "Sampled streamspraydf orbits do not have the same roSet as the progenitor orbit"
        )
        assert obs._voSet is sam._voSet, (
            "Sampled streamspraydf orbits do not have the same voSet as the progenitor orbit"
        )
        assert numpy.fabs(obs._ro - sam._ro) < 1e-10, (
            "Sampled streamspraydf orbits do not have the same ro as the progenitor orbit"
        )
        assert numpy.fabs(obs._vo - sam._vo) < 1e-10, (
            "Sampled streamspraydf orbits do not have the same vo as the progenitor orbit"
        )
        assert numpy.fabs(obs._zo - sam._zo) < 1e-10, (
            "Sampled streamspraydf orbits do not have the same zo as the progenitor orbit"
        )
        assert numpy.all(numpy.fabs(obs._solarmotion - sam._solarmotion) < 1e-10), (
            "Sampled streamspraydf orbits do not have the same solarmotion as the progenitor orbit"
        )
    # And another one
    vo = 230.0
    zo, solarmotion = 0.03, [-20.0, 30.0, 40.0]
    obs = Orbit(
        [1.56148083, 0.35081535, -1.15481504, 0.88719443, -0.47713334, 0.12019596],
        vo=vo,
        zo=zo,
        solarmotion=solarmotion,
    )
    for streamspraydf in [fardal15spraydf, chen24spraydf]:
        # Set up streamspraydf
        spdf_bovy14 = streamspraydf(
            2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
            progenitor=obs,
            pot=lp,
            tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
        )
        sam = spdf_bovy14.sample(n=10)
        assert obs._voSet, (
            "Test requires that vo be set for the progenitor orbit, but it appears not to have been set"
        )
        assert not obs._roSet, (
            "Test requires that ro not be set for the progenitor orbit, but it appears to have been set"
        )
        assert obs._roSet is sam._roSet, (
            "Sampled streamspraydf orbits do not have the same roSet as the progenitor orbit"
        )
        assert obs._voSet is sam._voSet, (
            "Sampled streamspraydf orbits do not have the same voSet as the progenitor orbit"
        )
        assert numpy.fabs(obs._ro - sam._ro) < 1e-10, (
            "Sampled streamspraydf orbits do not have the same ro as the progenitor orbit"
        )
        assert numpy.fabs(obs._vo - sam._vo) < 1e-10, (
            "Sampled streamspraydf orbits do not have the same vo as the progenitor orbit"
        )
        assert numpy.fabs(obs._zo - sam._zo) < 1e-10, (
            "Sampled streamspraydf orbits do not have the same zo as the progenitor orbit"
        )
        assert numpy.all(numpy.fabs(obs._solarmotion - sam._solarmotion) < 1e-10), (
            "Sampled streamspraydf orbits do not have the same solarmotion as the progenitor orbit"
        )
    return None


def test_integrate_with_prog():
    # Test integrating orbits with the progenitor's potential
    lp = LogarithmicHaloPotential(normalize=1.0, q=0.9)
    obs = Orbit(
        [1.56148083, 0.35081535, -1.15481504, 0.88719443, -0.47713334, 0.12019596]
    )
    ro, vo = 8.0, 220.0
    # Without the progenitor's potential
    spdf = chen24spraydf(
        2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
        progenitor=obs,
        pot=lp,
        tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
    )
    numpy.random.seed(4)
    RvR, dt = spdf.sample(n=100, return_orbit=False, returndt=True, integrate=True)
    # With the progenitor's potential, but set to zero-mass
    spdf = chen24spraydf(
        2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
        progenitor=obs,
        pot=lp,
        tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
        progpot=PlummerPotential(0, 0),
    )
    numpy.random.seed(4)
    RvR_withprog, dt_withprog = spdf.sample(
        n=100, return_orbit=False, returndt=True, integrate=True
    )
    # Should agree
    assert numpy.amax(numpy.fabs(dt - dt_withprog)) < 1e-10, (
        "Times not the same when sampling with and without prognitor's potential"
    )
    assert numpy.amax(numpy.fabs(RvR - RvR_withprog)) < 1e-7, (
        "Phase-space points not the same when sampling with and without prognitor's potential"
    )
    return None


def test_chen24spraydf_default_parameters():
    # Test the default parameters of chen24spraydf can be changed
    lp = LogarithmicHaloPotential(normalize=1.0, q=0.9)
    obs = Orbit(
        [1.56148083, 0.35081535, -1.15481504, 0.88719443, -0.47713334, 0.12019596]
    )
    ro, vo = 8.0, 220.0
    # Default parameters
    spdf = chen24spraydf(
        2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
        progenitor=obs,
        pot=lp,
        tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
    )
    numpy.random.seed(4)
    RvR_default, dt_default = spdf.sample(
        n=100, return_orbit=False, returndt=True, integrate=True
    )
    # Modified parameters, but only slightly
    spdf = chen24spraydf(
        2 * 10.0**4.0 / conversion.mass_in_msol(vo, ro),
        progenitor=obs,
        pot=lp,
        tdisrupt=4.5 / conversion.time_in_Gyr(vo, ro),
        mean=numpy.array([1.6, -0.525344, 0, 1, 0.349066, 0]),
        cov=numpy.array(
            [
                [0.1225, 0, 0, 0, -0.085521, 0],
                [0, 0.161143, 0, 0, 0, 0],
                [0, 0, 0.043865, 0, 0, 0],
                [0, 0, 0, 0, 0, 0],
                [-0.085521, 0, 0, 0, 0.121847, 0],
                [0, 0, 0, 0, 0, 0.147435],
            ]
        ),
    )
    numpy.random.seed(4)
    RvR, dt = spdf.sample(n=100, return_orbit=False, returndt=True, integrate=True)
    # Should agree
    assert numpy.amax(numpy.fabs(dt_default - dt)) < 1e-10, (
        "Times not the same when changing the default parameters"
    )
    assert numpy.amax(numpy.fabs(RvR_default - RvR)) > 1e-7, (
        "Phase-space points should not be the same when changing the default parameters"
    )
    assert numpy.amax(numpy.fabs(RvR_default - RvR)) < 1e-2, (
        "Phase-space points too different when sampling with and without prognitor's potential"
    )
    return None