# Licensed under a 3-clause BSD style license - see LICENSE.rst

import operator

import numpy as np
import pytest

from astropy import units as u
from astropy.coordinates import (
    CartesianDifferential,
    CartesianRepresentation,
    CylindricalDifferential,
    CylindricalRepresentation,
    Latitude,
    Longitude,
    PhysicsSphericalDifferential,
    PhysicsSphericalRepresentation,
    RadialDifferential,
    RadialRepresentation,
    SphericalCosLatDifferential,
    SphericalDifferential,
    SphericalRepresentation,
    UnitSphericalCosLatDifferential,
    UnitSphericalDifferential,
    UnitSphericalRepresentation,
    angular_separation,
)
from astropy.coordinates.representation import DIFFERENTIAL_CLASSES
from astropy.tests.helper import assert_quantity_allclose, quantity_allclose


def assert_representation_allclose(actual, desired, rtol=1.0e-7, atol=None, **kwargs):
    __tracebackhide__ = True
    actual_xyz = actual.to_cartesian().get_xyz(xyz_axis=-1)
    desired_xyz = desired.to_cartesian().get_xyz(xyz_axis=-1)
    actual_xyz, desired_xyz = np.broadcast_arrays(actual_xyz, desired_xyz, subok=True)
    assert_quantity_allclose(actual_xyz, desired_xyz, rtol, atol, **kwargs)


def assert_differential_allclose(actual, desired, rtol=1.0e-7, **kwargs):
    assert actual.components == desired.components
    for component in actual.components:
        actual_c = getattr(actual, component)
        atol = 1.0e-10 * actual_c.unit
        assert_quantity_allclose(
            actual_c, getattr(desired, component), rtol, atol, **kwargs
        )


def representation_equal(first, second):
    return np.all(
        getattr(first, comp) == getattr(second, comp) for comp in first.components
    )


