import os
from urllib.error import HTTPError, URLError

import numpy as np
import pytest

from astropy import units as u
from astropy.constants import c
from astropy.coordinates.builtin_frames import TETE
from astropy.coordinates.earth import EarthLocation
from astropy.coordinates.funcs import get_sun
from astropy.coordinates.representation import (
    CartesianRepresentation,
    UnitSphericalRepresentation,
)
from astropy.coordinates.sky_coordinate import SkyCoord
from astropy.coordinates.solar_system import (
    BODY_NAME_TO_KERNEL_SPEC,
    _get_apparent_body_position,
    get_body,
    get_body_barycentric,
    get_body_barycentric_posvel,
    get_moon,
    solar_system_ephemeris,
)
from astropy.tests.helper import assert_quantity_allclose
from astropy.time import Time
from astropy.units import allclose as quantity_allclose
from astropy.utils.compat.optional_deps import HAS_JPLEPHEM, HAS_SKYFIELD
from astropy.utils.data import download_file, get_pkg_data_filename

if HAS_SKYFIELD:
    from skyfield.api import Loader, Topos

de432s_separation_tolerance_planets = 5 * u.arcsec
de432s_separation_tolerance_moon = 5 * u.arcsec
de432s_distance_tolerance = 20 * u.km

skyfield_angular_separation_tolerance = 1 * u.arcsec
skyfield_separation_tolerance = 10 * u.km


@pytest.mark.remote_data
@pytest.mark.skipif(not HAS_SKYFIELD, reason="requires skyfield")
def test_positions_skyfield(tmp_path):
    """
    Test positions against those generated by skyfield.
    """
    load = Loader(tmp_path)
    t = Time("1980-03-25 00:00")
    location = None

    # skyfield ephemeris
    try:
        planets = load("de421.bsp")
        ts = load.timescale()
    except OSError as e:
        if os.environ.get("CI", False) and "timed out" in str(e):
            pytest.xfail("Timed out in CI")
        else:
            raise

    mercury, jupiter, moon = (
        planets["mercury"],
        planets["jupiter barycenter"],
        planets["moon"],
    )
    earth = planets["earth"]

    skyfield_t = ts.from_astropy(t)

    if location is not None:
        earth = earth + Topos(
            latitude_degrees=location.lat.to_value(u.deg),
            longitude_degrees=location.lon.to_value(u.deg),
            elevation_m=location.height.to_value(u.m),
        )

    skyfield_mercury = earth.at(skyfield_t).observe(mercury).apparent()
    skyfield_jupiter = earth.at(skyfield_t).observe(jupiter).apparent()
    skyfield_moon = earth.at(skyfield_t).observe(moon).apparent()

    if location is not None:
        frame = TETE(obstime=t, location=location)
    else:
        frame = TETE(obstime=t)

    ra, dec, dist = skyfield_mercury.radec(epoch="date")
    skyfield_mercury = SkyCoord(
        ra.to(u.deg), dec.to(u.deg), distance=dist.to(u.km), frame=frame
    )
    ra, dec, dist = skyfield_jupiter.radec(epoch="date")
    skyfield_jupiter = SkyCoord(
        ra.to(u.deg), dec.to(u.deg), distance=dist.to(u.km), frame=frame
    )
    ra, dec, dist = skyfield_moon.radec(epoch="date")
    skyfield_moon = SkyCoord(
        ra.to(u.deg), dec.to(u.deg), distance=dist.to(u.km), frame=frame
    )

    # planet positions w.r.t true equator and equinox
    moon_astropy = get_moon(t, location, ephemeris="de430").transform_to(frame)
    mercury_astropy = get_body("mercury", t, location, ephemeris="de430").transform_to(
        frame
    )
    jupiter_astropy = get_body("jupiter", t, location, ephemeris="de430").transform_to(
        frame
    )

    assert (
        moon_astropy.separation(skyfield_moon) < skyfield_angular_separation_tolerance
    )
    assert moon_astropy.separation_3d(skyfield_moon) < skyfield_separation_tolerance

    assert (
        jupiter_astropy.separation(skyfield_jupiter)
        < skyfield_angular_separation_tolerance
    )
    assert (
        jupiter_astropy.separation_3d(skyfield_jupiter) < skyfield_separation_tolerance
    )

    assert (
        mercury_astropy.separation(skyfield_mercury)
        < skyfield_angular_separation_tolerance
    )
    assert (
        mercury_astropy.separation_3d(skyfield_mercury) < skyfield_separation_tolerance
    )
    planets.close()