class TestArithmetic:
    def setup_method(self):
        # Choose some specific coordinates, for which ``sum`` and ``dot``
        # works out nicely.
        self.lon = Longitude(np.arange(0, 12.1, 2), u.hourangle)
        self.lat = Latitude(np.arange(-90, 91, 30), u.deg)
        self.distance = [5.0, 12.0, 4.0, 2.0, 4.0, 12.0, 5.0] * u.kpc
        self.spherical = SphericalRepresentation(self.lon, self.lat, self.distance)
        self.unit_spherical = self.spherical.represent_as(UnitSphericalRepresentation)
        self.cartesian = self.spherical.to_cartesian()

    def test_norm_spherical(self):
        norm_s = self.spherical.norm()
        assert isinstance(norm_s, u.Quantity)
        # Just to be sure, test against getting object arrays.
        assert norm_s.dtype.kind == "f"
        assert np.all(norm_s == self.distance)

    @pytest.mark.parametrize(
        "representation",
        (
            PhysicsSphericalRepresentation,
            CartesianRepresentation,
            CylindricalRepresentation,
        ),
    )
    def test_norm(self, representation):
        in_rep = self.spherical.represent_as(representation)
        norm_rep = in_rep.norm()
        assert isinstance(norm_rep, u.Quantity)
        assert_quantity_allclose(norm_rep, self.distance)

    def test_norm_unitspherical(self):
        norm_rep = self.unit_spherical.norm()
        assert norm_rep.unit == u.dimensionless_unscaled
        assert np.all(norm_rep == 1.0 * u.dimensionless_unscaled)

    @pytest.mark.parametrize(
        "representation",
        (
            SphericalRepresentation,
            PhysicsSphericalRepresentation,
            CartesianRepresentation,
            CylindricalRepresentation,
            UnitSphericalRepresentation,
        ),
    )
    def test_neg_pos(self, representation):
        in_rep = self.cartesian.represent_as(representation)
        pos_rep = +in_rep
        assert type(pos_rep) is type(in_rep)
        assert pos_rep is not in_rep
        assert np.all(representation_equal(pos_rep, in_rep))
        neg_rep = -in_rep
        assert type(neg_rep) is type(in_rep)
        assert np.all(neg_rep.norm() == in_rep.norm())
        in_rep_xyz = in_rep.to_cartesian().xyz
        assert_quantity_allclose(
            neg_rep.to_cartesian().xyz, -in_rep_xyz, atol=1.0e-10 * in_rep_xyz.unit
        )

    def test_mul_div_spherical(self):
        s0 = self.spherical / (1.0 * u.Myr)
        assert isinstance(s0, SphericalRepresentation)
        assert s0.distance.dtype.kind == "f"
        assert np.all(s0.lon == self.spherical.lon)
        assert np.all(s0.lat == self.spherical.lat)
        assert np.all(s0.distance == self.distance / (1.0 * u.Myr))
        s1 = (1.0 / u.Myr) * self.spherical
        assert isinstance(s1, SphericalRepresentation)
        assert np.all(representation_equal(s1, s0))
        s2 = self.spherical * np.array([[1.0], [2.0]])
        assert isinstance(s2, SphericalRepresentation)
        assert s2.shape == (2, self.spherical.shape[0])
        assert np.all(s2.lon == self.spherical.lon)
        assert np.all(s2.lat == self.spherical.lat)
        assert np.all(s2.distance == self.spherical.distance * np.array([[1.0], [2.0]]))
        s3 = np.array([[1.0], [2.0]]) * self.spherical
        assert isinstance(s3, SphericalRepresentation)
        assert np.all(representation_equal(s3, s2))
        s4 = -self.spherical
        assert isinstance(s4, SphericalRepresentation)
        assert quantity_allclose(
            s4.to_cartesian().xyz,
            -self.spherical.to_cartesian().xyz,
            atol=1e-15 * self.spherical.distance.unit,
        )
        assert np.all(s4.distance == self.spherical.distance)
        s5 = +self.spherical
        assert s5 is not self.spherical
        assert np.all(representation_equal(s5, self.spherical))

    @pytest.mark.parametrize(
        "representation",
        (
            PhysicsSphericalRepresentation,
            CartesianRepresentation,
            CylindricalRepresentation,
        ),
    )
    def test_mul_div(self, representation):
        in_rep = self.spherical.represent_as(representation)
        r1 = in_rep / (1.0 * u.Myr)
        assert isinstance(r1, representation)
        for component in in_rep.components:
            in_rep_comp = getattr(in_rep, component)
            r1_comp = getattr(r1, component)
            if in_rep_comp.unit == self.distance.unit:
                assert np.all(r1_comp == in_rep_comp / (1.0 * u.Myr))
            else:
                assert np.all(r1_comp == in_rep_comp)

        r2 = np.array([[1.0], [2.0]]) * in_rep
        assert isinstance(r2, representation)
        assert r2.shape == (2, in_rep.shape[0])
        assert_quantity_allclose(r2.norm(), self.distance * np.array([[1.0], [2.0]]))
        r3 = -in_rep
        assert_representation_allclose(
            r3.to_cartesian(), (in_rep * -1.0).to_cartesian(), atol=1e-5 * u.pc
        )
        with pytest.raises(TypeError):
            in_rep * in_rep
        with pytest.raises(TypeError):
            {} * in_rep

    def test_mul_div_unit_spherical(self):
        s1 = self.unit_spherical * self.distance
        assert isinstance(s1, SphericalRepresentation)
        assert np.all(s1.lon == self.unit_spherical.lon)
        assert np.all(s1.lat == self.unit_spherical.lat)
        assert np.all(s1.distance == self.spherical.distance)
        s2 = self.unit_spherical / u.s
        assert isinstance(s2, SphericalRepresentation)
        assert np.all(s2.lon == self.unit_spherical.lon)
        assert np.all(s2.lat == self.unit_spherical.lat)
        assert np.all(s2.distance == 1.0 / u.s)
        u3 = -self.unit_spherical
        assert isinstance(u3, UnitSphericalRepresentation)
        assert_quantity_allclose(u3.lon, self.unit_spherical.lon + 180.0 * u.deg)
        assert np.all(u3.lat == -self.unit_spherical.lat)
        assert_quantity_allclose(
            u3.to_cartesian().xyz,
            -self.unit_spherical.to_cartesian().xyz,
            atol=1.0e-10 * u.dimensionless_unscaled,
        )
        u4 = +self.unit_spherical
        assert isinstance(u4, UnitSphericalRepresentation)
        assert u4 is not self.unit_spherical
        assert np.all(representation_equal(u4, self.unit_spherical))

    def test_add_sub_cartesian(self):
        c1 = self.cartesian + self.cartesian
        assert isinstance(c1, CartesianRepresentation)
        assert c1.x.dtype.kind == "f"
        assert np.all(representation_equal(c1, 2.0 * self.cartesian))
        with pytest.raises(TypeError):
            self.cartesian + 10.0 * u.m
        with pytest.raises(u.UnitsError):
            self.cartesian + (self.cartesian / u.s)
        c2 = self.cartesian - self.cartesian
        assert isinstance(c2, CartesianRepresentation)
        assert np.all(
            representation_equal(
                c2, CartesianRepresentation(0.0 * u.m, 0.0 * u.m, 0.0 * u.m)
            )
        )
        c3 = self.cartesian - self.cartesian / 2.0
        assert isinstance(c3, CartesianRepresentation)
        assert np.all(representation_equal(c3, self.cartesian / 2.0))

    @pytest.mark.parametrize(
        "representation",
        (
            PhysicsSphericalRepresentation,
            SphericalRepresentation,
            CylindricalRepresentation,
        ),
    )
    def test_add_sub(self, representation):
        in_rep = self.cartesian.represent_as(representation)
        r1 = in_rep + in_rep
        assert isinstance(r1, representation)
        expected = 2.0 * in_rep
        for component in in_rep.components:
            assert_quantity_allclose(
                getattr(r1, component), getattr(expected, component)
            )
        with pytest.raises(TypeError):
            10.0 * u.m + in_rep
        with pytest.raises(u.UnitsError):
            in_rep + (in_rep / u.s)
        r2 = in_rep - in_rep
        assert isinstance(r2, representation)
        assert_representation_allclose(
            r2.to_cartesian(),
            CartesianRepresentation(0.0 * u.m, 0.0 * u.m, 0.0 * u.m),
            atol=1e-15 * u.kpc,
        )
        r3 = in_rep - in_rep / 2.0
        assert isinstance(r3, representation)
        expected = in_rep / 2.0
        assert_representation_allclose(r3, expected)

    def test_add_sub_unit_spherical(self):
        s1 = self.unit_spherical + self.unit_spherical
        assert isinstance(s1, SphericalRepresentation)
        expected = 2.0 * self.unit_spherical
        for component in s1.components:
            assert_quantity_allclose(
                getattr(s1, component), getattr(expected, component)
            )
        with pytest.raises(TypeError):
            10.0 * u.m - self.unit_spherical
        with pytest.raises(u.UnitsError):
            self.unit_spherical + (self.unit_spherical / u.s)
        s2 = self.unit_spherical - self.unit_spherical / 2.0
        assert isinstance(s2, SphericalRepresentation)
        expected = self.unit_spherical / 2.0
        for component in s2.components:
            assert_quantity_allclose(
                getattr(s2, component), getattr(expected, component)
            )

    @pytest.mark.parametrize(
        "representation",
        (
            CartesianRepresentation,
            PhysicsSphericalRepresentation,
            SphericalRepresentation,
            CylindricalRepresentation,
        ),
    )
    def test_sum_mean(self, representation):
        in_rep = self.spherical.represent_as(representation)
        r_sum = in_rep.sum()
        assert isinstance(r_sum, representation)
        expected = SphericalRepresentation(
            90.0 * u.deg, 0.0 * u.deg, 14.0 * u.kpc
        ).represent_as(representation)
        for component in expected.components:
            exp_component = getattr(expected, component)
            assert_quantity_allclose(
                getattr(r_sum, component),
                exp_component,
                atol=1e-10 * exp_component.unit,
            )

        r_mean = in_rep.mean()
        assert isinstance(r_mean, representation)
        expected = expected / len(in_rep)
        for component in expected.components:
            exp_component = getattr(expected, component)
            assert_quantity_allclose(
                getattr(r_mean, component),
                exp_component,
                atol=1e-10 * exp_component.unit,
            )

    def test_sum_mean_unit_spherical(self):
        s_sum = self.unit_spherical.sum()
        assert isinstance(s_sum, SphericalRepresentation)
        expected = SphericalRepresentation(
            90.0 * u.deg, 0.0 * u.deg, 3.0 * u.dimensionless_unscaled
        )
        for component in expected.components:
            exp_component = getattr(expected, component)
            assert_quantity_allclose(
                getattr(s_sum, component),
                exp_component,
                atol=1e-10 * exp_component.unit,
            )

        s_mean = self.unit_spherical.mean()
        assert isinstance(s_mean, SphericalRepresentation)
        expected = expected / len(self.unit_spherical)
        for component in expected.components:
            exp_component = getattr(expected, component)
            assert_quantity_allclose(
                getattr(s_mean, component),
                exp_component,
                atol=1e-10 * exp_component.unit,
            )

    @pytest.mark.parametrize(
        "representation",
        (
            CartesianRepresentation,
            PhysicsSphericalRepresentation,
            SphericalRepresentation,
            CylindricalRepresentation,
        ),
    )
    def test_dot(self, representation):
        in_rep = self.cartesian.represent_as(representation)
        r_dot_r = in_rep.dot(in_rep)
        assert isinstance(r_dot_r, u.Quantity)
        assert r_dot_r.shape == in_rep.shape
        assert_quantity_allclose(np.sqrt(r_dot_r), self.distance)
        r_dot_r_rev = in_rep.dot(in_rep[::-1])
        assert isinstance(r_dot_r_rev, u.Quantity)
        assert r_dot_r_rev.shape == in_rep.shape
        expected = [-25.0, -126.0, 2.0, 4.0, 2.0, -126.0, -25.0] * u.kpc**2
        assert_quantity_allclose(r_dot_r_rev, expected)
        for axis in "xyz":
            project = CartesianRepresentation(
                *(
                    (1.0 if axis == _axis else 0.0) * u.dimensionless_unscaled
                    for _axis in "xyz"
                )
            )
            assert_quantity_allclose(
                in_rep.dot(project), getattr(self.cartesian, axis), atol=1.0 * u.upc
            )
        with pytest.raises(TypeError):
            in_rep.dot(self.cartesian.xyz)

    def test_dot_unit_spherical(self):
        u_dot_u = self.unit_spherical.dot(self.unit_spherical)
        assert isinstance(u_dot_u, u.Quantity)
        assert u_dot_u.shape == self.unit_spherical.shape
        assert_quantity_allclose(u_dot_u, 1.0 * u.dimensionless_unscaled)
        cartesian = self.unit_spherical.to_cartesian()
        for axis in "xyz":
            project = CartesianRepresentation(
                *(
                    (1.0 if axis == _axis else 0.0) * u.dimensionless_unscaled
                    for _axis in "xyz"
                )
            )
            assert_quantity_allclose(
                self.unit_spherical.dot(project), getattr(cartesian, axis), atol=1.0e-10
            )

    @pytest.mark.parametrize(
        "representation",
        (
            CartesianRepresentation,
            PhysicsSphericalRepresentation,
            SphericalRepresentation,
            CylindricalRepresentation,
        ),
    )
    def test_cross(self, representation):
        in_rep = self.cartesian.represent_as(representation)
        r_cross_r = in_rep.cross(in_rep)
        assert isinstance(r_cross_r, representation)
        assert_quantity_allclose(r_cross_r.norm(), 0.0 * u.kpc**2, atol=1.0 * u.mpc**2)
        r_cross_r_rev = in_rep.cross(in_rep[::-1])
        sep = angular_separation(self.lon, self.lat, self.lon[::-1], self.lat[::-1])
        expected = self.distance * self.distance[::-1] * np.sin(sep)
        assert_quantity_allclose(r_cross_r_rev.norm(), expected, atol=1.0 * u.mpc**2)
        unit_vectors = CartesianRepresentation(
            [1.0, 0.0, 0.0] * u.one, [0.0, 1.0, 0.0] * u.one, [0.0, 0.0, 1.0] * u.one
        )[:, np.newaxis]
        r_cross_uv = in_rep.cross(unit_vectors)
        assert r_cross_uv.shape == (3, 7)
        assert_quantity_allclose(
            r_cross_uv.dot(unit_vectors), 0.0 * u.kpc, atol=1.0 * u.upc
        )
        assert_quantity_allclose(
            r_cross_uv.dot(in_rep), 0.0 * u.kpc**2, atol=1.0 * u.mpc**2
        )
        zeros = np.zeros(len(in_rep)) * u.kpc
        expected = CartesianRepresentation(
            u.Quantity((zeros, -self.cartesian.z, self.cartesian.y)),
            u.Quantity((self.cartesian.z, zeros, -self.cartesian.x)),
            u.Quantity((-self.cartesian.y, self.cartesian.x, zeros)),
        )
        # Comparison with spherical is hard since some distances are zero,
        # implying the angles are undefined.
        r_cross_uv_cartesian = r_cross_uv.to_cartesian()
        assert_representation_allclose(r_cross_uv_cartesian, expected, atol=1.0 * u.upc)
        # A final check, with the side benefit of ensuring __truediv__ and norm
        # work on multi-D representations.
        r_cross_uv_by_distance = r_cross_uv / self.distance
        uv_sph = unit_vectors.represent_as(UnitSphericalRepresentation)
        sep = angular_separation(self.lon, self.lat, uv_sph.lon, uv_sph.lat)
        assert_quantity_allclose(r_cross_uv_by_distance.norm(), np.sin(sep), atol=1e-9)

        with pytest.raises(TypeError):
            in_rep.cross(self.cartesian.xyz)

    def test_cross_unit_spherical(self):
        u_cross_u = self.unit_spherical.cross(self.unit_spherical)
        assert isinstance(u_cross_u, SphericalRepresentation)
        assert_quantity_allclose(u_cross_u.norm(), 0.0 * u.one, atol=1.0e-10 * u.one)
        u_cross_u_rev = self.unit_spherical.cross(self.unit_spherical[::-1])
        assert isinstance(u_cross_u_rev, SphericalRepresentation)
        sep = angular_separation(self.lon, self.lat, self.lon[::-1], self.lat[::-1])
        expected = np.sin(sep)
        assert_quantity_allclose(u_cross_u_rev.norm(), expected, atol=1.0e-10 * u.one)