class TestPositionsGeocentric:
    """
    Test positions against those generated by JPL Horizons accessed on
    2016-03-28, with refraction turned on.
    """

    def setup_method(self):
        self.t = Time("1980-03-25 00:00")
        self.apparent_frame = TETE(obstime=self.t)
        # Results returned by JPL Horizons web interface
        self.horizons = {
            "mercury": SkyCoord(
                ra="22h41m47.78s",
                dec="-08d29m32.0s",
                distance=c * 6.323037 * u.min,
                frame=self.apparent_frame,
            ),
            "moon": SkyCoord(
                ra="07h32m02.62s",
                dec="+18d34m05.0s",
                distance=c * 0.021921 * u.min,
                frame=self.apparent_frame,
            ),
            "jupiter": SkyCoord(
                ra="10h17m12.82s",
                dec="+12d02m57.0s",
                distance=c * 37.694557 * u.min,
                frame=self.apparent_frame,
            ),
            "sun": SkyCoord(
                ra="00h16m31.00s",
                dec="+01d47m16.9s",
                distance=c * 8.294858 * u.min,
                frame=self.apparent_frame,
            ),
        }

    @pytest.mark.parametrize(
        ("body", "sep_tol", "dist_tol"),
        (
            ("mercury", 7.0 * u.arcsec, 1000 * u.km),
            ("jupiter", 78.0 * u.arcsec, 76000 * u.km),
            ("moon", 20.0 * u.arcsec, 80 * u.km),
            ("sun", 5.0 * u.arcsec, 11.0 * u.km),
        ),
    )
    def test_erfa_planet(self, body, sep_tol, dist_tol):
        """Test predictions using erfa/plan94.

        Accuracies are maximum deviations listed in erfa/plan94.c, for Jupiter and
        Mercury, and that quoted in Meeus "Astronomical Algorithms" (1998) for the Moon.
        """
        astropy = get_body(body, self.t, ephemeris="builtin")
        horizons = self.horizons[body]

        # convert to true equator and equinox
        astropy = astropy.transform_to(self.apparent_frame)

        # Assert sky coordinates are close.
        assert astropy.separation(horizons) < sep_tol

        # Assert distances are close.
        assert_quantity_allclose(astropy.distance, horizons.distance, atol=dist_tol)

    @pytest.mark.remote_data
    @pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
    @pytest.mark.parametrize("body", ("mercury", "jupiter", "sun"))
    def test_de432s_planet(self, body):
        astropy = get_body(body, self.t, ephemeris="de432s")
        horizons = self.horizons[body]

        # convert to true equator and equinox
        astropy = astropy.transform_to(self.apparent_frame)

        # Assert sky coordinates are close.
        assert astropy.separation(horizons) < de432s_separation_tolerance_planets

        # Assert distances are close.
        assert_quantity_allclose(
            astropy.distance, horizons.distance, atol=de432s_distance_tolerance
        )

    @pytest.mark.remote_data
    @pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
    def test_de432s_moon(self):
        astropy = get_moon(self.t, ephemeris="de432s")
        horizons = self.horizons["moon"]

        # convert to true equator and equinox
        astropy = astropy.transform_to(self.apparent_frame)

        # Assert sky coordinates are close.
        assert astropy.separation(horizons) < de432s_separation_tolerance_moon

        # Assert distances are close.
        assert_quantity_allclose(
            astropy.distance, horizons.distance, atol=de432s_distance_tolerance
        )


class TestPositionKittPeak:
    """
    Test positions against those generated by JPL Horizons accessed on
    2016-03-28, with refraction turned on.
    """

    def setup_method(self):
        kitt_peak = EarthLocation.from_geodetic(
            lon=-111.6 * u.deg, lat=31.963333333333342 * u.deg, height=2120 * u.m
        )
        self.t = Time("2014-09-25T00:00", location=kitt_peak)
        self.apparent_frame = TETE(obstime=self.t, location=kitt_peak)
        # Results returned by JPL Horizons web interface
        self.horizons = {
            "mercury": SkyCoord(
                ra="13h38m58.50s",
                dec="-13d34m42.6s",
                distance=c * 7.699020 * u.min,
                frame=self.apparent_frame,
            ),
            "moon": SkyCoord(
                ra="12h33m12.85s",
                dec="-05d17m54.4s",
                distance=c * 0.022054 * u.min,
                frame=self.apparent_frame,
            ),
            "jupiter": SkyCoord(
                ra="09h09m55.55s",
                dec="+16d51m57.8s",
                distance=c * 49.244937 * u.min,
                frame=self.apparent_frame,
            ),
        }

    @pytest.mark.parametrize(
        ("body", "sep_tol", "dist_tol"),
        (
            ("mercury", 7.0 * u.arcsec, 500 * u.km),
            ("jupiter", 78.0 * u.arcsec, 82000 * u.km),
        ),
    )
    def test_erfa_planet(self, body, sep_tol, dist_tol):
        """Test predictions using erfa/plan94.

        Accuracies are maximum deviations listed in erfa/plan94.c.
        """
        # Add uncertainty in position of Earth
        dist_tol = dist_tol + 1300 * u.km

        astropy = get_body(body, self.t, ephemeris="builtin")
        horizons = self.horizons[body]

        # convert to true equator and equinox
        astropy = astropy.transform_to(self.apparent_frame)

        # Assert sky coordinates are close.
        assert astropy.separation(horizons) < sep_tol

        # Assert distances are close.
        assert_quantity_allclose(astropy.distance, horizons.distance, atol=dist_tol)

    @pytest.mark.remote_data
    @pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
    @pytest.mark.parametrize("body", ("mercury", "jupiter"))
    def test_de432s_planet(self, body):
        astropy = get_body(body, self.t, ephemeris="de432s")
        horizons = self.horizons[body]

        # convert to true equator and equinox
        astropy = astropy.transform_to(self.apparent_frame)

        # Assert sky coordinates are close.
        assert astropy.separation(horizons) < de432s_separation_tolerance_planets

        # Assert distances are close.
        assert_quantity_allclose(
            astropy.distance, horizons.distance, atol=de432s_distance_tolerance
        )

    @pytest.mark.remote_data
    @pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
    def test_de432s_moon(self):
        astropy = get_moon(self.t, ephemeris="de432s")
        horizons = self.horizons["moon"]

        # convert to true equator and equinox
        astropy = astropy.transform_to(self.apparent_frame)

        # Assert sky coordinates are close.
        assert astropy.separation(horizons) < de432s_separation_tolerance_moon

        # Assert distances are close.
        assert_quantity_allclose(
            astropy.distance, horizons.distance, atol=de432s_distance_tolerance
        )

    @pytest.mark.remote_data
    @pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
    @pytest.mark.parametrize("bodyname", ("mercury", "jupiter"))
    def test_custom_kernel_spec_body(self, bodyname):
        """
        Checks that giving a kernel specifier instead of a body name works
        """
        coord_by_name = get_body(bodyname, self.t, ephemeris="de432s")
        kspec = BODY_NAME_TO_KERNEL_SPEC[bodyname]
        coord_by_kspec = get_body(kspec, self.t, ephemeris="de432s")

        assert_quantity_allclose(coord_by_name.ra, coord_by_kspec.ra)
        assert_quantity_allclose(coord_by_name.dec, coord_by_kspec.dec)
        assert_quantity_allclose(coord_by_name.distance, coord_by_kspec.distance)


@pytest.mark.remote_data
@pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
def test_horizons_consistency_with_precision():
    """
    A test to compare at high precision against output of JPL horizons.

    Tests ephemerides, and conversions from ICRS to GCRS to TETE. We are aiming for
    better than 2 milli-arcsecond precision.

    We use the Moon since it is nearby, and moves fast in the sky so we are
    testing for parallax, proper handling of light deflection and aberration.
    """
    # JPL Horizon values for 2020_04_06 00:00 to 23:00 in 1 hour steps
    # JPL Horizons has a known offset (frame bias) of 51.02 mas in RA. We correct that here
    ra_apparent_horizons = [
        170.167332531,
        170.560688674,
        170.923834838,
        171.271663481,
        171.620188972,
        171.985340827,
        172.381766539,
        172.821772139,
        173.314502650,
        173.865422398,
        174.476108551,
        175.144332386,
        175.864375310,
        176.627519827,
        177.422655853,
        178.236955730,
        179.056584831,
        179.867427392,
        180.655815385,
        181.409252074,
        182.117113814,
        182.771311578,
        183.366872837,
        183.902395443,
    ] * u.deg + 51.02376467 * u.mas
    dec_apparent_horizons = [
        10.269112037,
        10.058820647,
        9.837152044,
        9.603724551,
        9.358956528,
        9.104012390,
        8.840674927,
        8.571162442,
        8.297917326,
        8.023394488,
        7.749873882,
        7.479312991,
        7.213246666,
        6.952732614,
        6.698336823,
        6.450150213,
        6.207828142,
        5.970645962,
        5.737565957,
        5.507313851,
        5.278462034,
        5.049521497,
        4.819038911,
        4.585696512,
    ] * u.deg
    with solar_system_ephemeris.set("de430"):
        loc = EarthLocation.from_geodetic(
            -67.787260 * u.deg, -22.959748 * u.deg, 5186 * u.m
        )
        times = Time("2020-04-06 00:00") + np.arange(0, 24, 1) * u.hour
        astropy = get_body("moon", times, loc)

        apparent_frame = TETE(obstime=times, location=loc)
        astropy = astropy.transform_to(apparent_frame)
        usrepr = UnitSphericalRepresentation(
            ra_apparent_horizons, dec_apparent_horizons
        )
        horizons = apparent_frame.realize_frame(usrepr)
    assert_quantity_allclose(astropy.separation(horizons), 0 * u.mas, atol=1.5 * u.mas)