class TestUnitVectorsAndScales:
    @staticmethod
    def check_unit_vectors(e):
        for v in e.values():
            assert type(v) is CartesianRepresentation
            assert_quantity_allclose(v.norm(), 1.0 * u.one)
        return e

    @staticmethod
    def check_scale_factors(sf, rep):
        unit = rep.norm().unit
        for c, f in sf.items():
            assert type(f) is u.Quantity
            assert (f.unit * getattr(rep, c).unit).is_equivalent(unit)

    def test_spherical(self):
        s = SphericalRepresentation(
            lon=[0.0, 6.0, 21.0] * u.hourangle,
            lat=[0.0, -30.0, 85.0] * u.deg,
            distance=[1, 2, 3] * u.kpc,
        )
        e = s.unit_vectors()
        self.check_unit_vectors(e)
        sf = s.scale_factors()
        self.check_scale_factors(sf, s)

        s_lon = s + s.distance * 1e-5 * np.cos(s.lat) * e["lon"]
        assert_quantity_allclose(s_lon.lon, s.lon + 1e-5 * u.rad, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_lon.lat, s.lat, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_lon.distance, s.distance)
        s_lon2 = s + 1e-5 * u.radian * sf["lon"] * e["lon"]
        assert_representation_allclose(s_lon2, s_lon)

        s_lat = s + s.distance * 1e-5 * e["lat"]
        assert_quantity_allclose(s_lat.lon, s.lon)
        assert_quantity_allclose(s_lat.lat, s.lat + 1e-5 * u.rad, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_lon.distance, s.distance)
        s_lat2 = s + 1.0e-5 * u.radian * sf["lat"] * e["lat"]
        assert_representation_allclose(s_lat2, s_lat)

        s_distance = s + 1.0 * u.pc * e["distance"]
        assert_quantity_allclose(s_distance.lon, s.lon, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_distance.lat, s.lat, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_distance.distance, s.distance + 1.0 * u.pc)
        s_distance2 = s + 1.0 * u.pc * sf["distance"] * e["distance"]
        assert_representation_allclose(s_distance2, s_distance)

    def test_unit_spherical(self):
        s = UnitSphericalRepresentation(
            lon=[0.0, 6.0, 21.0] * u.hourangle, lat=[0.0, -30.0, 85.0] * u.deg
        )

        e = s.unit_vectors()
        self.check_unit_vectors(e)
        sf = s.scale_factors()
        self.check_scale_factors(sf, s)

        s_lon = s + 1e-5 * np.cos(s.lat) * e["lon"]
        assert_quantity_allclose(s_lon.lon, s.lon + 1e-5 * u.rad, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_lon.lat, s.lat, atol=1e-10 * u.rad)
        s_lon2 = s + 1e-5 * u.radian * sf["lon"] * e["lon"]
        assert_representation_allclose(s_lon2, s_lon)

        s_lat = s + 1e-5 * e["lat"]
        assert_quantity_allclose(s_lat.lon, s.lon)
        assert_quantity_allclose(s_lat.lat, s.lat + 1e-5 * u.rad, atol=1e-10 * u.rad)
        s_lat2 = s + 1.0e-5 * u.radian * sf["lat"] * e["lat"]
        assert_representation_allclose(s_lat2, s_lat)

    def test_radial(self):
        r = RadialRepresentation(10.0 * u.kpc)
        with pytest.raises(NotImplementedError):
            r.unit_vectors()
        sf = r.scale_factors()
        assert np.all(sf["distance"] == 1.0 * u.one)
        assert np.all(r.norm() == r.distance)
        with pytest.raises(TypeError):
            r + r

    def test_physical_spherical(self):
        s = PhysicsSphericalRepresentation(
            phi=[0.0, 6.0, 21.0] * u.hourangle,
            theta=[90.0, 120.0, 5.0] * u.deg,
            r=[1, 2, 3] * u.kpc,
        )

        e = s.unit_vectors()
        self.check_unit_vectors(e)
        sf = s.scale_factors()
        self.check_scale_factors(sf, s)

        s_phi = s + s.r * 1e-5 * np.sin(s.theta) * e["phi"]
        assert_quantity_allclose(s_phi.phi, s.phi + 1e-5 * u.rad, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_phi.theta, s.theta, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_phi.r, s.r)
        s_phi2 = s + 1e-5 * u.radian * sf["phi"] * e["phi"]
        assert_representation_allclose(s_phi2, s_phi)

        s_theta = s + s.r * 1e-5 * e["theta"]
        assert_quantity_allclose(s_theta.phi, s.phi)
        assert_quantity_allclose(
            s_theta.theta, s.theta + 1e-5 * u.rad, atol=1e-10 * u.rad
        )
        assert_quantity_allclose(s_theta.r, s.r)
        s_theta2 = s + 1.0e-5 * u.radian * sf["theta"] * e["theta"]
        assert_representation_allclose(s_theta2, s_theta)

        s_r = s + 1.0 * u.pc * e["r"]
        assert_quantity_allclose(s_r.phi, s.phi, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_r.theta, s.theta, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_r.r, s.r + 1.0 * u.pc)
        s_r2 = s + 1.0 * u.pc * sf["r"] * e["r"]
        assert_representation_allclose(s_r2, s_r)

    def test_cartesian(self):
        s = CartesianRepresentation(
            x=[1, 2, 3] * u.pc, y=[2, 3, 4] * u.Mpc, z=[3, 4, 5] * u.kpc
        )

        e = s.unit_vectors()
        sf = s.scale_factors()
        for v, expected in zip(
            e.values(),
            ([1.0, 0.0, 0.0] * u.one, [0.0, 1.0, 0.0] * u.one, [0.0, 0.0, 1.0] * u.one),
        ):
            assert np.all(v.get_xyz(xyz_axis=-1) == expected)
        for f in sf.values():
            assert np.all(f == 1.0 * u.one)

    def test_cylindrical(self):
        s = CylindricalRepresentation(
            rho=[1, 2, 3] * u.pc, phi=[0.0, 90.0, -45.0] * u.deg, z=[3, 4, 5] * u.kpc
        )
        e = s.unit_vectors()
        self.check_unit_vectors(e)
        sf = s.scale_factors()
        self.check_scale_factors(sf, s)

        s_rho = s + 1.0 * u.pc * e["rho"]
        assert_quantity_allclose(s_rho.rho, s.rho + 1.0 * u.pc)
        assert_quantity_allclose(s_rho.phi, s.phi)
        assert_quantity_allclose(s_rho.z, s.z)
        s_rho2 = s + 1.0 * u.pc * sf["rho"] * e["rho"]
        assert_representation_allclose(s_rho2, s_rho)

        s_phi = s + s.rho * 1e-5 * e["phi"]
        assert_quantity_allclose(s_phi.rho, s.rho)
        assert_quantity_allclose(s_phi.phi, s.phi + 1e-5 * u.rad)
        assert_quantity_allclose(s_phi.z, s.z)
        s_phi2 = s + 1e-5 * u.radian * sf["phi"] * e["phi"]
        assert_representation_allclose(s_phi2, s_phi)

        s_z = s + 1.0 * u.pc * e["z"]
        assert_quantity_allclose(s_z.rho, s.rho)
        assert_quantity_allclose(s_z.phi, s.phi, atol=1e-10 * u.rad)
        assert_quantity_allclose(s_z.z, s.z + 1.0 * u.pc)
        s_z2 = s + 1.0 * u.pc * sf["z"] * e["z"]
        assert_representation_allclose(s_z2, s_z)