@pytest.mark.remote_data
@pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
@pytest.mark.parametrize(
    "time",
    (Time("1960-01-12 00:00"), Time("1980-03-25 00:00"), Time("2010-10-13 00:00")),
)
def test_get_sun_consistency(time):
    """
    Test that the sun from JPL and the builtin get_sun match
    """
    sun_jpl_gcrs = get_body("sun", time, ephemeris="de432s")
    builtin_get_sun = get_sun(time)
    sep = builtin_get_sun.separation(sun_jpl_gcrs)
    assert sep < 0.1 * u.arcsec


def test_get_moon_nonscalar_regression():
    """
    Test that the builtin ephemeris works with non-scalar times.

    See Issue #5069.
    """
    times = Time(["2015-08-28 03:30", "2015-09-05 10:30"])
    # the following line will raise an Exception if the bug recurs.
    get_moon(times, ephemeris="builtin")


def test_barycentric_pos_posvel_same():
    # Check that the two routines give identical results.
    ep1 = get_body_barycentric("earth", Time("2016-03-20T12:30:00"))
    ep2, _ = get_body_barycentric_posvel("earth", Time("2016-03-20T12:30:00"))
    assert np.all(ep1.xyz == ep2.xyz)


def test_earth_barycentric_velocity_rough():
    # Check that a time near the equinox gives roughly the right result.
    ep, ev = get_body_barycentric_posvel("earth", Time("2016-03-20T12:30:00"))
    assert_quantity_allclose(ep.xyz, [-1.0, 0.0, 0.0] * u.AU, atol=0.01 * u.AU)
    expected = (
        u.Quantity([0.0 * u.one, np.cos(23.5 * u.deg), np.sin(23.5 * u.deg)])
        * -30.0
        * u.km
        / u.s
    )
    assert_quantity_allclose(ev.xyz, expected, atol=1.0 * u.km / u.s)


def test_earth_barycentric_velocity_multi_d():
    # Might as well test it with a multidimensional array too.
    t = Time("2016-03-20T12:30:00") + np.arange(8.0).reshape(2, 2, 2) * u.yr / 2.0
    ep, ev = get_body_barycentric_posvel("earth", t)
    # note: assert_quantity_allclose doesn't like the shape mismatch.
    # this is a problem with np.testing.assert_allclose.
    assert quantity_allclose(
        ep.get_xyz(xyz_axis=-1),
        [[-1.0, 0.0, 0.0], [+1.0, 0.0, 0.0]] * u.AU,
        atol=0.06 * u.AU,
    )
    expected = u.Quantity([0.0 * u.one, np.cos(23.5 * u.deg), np.sin(23.5 * u.deg)]) * (
        [[-30.0], [30.0]] * u.km / u.s
    )
    assert quantity_allclose(ev.get_xyz(xyz_axis=-1), expected, atol=2.0 * u.km / u.s)