@pytest.mark.parametrize("omit_coslat", [False, True], scope="class")
class TestSphericalDifferential:
    # these test cases are subclassed for SphericalCosLatDifferential,
    # hence some tests depend on omit_coslat.

    def _setup(self, omit_coslat):
        if omit_coslat:
            self.SD_cls = SphericalCosLatDifferential
        else:
            self.SD_cls = SphericalDifferential

        s = SphericalRepresentation(
            lon=[0.0, 6.0, 21.0] * u.hourangle,
            lat=[0.0, -30.0, 85.0] * u.deg,
            distance=[1, 2, 3] * u.kpc,
        )
        self.s = s
        self.e = s.unit_vectors()
        self.sf = s.scale_factors(omit_coslat=omit_coslat)

    def test_name_coslat(self, omit_coslat):
        self._setup(omit_coslat)
        if omit_coslat:
            assert self.SD_cls is SphericalCosLatDifferential
            assert self.SD_cls.name == "sphericalcoslat"
        else:
            assert self.SD_cls is SphericalDifferential
            assert self.SD_cls.name == "spherical"
        assert self.SD_cls.name in DIFFERENTIAL_CLASSES

    def test_simple_differentials(self, omit_coslat):
        self._setup(omit_coslat)
        s, e, sf = self.s, self.e, self.sf

        o_lon = self.SD_cls(1.0 * u.arcsec, 0.0 * u.arcsec, 0.0 * u.kpc)
        o_lonc = o_lon.to_cartesian(base=s)
        o_lon2 = self.SD_cls.from_cartesian(o_lonc, base=s)
        assert_differential_allclose(o_lon, o_lon2)
        # simple check by hand for first element.
        # lat[0] is 0, so cos(lat) term doesn't matter.
        assert_quantity_allclose(
            o_lonc[0].xyz, [0.0, np.pi / 180.0 / 3600.0, 0.0] * u.kpc
        )
        # check all using unit vectors and scale factors.
        s_lon = s + 1.0 * u.arcsec * sf["lon"] * e["lon"]
        assert_representation_allclose(o_lonc, s_lon - s, atol=1 * u.npc)
        s_lon2 = s + o_lon
        assert_representation_allclose(s_lon2, s_lon, atol=1 * u.npc)

        o_lat = self.SD_cls(0.0 * u.arcsec, 1.0 * u.arcsec, 0.0 * u.kpc)
        o_latc = o_lat.to_cartesian(base=s)
        assert_quantity_allclose(
            o_latc[0].xyz, [0.0, 0.0, np.pi / 180.0 / 3600.0] * u.kpc, atol=1.0 * u.npc
        )
        s_lat = s + 1.0 * u.arcsec * sf["lat"] * e["lat"]
        assert_representation_allclose(o_latc, s_lat - s, atol=1 * u.npc)
        s_lat2 = s + o_lat
        assert_representation_allclose(s_lat2, s_lat, atol=1 * u.npc)

        o_distance = self.SD_cls(0.0 * u.arcsec, 0.0 * u.arcsec, 1.0 * u.mpc)
        o_distancec = o_distance.to_cartesian(base=s)
        assert_quantity_allclose(
            o_distancec[0].xyz, [1e-6, 0.0, 0.0] * u.kpc, atol=1.0 * u.npc
        )
        s_distance = s + 1.0 * u.mpc * sf["distance"] * e["distance"]
        assert_representation_allclose(o_distancec, s_distance - s, atol=1 * u.npc)
        s_distance2 = s + o_distance
        assert_representation_allclose(s_distance2, s_distance)

    def test_differential_arithmetic(self, omit_coslat):
        self._setup(omit_coslat)
        s = self.s

        o_lon = self.SD_cls(1.0 * u.arcsec, 0.0 * u.arcsec, 0.0 * u.kpc)
        o_lon_by_2 = o_lon / 2.0
        assert_representation_allclose(
            o_lon_by_2.to_cartesian(s) * 2.0, o_lon.to_cartesian(s), atol=1e-10 * u.kpc
        )
        assert_representation_allclose(
            s + o_lon, s + 2 * o_lon_by_2, atol=1e-10 * u.kpc
        )
        o_lon_rec = o_lon_by_2 + o_lon_by_2
        assert_representation_allclose(s + o_lon, s + o_lon_rec, atol=1e-10 * u.kpc)
        o_lon_0 = o_lon - o_lon
        for c in o_lon_0.components:
            assert np.all(getattr(o_lon_0, c) == 0.0)
        o_lon2 = self.SD_cls(1 * u.mas / u.yr, 0 * u.mas / u.yr, 0 * u.km / u.s)
        assert_quantity_allclose(
            o_lon2.norm(s)[0], 4.74 * u.km / u.s, atol=0.01 * u.km / u.s
        )
        assert_representation_allclose(
            o_lon2.to_cartesian(s) * 1000.0 * u.yr,
            o_lon.to_cartesian(s),
            atol=1e-10 * u.kpc,
        )
        s_off = s + o_lon
        s_off2 = s + o_lon2 * 1000.0 * u.yr
        assert_representation_allclose(s_off, s_off2, atol=1e-10 * u.kpc)

        factor = 1e5 * u.radian / u.arcsec
        if not omit_coslat:
            factor = factor / np.cos(s.lat)
        s_off_big = s + o_lon * factor

        assert_representation_allclose(
            s_off_big,
            SphericalRepresentation(
                s.lon + 90.0 * u.deg, 0.0 * u.deg, 1e5 * s.distance
            ),
            atol=5.0 * u.kpc,
        )

        o_lon3c = CartesianRepresentation(0.0, 4.74047, 0.0, unit=u.km / u.s)
        o_lon3 = self.SD_cls.from_cartesian(o_lon3c, base=s)
        expected0 = self.SD_cls(
            1.0 * u.mas / u.yr, 0.0 * u.mas / u.yr, 0.0 * u.km / u.s
        )
        assert_differential_allclose(o_lon3[0], expected0)
        s_off_big2 = s + o_lon3 * 1e5 * u.yr * u.radian / u.mas
        assert_representation_allclose(
            s_off_big2,
            SphericalRepresentation(90.0 * u.deg, 0.0 * u.deg, 1e5 * u.kpc),
            atol=5.0 * u.kpc,
        )

        with pytest.raises(TypeError):
            o_lon - s
        with pytest.raises(TypeError):
            s.to_cartesian() + o_lon

    def test_differential_init_errors(self, omit_coslat):
        self._setup(omit_coslat)
        s = self.s
        with pytest.raises(u.UnitsError):
            self.SD_cls(1.0 * u.arcsec, 0.0, 0.0)
        with pytest.raises(TypeError):
            self.SD_cls(1.0 * u.arcsec, 0.0 * u.arcsec, 0.0 * u.kpc, False, False)
        with pytest.raises(TypeError):
            self.SD_cls(
                1.0 * u.arcsec,
                0.0 * u.arcsec,
                0.0 * u.kpc,
                copy=False,
                d_lat=0.0 * u.arcsec,
            )
        with pytest.raises(TypeError):
            self.SD_cls(
                1.0 * u.arcsec, 0.0 * u.arcsec, 0.0 * u.kpc, copy=False, flying="circus"
            )
        with pytest.raises(ValueError):
            self.SD_cls(
                np.ones(2) * u.arcsec, np.zeros(3) * u.arcsec, np.zeros(2) * u.kpc
            )
        with pytest.raises(u.UnitsError):
            self.SD_cls(1.0 * u.arcsec, 1.0 * u.s, 0.0 * u.kpc)
        with pytest.raises(u.UnitsError):
            self.SD_cls(1.0 * u.kpc, 1.0 * u.arcsec, 0.0 * u.kpc)
        o = self.SD_cls(1.0 * u.arcsec, 1.0 * u.arcsec, 0.0 * u.km / u.s)
        with pytest.raises(u.UnitsError):
            o.to_cartesian(s)
        with pytest.raises(AttributeError):
            o.d_lat = 0.0 * u.arcsec
        with pytest.raises(AttributeError):
            del o.d_lat

        o = self.SD_cls(1.0 * u.arcsec, 1.0 * u.arcsec, 0.0 * u.km)
        with pytest.raises(TypeError):
            o.to_cartesian()
        c = CartesianRepresentation(10.0, 0.0, 0.0, unit=u.km)
        with pytest.raises(TypeError):
            self.SD_cls.to_cartesian(c)
        with pytest.raises(TypeError):
            self.SD_cls.from_cartesian(c)
        with pytest.raises(TypeError):
            self.SD_cls.from_cartesian(c, SphericalRepresentation)