@pytest.mark.remote_data
@pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
@pytest.mark.parametrize(
    ("body", "pos_tol", "vel_tol"),
    (
        ("mercury", 1000.0 * u.km, 1.0 * u.km / u.s),
        ("jupiter", 100000.0 * u.km, 2.0 * u.km / u.s),
        ("earth", 10 * u.km, 10 * u.mm / u.s),
        ("moon", 18 * u.km, 50 * u.mm / u.s),
    ),
)
def test_barycentric_velocity_consistency(body, pos_tol, vel_tol):
    # Tolerances are about 1.5 times the rms listed for plan94 and epv00,
    # except for Mercury (which nominally is 334 km rms), and the Moon
    # (which nominally is 6 km rms).
    t = Time("2016-03-20T12:30:00")
    ep, ev = get_body_barycentric_posvel(body, t, ephemeris="builtin")
    dp, dv = get_body_barycentric_posvel(body, t, ephemeris="de432s")
    assert_quantity_allclose(ep.xyz, dp.xyz, atol=pos_tol)
    assert_quantity_allclose(ev.xyz, dv.xyz, atol=vel_tol)
    # Might as well test it with a multidimensional array too.
    t = Time("2016-03-20T12:30:00") + np.arange(8.0).reshape(2, 2, 2) * u.yr / 2.0
    ep, ev = get_body_barycentric_posvel(body, t, ephemeris="builtin")
    dp, dv = get_body_barycentric_posvel(body, t, ephemeris="de432s")
    assert_quantity_allclose(ep.xyz, dp.xyz, atol=pos_tol)
    assert_quantity_allclose(ev.xyz, dv.xyz, atol=vel_tol)


@pytest.mark.remote_data
@pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
@pytest.mark.parametrize(
    "time",
    (Time("1960-01-12 00:00"), Time("1980-03-25 00:00"), Time("2010-10-13 00:00")),
)
def test_url_or_file_ephemeris(time):
    # URL for ephemeris de432s used for testing:
    url = "http://naif.jpl.nasa.gov/pub/naif/generic_kernels/spk/planets/de432s.bsp"

    # Pass the ephemeris directly as a URL.
    coord_by_url = get_body("earth", time, ephemeris=url)

    # Translate the URL to the cached location on the filesystem.
    # Since we just used the url above, it should already have been downloaded.
    filepath = download_file(url, cache=True)

    # Get the coordinates using the file path directly:
    coord_by_filepath = get_body("earth", time, ephemeris=filepath)

    # Using the URL or filepath should give exactly the same results:
    assert_quantity_allclose(coord_by_url.ra, coord_by_filepath.ra)
    assert_quantity_allclose(coord_by_url.dec, coord_by_filepath.dec)
    assert_quantity_allclose(coord_by_url.distance, coord_by_filepath.distance)


@pytest.mark.remote_data
@pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
def test_url_ephemeris_wrong_input():
    time = Time("1960-01-12 00:00")
    with pytest.raises((HTTPError, URLError)):
        # A non-existent URL
        get_body(
            "earth",
            time,
            ephemeris=get_pkg_data_filename("path/to/nonexisting/file.bsp"),
        )

    with pytest.raises(HTTPError):
        # A non-existent version of the JPL ephemeris
        get_body("earth", time, ephemeris="de001")

    with pytest.raises(ValueError):
        # An invalid string
        get_body("earth", time, ephemeris="not_an_ephemeris")


@pytest.mark.skipif(not HAS_JPLEPHEM, reason="requires jplephem")
def test_file_ephemeris_wrong_input():
    time = Time("1960-01-12 00:00")
    # Try loading a non-existing file:
    with pytest.raises(ValueError):
        get_body("earth", time, ephemeris="/path/to/nonexisting/file.bsp")

    # NOTE: This test currently leaves the file open (ResourceWarning).
    # To fix this issue, an upstream fix is required in jplephem
    # package.
    # Try loading a file that does exist, but is not an ephemeris file:
    with pytest.warns(ResourceWarning), pytest.raises(ValueError):
        get_body("earth", time, ephemeris=__file__)


def test_regression_10271():
    t = Time(58973.534052125986, format="mjd")
    # GCRS position of ALMA at this time
    obs_p = CartesianRepresentation(
        5724535.74068625, -1311071.58985697, -2492738.93017009, u.m
    )

    geocentre = CartesianRepresentation(0, 0, 0, u.m)
    icrs_sun_from_alma = _get_apparent_body_position("sun", t, "builtin", obs_p)
    icrs_sun_from_geocentre = _get_apparent_body_position(
        "sun", t, "builtin", geocentre
    )

    difference = (icrs_sun_from_alma - icrs_sun_from_geocentre).norm()
    assert_quantity_allclose(difference, 0.13046941 * u.m, atol=1 * u.mm)