@pytest.mark.parametrize("omit_coslat", [False, True], scope="class")
class TestUnitSphericalDifferential:
    def _setup(self, omit_coslat):
        if omit_coslat:
            self.USD_cls = UnitSphericalCosLatDifferential
        else:
            self.USD_cls = UnitSphericalDifferential

        s = UnitSphericalRepresentation(
            lon=[0.0, 6.0, 21.0] * u.hourangle, lat=[0.0, -30.0, 85.0] * u.deg
        )
        self.s = s
        self.e = s.unit_vectors()
        self.sf = s.scale_factors(omit_coslat=omit_coslat)

    def test_name_coslat(self, omit_coslat):
        self._setup(omit_coslat)
        if omit_coslat:
            assert self.USD_cls is UnitSphericalCosLatDifferential
            assert self.USD_cls.name == "unitsphericalcoslat"
        else:
            assert self.USD_cls is UnitSphericalDifferential
            assert self.USD_cls.name == "unitspherical"
        assert self.USD_cls.name in DIFFERENTIAL_CLASSES

    def test_simple_differentials(self, omit_coslat):
        self._setup(omit_coslat)
        s, e, sf = self.s, self.e, self.sf

        o_lon = self.USD_cls(1.0 * u.arcsec, 0.0 * u.arcsec)
        o_lonc = o_lon.to_cartesian(base=s)
        o_lon2 = self.USD_cls.from_cartesian(o_lonc, base=s)
        assert_differential_allclose(o_lon, o_lon2)
        # simple check by hand for first element
        # (lat[0]=0, so works for both normal and CosLat differential)
        assert_quantity_allclose(
            o_lonc[0].xyz, [0.0, np.pi / 180.0 / 3600.0, 0.0] * u.one
        )
        # check all using unit vectors and scale factors.
        s_lon = s + 1.0 * u.arcsec * sf["lon"] * e["lon"]
        assert type(s_lon) is SphericalRepresentation
        assert_representation_allclose(o_lonc, s_lon - s, atol=1e-10 * u.one)
        s_lon2 = s + o_lon
        assert_representation_allclose(s_lon2, s_lon, atol=1e-10 * u.one)

        o_lat = self.USD_cls(0.0 * u.arcsec, 1.0 * u.arcsec)
        o_latc = o_lat.to_cartesian(base=s)
        assert_quantity_allclose(
            o_latc[0].xyz,
            [0.0, 0.0, np.pi / 180.0 / 3600.0] * u.one,
            atol=1e-10 * u.one,
        )
        s_lat = s + 1.0 * u.arcsec * sf["lat"] * e["lat"]
        assert type(s_lat) is SphericalRepresentation
        assert_representation_allclose(o_latc, s_lat - s, atol=1e-10 * u.one)
        s_lat2 = s + o_lat
        assert_representation_allclose(s_lat2, s_lat, atol=1e-10 * u.one)

    def test_differential_arithmetic(self, omit_coslat):
        self._setup(omit_coslat)
        s = self.s

        o_lon = self.USD_cls(1.0 * u.arcsec, 0.0 * u.arcsec)
        o_lon_by_2 = o_lon / 2.0
        assert type(o_lon_by_2) is self.USD_cls
        assert_representation_allclose(
            o_lon_by_2.to_cartesian(s) * 2.0, o_lon.to_cartesian(s), atol=1e-10 * u.one
        )
        s_lon = s + o_lon
        s_lon2 = s + 2 * o_lon_by_2
        assert type(s_lon) is SphericalRepresentation
        assert_representation_allclose(s_lon, s_lon2, atol=1e-10 * u.one)
        o_lon_rec = o_lon_by_2 + o_lon_by_2
        assert type(o_lon_rec) is self.USD_cls
        assert representation_equal(o_lon, o_lon_rec)
        assert_representation_allclose(s + o_lon, s + o_lon_rec, atol=1e-10 * u.one)
        o_lon_0 = o_lon - o_lon
        assert type(o_lon_0) is self.USD_cls
        for c in o_lon_0.components:
            assert np.all(getattr(o_lon_0, c) == 0.0)

        o_lon2 = self.USD_cls(1.0 * u.mas / u.yr, 0.0 * u.mas / u.yr)
        kks = u.km / u.kpc / u.s
        assert_quantity_allclose(o_lon2.norm(s)[0], 4.74047 * kks, atol=1e-4 * kks)
        assert_representation_allclose(
            o_lon2.to_cartesian(s) * 1000.0 * u.yr,
            o_lon.to_cartesian(s),
            atol=1e-10 * u.one,
        )
        s_off = s + o_lon
        s_off2 = s + o_lon2 * 1000.0 * u.yr
        assert_representation_allclose(s_off, s_off2, atol=1e-10 * u.one)

        factor = 1e5 * u.radian / u.arcsec
        if not omit_coslat:
            factor = factor / np.cos(s.lat)
        s_off_big = s + o_lon * factor

        assert_representation_allclose(
            s_off_big,
            SphericalRepresentation(s.lon + 90.0 * u.deg, 0.0 * u.deg, 1e5),
            atol=5.0 * u.one,
        )

        o_lon3c = CartesianRepresentation(0.0, 4.74047, 0.0, unit=kks)
        # This looses information!!
        o_lon3 = self.USD_cls.from_cartesian(o_lon3c, base=s)
        expected0 = self.USD_cls(1.0 * u.mas / u.yr, 0.0 * u.mas / u.yr)
        assert_differential_allclose(o_lon3[0], expected0)
        # Part of motion kept.
        part_kept = s.cross(CartesianRepresentation(0, 1, 0, unit=u.one)).norm()
        assert_quantity_allclose(
            o_lon3.norm(s), 4.74047 * part_kept * kks, atol=1e-10 * kks
        )
        # (lat[0]=0, so works for both normal and CosLat differential)
        s_off_big2 = s + o_lon3 * 1e5 * u.yr * u.radian / u.mas
        expected0 = SphericalRepresentation(90.0 * u.deg, 0.0 * u.deg, 1e5 * u.one)
        assert_representation_allclose(s_off_big2[0], expected0, atol=5.0 * u.one)

    def test_differential_init_errors(self, omit_coslat):
        self._setup(omit_coslat)
        with pytest.raises(u.UnitsError):
            self.USD_cls(0.0 * u.deg, 10.0 * u.deg / u.yr)


class TestRadialDifferential:
    def setup_method(self):
        s = SphericalRepresentation(
            lon=[0.0, 6.0, 21.0] * u.hourangle,
            lat=[0.0, -30.0, 85.0] * u.deg,
            distance=[1, 2, 3] * u.kpc,
        )
        self.s = s
        self.r = s.represent_as(RadialRepresentation)
        self.e = s.unit_vectors()
        self.sf = s.scale_factors()

    def test_name(self):
        assert RadialDifferential.name == "radial"
        assert RadialDifferential.name in DIFFERENTIAL_CLASSES

    def test_simple_differentials(self):
        r, s, e, sf = self.r, self.s, self.e, self.sf

        o_distance = RadialDifferential(1.0 * u.mpc)
        # Can be applied to RadialRepresentation, though not most useful.
        r_distance = r + o_distance
        assert_quantity_allclose(
            r_distance.distance, r.distance + o_distance.d_distance
        )
        r_distance2 = o_distance + r
        assert_quantity_allclose(
            r_distance2.distance, r.distance + o_distance.d_distance
        )
        # More sense to apply it relative to spherical representation.
        o_distancec = o_distance.to_cartesian(base=s)
        assert_quantity_allclose(
            o_distancec[0].xyz, [1e-6, 0.0, 0.0] * u.kpc, atol=1.0 * u.npc
        )
        o_recover = RadialDifferential.from_cartesian(o_distancec, base=s)
        assert_quantity_allclose(o_recover.d_distance, o_distance.d_distance)

        s_distance = s + 1.0 * u.mpc * sf["distance"] * e["distance"]
        assert_representation_allclose(o_distancec, s_distance - s, atol=1 * u.npc)
        s_distance2 = s + o_distance
        assert_representation_allclose(s_distance2, s_distance)


class TestPhysicsSphericalDifferential:
    """Test copied from SphericalDifferential, so less extensive."""

    def setup_method(self):
        s = PhysicsSphericalRepresentation(
            phi=[0.0, 90.0, 315.0] * u.deg,
            theta=[90.0, 120.0, 5.0] * u.deg,
            r=[1, 2, 3] * u.kpc,
        )
        self.s = s
        self.e = s.unit_vectors()
        self.sf = s.scale_factors()

    def test_name(self):
        assert PhysicsSphericalDifferential.name == "physicsspherical"
        assert PhysicsSphericalDifferential.name in DIFFERENTIAL_CLASSES

    def test_simple_differentials(self):
        s, e, sf = self.s, self.e, self.sf

        o_phi = PhysicsSphericalDifferential(1 * u.arcsec, 0 * u.arcsec, 0 * u.kpc)
        o_phic = o_phi.to_cartesian(base=s)
        o_phi2 = PhysicsSphericalDifferential.from_cartesian(o_phic, base=s)
        assert_quantity_allclose(o_phi.d_phi, o_phi2.d_phi, atol=1.0 * u.narcsec)
        assert_quantity_allclose(o_phi.d_theta, o_phi2.d_theta, atol=1.0 * u.narcsec)
        assert_quantity_allclose(o_phi.d_r, o_phi2.d_r, atol=1.0 * u.npc)
        # simple check by hand for first element.
        assert_quantity_allclose(
            o_phic[0].xyz, [0.0, np.pi / 180.0 / 3600.0, 0.0] * u.kpc, atol=1.0 * u.npc
        )
        # check all using unit vectors and scale factors.
        s_phi = s + 1.0 * u.arcsec * sf["phi"] * e["phi"]
        assert_representation_allclose(o_phic, s_phi - s, atol=1e-10 * u.kpc)

        o_theta = PhysicsSphericalDifferential(0 * u.arcsec, 1 * u.arcsec, 0 * u.kpc)
        o_thetac = o_theta.to_cartesian(base=s)
        assert_quantity_allclose(
            o_thetac[0].xyz,
            [0.0, 0.0, -np.pi / 180.0 / 3600.0] * u.kpc,
            atol=1.0 * u.npc,
        )
        s_theta = s + 1.0 * u.arcsec * sf["theta"] * e["theta"]
        assert_representation_allclose(o_thetac, s_theta - s, atol=1e-10 * u.kpc)
        s_theta2 = s + o_theta
        assert_representation_allclose(s_theta2, s_theta, atol=1e-10 * u.kpc)

        o_r = PhysicsSphericalDifferential(0 * u.arcsec, 0 * u.arcsec, 1 * u.mpc)
        o_rc = o_r.to_cartesian(base=s)
        assert_quantity_allclose(
            o_rc[0].xyz, [1e-6, 0.0, 0.0] * u.kpc, atol=1.0 * u.npc
        )
        s_r = s + 1.0 * u.mpc * sf["r"] * e["r"]
        assert_representation_allclose(o_rc, s_r - s, atol=1e-10 * u.kpc)
        s_r2 = s + o_r
        assert_representation_allclose(s_r2, s_r)

    def test_differential_init_errors(self):
        with pytest.raises(u.UnitsError):
            PhysicsSphericalDifferential(1.0 * u.arcsec, 0.0, 0.0)


class TestCylindricalDifferential:
    """Test copied from SphericalDifferential, so less extensive."""

    def setup_method(self):
        s = CylindricalRepresentation(
            rho=[1, 2, 3] * u.kpc, phi=[0.0, 90.0, 315.0] * u.deg, z=[3, 2, 1] * u.kpc
        )
        self.s = s
        self.e = s.unit_vectors()
        self.sf = s.scale_factors()

    def test_name(self):
        assert CylindricalDifferential.name == "cylindrical"
        assert CylindricalDifferential.name in DIFFERENTIAL_CLASSES

    def test_simple_differentials(self):
        s, e, sf = self.s, self.e, self.sf

        o_rho = CylindricalDifferential(1.0 * u.mpc, 0.0 * u.arcsec, 0.0 * u.kpc)
        o_rhoc = o_rho.to_cartesian(base=s)
        assert_quantity_allclose(o_rhoc[0].xyz, [1.0e-6, 0.0, 0.0] * u.kpc)
        s_rho = s + 1.0 * u.mpc * sf["rho"] * e["rho"]
        assert_representation_allclose(o_rhoc, s_rho - s, atol=1e-10 * u.kpc)
        s_rho2 = s + o_rho
        assert_representation_allclose(s_rho2, s_rho)

        o_phi = CylindricalDifferential(0.0 * u.kpc, 1.0 * u.arcsec, 0.0 * u.kpc)
        o_phic = o_phi.to_cartesian(base=s)
        o_phi2 = CylindricalDifferential.from_cartesian(o_phic, base=s)
        assert_quantity_allclose(o_phi.d_rho, o_phi2.d_rho, atol=1.0 * u.npc)
        assert_quantity_allclose(o_phi.d_phi, o_phi2.d_phi, atol=1.0 * u.narcsec)
        assert_quantity_allclose(o_phi.d_z, o_phi2.d_z, atol=1.0 * u.npc)
        # simple check by hand for first element.
        assert_quantity_allclose(
            o_phic[0].xyz, [0.0, np.pi / 180.0 / 3600.0, 0.0] * u.kpc
        )
        # check all using unit vectors and scale factors.
        s_phi = s + 1.0 * u.arcsec * sf["phi"] * e["phi"]
        assert_representation_allclose(o_phic, s_phi - s, atol=1e-10 * u.kpc)

        o_z = CylindricalDifferential(0.0 * u.kpc, 0.0 * u.arcsec, 1.0 * u.mpc)
        o_zc = o_z.to_cartesian(base=s)
        assert_quantity_allclose(o_zc[0].xyz, [0.0, 0.0, 1.0e-6] * u.kpc)
        s_z = s + 1.0 * u.mpc * sf["z"] * e["z"]
        assert_representation_allclose(o_zc, s_z - s, atol=1e-10 * u.kpc)
        s_z2 = s + o_z
        assert_representation_allclose(s_z2, s_z)

    def test_differential_init_errors(self):
        with pytest.raises(u.UnitsError):
            CylindricalDifferential(1.0 * u.pc, 1.0 * u.arcsec, 3.0 * u.km / u.s)


class TestCartesianDifferential:
    """Test copied from SphericalDifferential, so less extensive."""

    def setup_method(self):
        s = CartesianRepresentation(
            x=[1, 2, 3] * u.kpc, y=[2, 3, 1] * u.kpc, z=[3, 1, 2] * u.kpc
        )
        self.s = s
        self.e = s.unit_vectors()
        self.sf = s.scale_factors()

    def test_name(self):
        assert CartesianDifferential.name == "cartesian"
        assert CartesianDifferential.name in DIFFERENTIAL_CLASSES

    def test_simple_differentials(self):
        s, e, sf = self.s, self.e, self.sf

        for d, differential in (  # test different inits while we're at it.
            ("x", CartesianDifferential(1.0 * u.pc, 0.0 * u.pc, 0.0 * u.pc)),
            ("y", CartesianDifferential([0.0, 1.0, 0.0], unit=u.pc)),
            (
                "z",
                CartesianDifferential(np.array([[0.0, 0.0, 1.0]]) * u.pc, xyz_axis=1),
            ),
        ):
            o_c = differential.to_cartesian(base=s)
            o_c2 = differential.to_cartesian()
            assert np.all(representation_equal(o_c, o_c2))
            assert all(
                np.all(getattr(differential, "d_" + c) == getattr(o_c, c))
                for c in ("x", "y", "z")
            )
            differential2 = CartesianDifferential.from_cartesian(o_c)
            assert np.all(representation_equal(differential2, differential))
            differential3 = CartesianDifferential.from_cartesian(o_c, base=o_c)
            assert np.all(representation_equal(differential3, differential))

            s_off = s + 1.0 * u.pc * sf[d] * e[d]
            assert_representation_allclose(o_c, s_off - s, atol=1e-10 * u.kpc)
            s_off2 = s + differential
            assert_representation_allclose(s_off2, s_off)

    def test_init_failures(self):
        with pytest.raises(ValueError):
            CartesianDifferential(1.0 * u.kpc / u.s, 2.0 * u.kpc)
        with pytest.raises(u.UnitsError):
            CartesianDifferential(1.0 * u.kpc / u.s, 2.0 * u.kpc, 3.0 * u.kpc)
        with pytest.raises(ValueError):
            CartesianDifferential(1.0 * u.kpc, 2.0 * u.kpc, 3.0 * u.kpc, xyz_axis=1)


class TestDifferentialConversion:
    def setup_method(self):
        self.s = SphericalRepresentation(
            lon=[0.0, 6.0, 21.0] * u.hourangle,
            lat=[0.0, -30.0, 85.0] * u.deg,
            distance=[1, 2, 3] * u.kpc,
        )

    @pytest.mark.parametrize(
        "sd_cls", [SphericalDifferential, SphericalCosLatDifferential]
    )
    def test_represent_as_own_class(self, sd_cls):
        so = sd_cls(1.0 * u.deg, 2.0 * u.deg, 0.1 * u.kpc)
        so2 = so.represent_as(sd_cls)
        assert so2 is so

    def test_represent_other_coslat(self):
        s = self.s
        coslat = np.cos(s.lat)
        so = SphericalDifferential(1.0 * u.deg, 2.0 * u.deg, 0.1 * u.kpc)
        so_coslat = so.represent_as(SphericalCosLatDifferential, base=s)
        assert_quantity_allclose(so.d_lon * coslat, so_coslat.d_lon_coslat)
        so2 = so_coslat.represent_as(SphericalDifferential, base=s)
        assert np.all(representation_equal(so2, so))
        so3 = SphericalDifferential.from_representation(so_coslat, base=s)
        assert np.all(representation_equal(so3, so))
        so_coslat2 = SphericalCosLatDifferential.from_representation(so, base=s)
        assert np.all(representation_equal(so_coslat2, so_coslat))
        # Also test UnitSpherical
        us = s.represent_as(UnitSphericalRepresentation)
        uo = so.represent_as(UnitSphericalDifferential)
        uo_coslat = so.represent_as(UnitSphericalCosLatDifferential, base=s)
        assert_quantity_allclose(uo.d_lon * coslat, uo_coslat.d_lon_coslat)
        uo2 = uo_coslat.represent_as(UnitSphericalDifferential, base=us)
        assert np.all(representation_equal(uo2, uo))
        uo3 = UnitSphericalDifferential.from_representation(uo_coslat, base=us)
        assert np.all(representation_equal(uo3, uo))
        uo_coslat2 = UnitSphericalCosLatDifferential.from_representation(uo, base=us)
        assert np.all(representation_equal(uo_coslat2, uo_coslat))
        uo_coslat3 = uo.represent_as(UnitSphericalCosLatDifferential, base=us)
        assert np.all(representation_equal(uo_coslat3, uo_coslat))

    @pytest.mark.parametrize(
        "sd_cls", [SphericalDifferential, SphericalCosLatDifferential]
    )
    @pytest.mark.parametrize(
        "r_cls",
        (
            SphericalRepresentation,
            UnitSphericalRepresentation,
            PhysicsSphericalRepresentation,
            CylindricalRepresentation,
        ),
    )
    def test_represent_regular_class(self, sd_cls, r_cls):
        so = sd_cls(1.0 * u.deg, 2.0 * u.deg, 0.1 * u.kpc)
        r = so.represent_as(r_cls, base=self.s)
        c = so.to_cartesian(self.s)
        r_check = c.represent_as(r_cls)
        assert np.all(representation_equal(r, r_check))
        so2 = sd_cls.from_representation(r, base=self.s)
        so3 = sd_cls.from_cartesian(r.to_cartesian(), self.s)
        assert np.all(representation_equal(so2, so3))

    @pytest.mark.parametrize(
        "sd_cls", [SphericalDifferential, SphericalCosLatDifferential]
    )
    def test_convert_physics(self, sd_cls):
        # Conversion needs no base for SphericalDifferential, but does
        # need one (to get the latitude) for SphericalCosLatDifferential.
        if sd_cls is SphericalDifferential:
            usd_cls = UnitSphericalDifferential
            base_s = base_u = base_p = None
        else:
            usd_cls = UnitSphericalCosLatDifferential
            base_s = self.s[1]
            base_u = base_s.represent_as(UnitSphericalRepresentation)
            base_p = base_s.represent_as(PhysicsSphericalRepresentation)

        so = sd_cls(1.0 * u.deg, 2.0 * u.deg, 0.1 * u.kpc)
        po = so.represent_as(PhysicsSphericalDifferential, base=base_s)
        so2 = sd_cls.from_representation(po, base=base_s)
        assert_differential_allclose(so, so2)
        po2 = PhysicsSphericalDifferential.from_representation(so, base=base_p)
        assert_differential_allclose(po, po2)
        so3 = po.represent_as(sd_cls, base=base_p)
        assert_differential_allclose(so, so3)

        s = self.s
        p = s.represent_as(PhysicsSphericalRepresentation)
        cso = so.to_cartesian(s[1])
        cpo = po.to_cartesian(p[1])
        assert_representation_allclose(cso, cpo)
        assert_representation_allclose(s[1] + so, p[1] + po)
        po2 = so.represent_as(
            PhysicsSphericalDifferential, base=None if base_s is None else s
        )
        assert_representation_allclose(s + so, p + po2)

        suo = usd_cls.from_representation(so)
        puo = usd_cls.from_representation(po, base=base_u)
        assert_differential_allclose(suo, puo)
        suo2 = so.represent_as(usd_cls)
        puo2 = po.represent_as(usd_cls, base=base_p)
        assert_differential_allclose(suo2, puo2)
        assert_differential_allclose(puo, puo2)

        sro = RadialDifferential.from_representation(so)
        pro = RadialDifferential.from_representation(po)
        assert representation_equal(sro, pro)
        sro2 = so.represent_as(RadialDifferential)
        pro2 = po.represent_as(RadialDifferential)
        assert representation_equal(sro2, pro2)
        assert representation_equal(pro, pro2)

    @pytest.mark.parametrize(
        ("sd_cls", "usd_cls"),
        [
            (SphericalDifferential, UnitSphericalDifferential),
            (SphericalCosLatDifferential, UnitSphericalCosLatDifferential),
        ],
    )
    def test_convert_unit_spherical_radial(self, sd_cls, usd_cls):
        s = self.s
        us = s.represent_as(UnitSphericalRepresentation)
        rs = s.represent_as(RadialRepresentation)
        assert_representation_allclose(rs * us, s)

        uo = usd_cls(2.0 * u.deg, 1.0 * u.deg)
        so = uo.represent_as(sd_cls, base=s)
        assert_quantity_allclose(so.d_distance, 0.0 * u.kpc, atol=1.0 * u.npc)
        uo2 = so.represent_as(usd_cls)
        assert_representation_allclose(uo.to_cartesian(us), uo2.to_cartesian(us))
        so1 = sd_cls(2.0 * u.deg, 1.0 * u.deg, 5.0 * u.pc)
        uo_r = so1.represent_as(usd_cls)
        ro_r = so1.represent_as(RadialDifferential)
        assert np.all(representation_equal(uo_r, uo))
        assert np.all(representation_equal(ro_r, RadialDifferential(5.0 * u.pc)))

    @pytest.mark.parametrize(
        "sd_cls", [SphericalDifferential, SphericalCosLatDifferential]
    )
    def test_convert_cylindrial(self, sd_cls):
        s = self.s
        so = sd_cls(1.0 * u.deg, 2.0 * u.deg, 0.1 * u.kpc)
        cyo = so.represent_as(CylindricalDifferential, base=s)
        cy = s.represent_as(CylindricalRepresentation)
        so1 = cyo.represent_as(sd_cls, base=cy)
        assert_representation_allclose(so.to_cartesian(s), so1.to_cartesian(s))
        cyo2 = CylindricalDifferential.from_representation(so, base=cy)
        assert_representation_allclose(
            cyo2.to_cartesian(base=cy), cyo.to_cartesian(base=cy)
        )
        so2 = sd_cls.from_representation(cyo2, base=s)
        assert_representation_allclose(so.to_cartesian(s), so2.to_cartesian(s))

    @pytest.mark.parametrize(
        "sd_cls", [SphericalDifferential, SphericalCosLatDifferential]
    )
    def test_combinations(self, sd_cls):
        if sd_cls is SphericalDifferential:
            uo = UnitSphericalDifferential(2.0 * u.deg, 1.0 * u.deg)
            uo_d_lon = uo.d_lon
        else:
            uo = UnitSphericalCosLatDifferential(2.0 * u.deg, 1.0 * u.deg)
            uo_d_lon = uo.d_lon_coslat
        ro = RadialDifferential(1.0 * u.mpc)
        so1 = uo + ro
        so1c = sd_cls(uo_d_lon, uo.d_lat, ro.d_distance)
        assert np.all(representation_equal(so1, so1c))

        so2 = uo - ro
        so2c = sd_cls(uo_d_lon, uo.d_lat, -ro.d_distance)
        assert np.all(representation_equal(so2, so2c))
        so3 = so2 + ro
        so3c = sd_cls(uo_d_lon, uo.d_lat, 0.0 * u.kpc)
        assert np.all(representation_equal(so3, so3c))
        so4 = so1 + ro
        so4c = sd_cls(uo_d_lon, uo.d_lat, 2 * ro.d_distance)
        assert np.all(representation_equal(so4, so4c))
        so5 = so1 - uo
        so5c = sd_cls(0 * u.deg, 0.0 * u.deg, ro.d_distance)
        assert np.all(representation_equal(so5, so5c))
        assert_representation_allclose(self.s + (uo + ro), self.s + so1)


@pytest.mark.parametrize(
    "op,args",
    [
        (operator.neg, ()),
        (operator.pos, ()),
        (operator.mul, (-8.0,)),
        (operator.truediv, ([4.0, 8.0] * u.s,)),
    ],
    scope="class",
)
class TestArithmeticWithDifferentials:
    def setup_class(self):
        self.cr = CartesianRepresentation([1, 2, 3] * u.kpc)
        self.cd = CartesianDifferential([0.1, -0.2, 0.3] * u.km / u.s)
        self.c = self.cr.with_differentials(self.cd)

    def test_operation_cartesian(self, op, args):
        ncr = op(self.c, *args)
        expected = (op(self.cr, *args)).with_differentials(op(self.cd, *args))
        assert np.all(ncr == expected)

    def test_operation_radial(self, op, args):
        rep = self.c.represent_as(RadialRepresentation, {"s": RadialDifferential})
        result = op(rep, *args)
        expected_distance = op(self.cr.norm(), *args)
        expected_rv = op((self.cr / self.cr.norm()).dot(self.cd), *args)
        assert u.allclose(result.distance, expected_distance)
        assert u.allclose(result.differentials["s"].d_distance, expected_rv)

    @pytest.mark.parametrize(
        "diff_cls",
        [
            SphericalDifferential,
            SphericalCosLatDifferential,
            PhysicsSphericalDifferential,
            CylindricalDifferential,
        ],
    )
    def test_operation_other(self, diff_cls, op, args):
        rep_cls = diff_cls.base_representation
        rep = self.c.represent_as(rep_cls, {"s": diff_cls})
        result = op(rep, *args)

        expected_c = op(self.c, *args)
        expected = expected_c.represent_as(rep_cls, {"s": diff_cls})
        # Check that we match in the representation itself.
        assert_representation_allclose(result, expected)
        assert_differential_allclose(
            result.differentials["s"], expected.differentials["s"]
        )

        # Check that we compare correctly in cartesian as well, just to be sure.
        result_c = result.represent_as(
            CartesianRepresentation, {"s": CartesianDifferential}
        )
        assert_representation_allclose(result_c, expected_c)
        assert_differential_allclose(
            result_c.differentials["s"], expected_c.differentials["s"]
        )

    @pytest.mark.parametrize(
        "rep_cls",
        [
            SphericalRepresentation,
            PhysicsSphericalRepresentation,
            CylindricalRepresentation,
        ],
    )
    def test_operation_cartesian_differential(self, rep_cls, op, args):
        rep = self.c.represent_as(rep_cls, {"s": CartesianDifferential})
        result = op(rep, *args)

        expected_c = op(self.c, *args)
        expected = expected_c.represent_as(rep_cls, {"s": CartesianDifferential})
        # Check that we match in the representation itself.
        assert_representation_allclose(result, expected)
        assert_differential_allclose(
            result.differentials["s"], expected.differentials["s"]
        )

    @pytest.mark.parametrize(
        "diff_cls", [UnitSphericalDifferential, UnitSphericalCosLatDifferential]
    )
    def test_operation_unit_spherical(self, diff_cls, op, args):
        rep_cls = diff_cls.base_representation
        rep = self.c.represent_as(rep_cls, {"s": diff_cls})
        result = op(rep, *args)

        if op not in (operator.neg, operator.pos):
            expected_cls = rep._dimensional_representation
        else:
            expected_cls = rep_cls

        assert type(result) is expected_cls
        assert type(result.differentials["s"]) is diff_cls

        # Have lost information, so unlike above we convert our initial
        # unit-spherical back to Cartesian, and check that applying
        # the operation on that cartesian representation gives the same result.
        # We do not compare the output directly, since for multiplication
        # and division there will be sign flips in the spherical distance.
        expected_c = op(
            rep.represent_as(CartesianRepresentation, {"s": CartesianDifferential}),
            *args,
        )
        result_c = result.represent_as(
            CartesianRepresentation, {"s": CartesianDifferential}
        )
        assert_representation_allclose(result_c, expected_c)
        assert_differential_allclose(
            result_c.differentials["s"], expected_c.differentials["s"]
        )

    @pytest.mark.parametrize(
        "diff_cls",
        [
            RadialDifferential,
            UnitSphericalDifferential,
            UnitSphericalCosLatDifferential,
        ],
    )
    def test_operation_spherical_with_rv_or_pm(self, diff_cls, op, args):
        rep = self.c.represent_as(SphericalRepresentation, {"s": diff_cls})
        result = op(rep, *args)
        assert type(result) is SphericalRepresentation
        assert type(result.differentials["s"]) is diff_cls

        expected_c = op(
            rep.represent_as(CartesianRepresentation, {"s": CartesianDifferential}),
            *args,
        )
        result_c = result.represent_as(
            CartesianRepresentation, {"s": CartesianDifferential}
        )
        assert_representation_allclose(result_c, expected_c)
        assert_differential_allclose(
            result_c.differentials["s"], expected_c.differentials["s"]
        )


@pytest.mark.parametrize("op,args", [(operator.neg, ()), (operator.mul, (10.0,))])
def test_operation_unitspherical_with_rv_fails(op, args):
    rep = UnitSphericalRepresentation(
        0 * u.deg, 0 * u.deg, differentials={"s": RadialDifferential(10 * u.km / u.s)}
    )
    with pytest.raises(ValueError, match="unit key"):
        op(rep, *args)


@pytest.mark.parametrize(
    "rep,dif",
    [
        [
            CartesianRepresentation([1, 2, 3] * u.kpc),
            CartesianDifferential([0.1, 0.2, 0.3] * u.km / u.s),
        ],
        [
            SphericalRepresentation(90 * u.deg, 0.0 * u.deg, 14.0 * u.kpc),
            SphericalDifferential(1.0 * u.deg, 2.0 * u.deg, 0.1 * u.kpc),
        ],
    ],
)
def test_arithmetic_with_differentials_fail(rep, dif):
    rep = rep.with_differentials(dif)

    with pytest.raises(TypeError):
        rep + rep

    with pytest.raises(TypeError):
        rep - rep

    with pytest.raises(TypeError):
        rep * rep

    with pytest.raises(TypeError):
        rep / rep
