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


import numpy as np
import pytest

from astropy import units as u
from astropy.coordinates import (
    Latitude,
    Longitude,
    SphericalDifferential,
    SphericalRepresentation,
)
from astropy.units.quantity_helper.function_helpers import ARRAY_FUNCTION_ENABLED

from .test_representation import representation_equal


@pytest.fixture(params=[True, False] if ARRAY_FUNCTION_ENABLED else [True])
def method(request):
    return request.param


needs_array_function = pytest.mark.xfail(
    not ARRAY_FUNCTION_ENABLED, reason="Needs __array_function__ support"
)


class ShapeSetup:
    """Manipulation of Representation shapes.

    Checking that attributes are manipulated correctly.

    Even more exhaustive tests are done in time.tests.test_methods
    """

    def setup_class(cls):
        # We set up some representations with, on purpose, copy=False,
        # so we can check that broadcasting is handled correctly.
        lon = Longitude(np.arange(0, 24, 4), u.hourangle)
        lat = Latitude(np.arange(-90, 91, 30), u.deg)

        # With same-sized arrays
        cls.s0 = SphericalRepresentation(
            lon[:, np.newaxis] * np.ones(lat.shape),
            lat * np.ones(lon.shape)[:, np.newaxis],
            np.ones(lon.shape + lat.shape) * u.kpc,
            copy=False,
        )

        cls.diff = SphericalDifferential(
            d_lon=np.ones(cls.s0.shape) * u.mas / u.yr,
            d_lat=np.ones(cls.s0.shape) * u.mas / u.yr,
            d_distance=np.ones(cls.s0.shape) * u.km / u.s,
            copy=False,
        )
        cls.s0 = cls.s0.with_differentials(cls.diff)

        # With unequal arrays -> these will be broadcasted.
        cls.s1 = SphericalRepresentation(
            lon[:, np.newaxis], lat, 1.0 * u.kpc, differentials=cls.diff, copy=False
        )

        # For completeness on some tests, also a cartesian one
        cls.c0 = cls.s0.to_cartesian()


class TestManipulation(ShapeSetup):
    """Manipulation of Representation shapes.

    Checking that attributes are manipulated correctly.

    Even more exhaustive tests are done in time.tests.test_methods
    """

    def test_ravel(self, method):
        if method:
            s0_ravel = self.s0.ravel()
        else:
            s0_ravel = np.ravel(self.s0)
        assert type(s0_ravel) is type(self.s0)
        assert s0_ravel.shape == (self.s0.size,)
        assert np.all(s0_ravel.lon == self.s0.lon.ravel())
        assert np.may_share_memory(s0_ravel.lon, self.s0.lon)
        assert np.may_share_memory(s0_ravel.lat, self.s0.lat)
        assert np.may_share_memory(s0_ravel.distance, self.s0.distance)
        assert s0_ravel.differentials["s"].shape == (self.s0.size,)

        # Since s1 was broadcast, ravel needs to make a copy.
        if method:
            s1_ravel = self.s1.ravel()
        else:
            s1_ravel = np.ravel(self.s1)
        assert type(s1_ravel) is type(self.s1)
        assert s1_ravel.shape == (self.s1.size,)
        assert s1_ravel.differentials["s"].shape == (self.s1.size,)
        assert np.all(s1_ravel.lon == self.s1.lon.ravel())
        assert not np.may_share_memory(s1_ravel.lat, self.s1.lat)

    def test_copy(self, method):
        if method:
            s0_copy = self.s0.copy()
        else:
            s0_copy = np.copy(self.s0)
        s0_copy_diff = s0_copy.differentials["s"]
        assert s0_copy.shape == self.s0.shape
        assert np.all(s0_copy.lon == self.s0.lon)
        assert np.all(s0_copy.lat == self.s0.lat)

        # Check copy was made of internal data.
        assert not np.may_share_memory(s0_copy.distance, self.s0.distance)
        assert not np.may_share_memory(s0_copy_diff.d_lon, self.diff.d_lon)

    def test_flatten(self):
        s0_flatten = self.s0.flatten()
        s0_diff = s0_flatten.differentials["s"]
        assert s0_flatten.shape == (self.s0.size,)
        assert s0_diff.shape == (self.s0.size,)
        assert np.all(s0_flatten.lon == self.s0.lon.flatten())
        assert np.all(s0_diff.d_lon == self.diff.d_lon.flatten())

        # Flatten always copies.
        assert not np.may_share_memory(s0_flatten.distance, self.s0.distance)
        assert not np.may_share_memory(s0_diff.d_lon, self.diff.d_lon)

        s1_flatten = self.s1.flatten()
        assert s1_flatten.shape == (self.s1.size,)
        assert np.all(s1_flatten.lon == self.s1.lon.flatten())
        assert not np.may_share_memory(s1_flatten.lat, self.s1.lat)

    def test_transpose(self):
        s0_transpose = self.s0.transpose()
        s0_diff = s0_transpose.differentials["s"]
        assert s0_transpose.shape == (7, 6)
        assert s0_diff.shape == s0_transpose.shape
        assert np.all(s0_transpose.lon == self.s0.lon.transpose())
        assert np.all(s0_diff.d_lon == self.diff.d_lon.transpose())
        assert np.may_share_memory(s0_transpose.distance, self.s0.distance)
        assert np.may_share_memory(s0_diff.d_lon, self.diff.d_lon)

        s1_transpose = self.s1.transpose()
        s1_diff = s1_transpose.differentials["s"]
        assert s1_transpose.shape == (7, 6)
        assert s1_diff.shape == s1_transpose.shape
        assert np.all(s1_transpose.lat == self.s1.lat.transpose())
        assert np.all(s1_diff.d_lon == self.diff.d_lon.transpose())
        assert np.may_share_memory(s1_transpose.lat, self.s1.lat)
        assert np.may_share_memory(s1_diff.d_lon, self.diff.d_lon)

        # Only one check on T, since it just calls transpose anyway.
        # Doing it on the CartesianRepresentation just for variety's sake.
        c0_T = self.c0.T
        assert c0_T.shape == (7, 6)
        assert np.all(c0_T.x == self.c0.x.T)
        assert np.may_share_memory(c0_T.y, self.c0.y)

    def test_diagonal(self):
        s0_diagonal = self.s0.diagonal()
        s0_diff = s0_diagonal.differentials["s"]
        assert s0_diagonal.shape == (6,)
        assert s0_diff.shape == s0_diagonal.shape
        assert np.all(s0_diagonal.lat == self.s0.lat.diagonal())
        assert np.all(s0_diff.d_lon == self.diff.d_lon.diagonal())
        assert np.may_share_memory(s0_diagonal.lat, self.s0.lat)
        assert np.may_share_memory(s0_diff.d_lon, self.diff.d_lon)

    def test_swapaxes(self, method):
        if method:
            s1_swapaxes = self.s1.swapaxes(0, 1)
        else:
            s1_swapaxes = np.swapaxes(self.s1, 0, 1)
        s1_diff = s1_swapaxes.differentials["s"]
        assert s1_swapaxes.shape == (7, 6)
        assert s1_diff.shape == s1_swapaxes.shape
        assert np.all(s1_swapaxes.lat == self.s1.lat.swapaxes(0, 1))
        assert np.all(s1_diff.d_lon == self.diff.d_lon.swapaxes(0, 1))
        assert np.may_share_memory(s1_swapaxes.lat, self.s1.lat)
        assert np.may_share_memory(s1_diff.d_lon, self.diff.d_lon)

    def test_reshape(self):
        s0_reshape = self.s0.reshape(2, 3, 7)
        s0_diff = s0_reshape.differentials["s"]
        assert s0_reshape.shape == (2, 3, 7)
        assert s0_diff.shape == s0_reshape.shape
        assert np.all(s0_reshape.lon == self.s0.lon.reshape(2, 3, 7))
        assert np.all(s0_reshape.lat == self.s0.lat.reshape(2, 3, 7))
        assert np.all(s0_reshape.distance == self.s0.distance.reshape(2, 3, 7))
        assert np.may_share_memory(s0_reshape.lon, self.s0.lon)
        assert np.may_share_memory(s0_reshape.lat, self.s0.lat)
        assert np.may_share_memory(s0_reshape.distance, self.s0.distance)

        s1_reshape = self.s1.reshape(3, 2, 7)
        s1_diff = s1_reshape.differentials["s"]
        assert s1_reshape.shape == (3, 2, 7)
        assert s1_diff.shape == s1_reshape.shape
        assert np.all(s1_reshape.lat == self.s1.lat.reshape(3, 2, 7))
        assert np.all(s1_diff.d_lon == self.diff.d_lon.reshape(3, 2, 7))
        assert np.may_share_memory(s1_reshape.lat, self.s1.lat)
        assert np.may_share_memory(s1_diff.d_lon, self.diff.d_lon)

        # For reshape(3, 14), copying is necessary for lon, lat, but not for d
        s1_reshape2 = self.s1.reshape(3, 14)
        assert s1_reshape2.shape == (3, 14)
        assert np.all(s1_reshape2.lon == self.s1.lon.reshape(3, 14))
        assert not np.may_share_memory(s1_reshape2.lon, self.s1.lon)
        assert s1_reshape2.distance.shape == (3, 14)
        assert np.may_share_memory(s1_reshape2.distance, self.s1.distance)

    def test_squeeze(self):
        s0_squeeze = self.s0.reshape(3, 1, 2, 1, 7).squeeze()
        s0_diff = s0_squeeze.differentials["s"]
        assert s0_squeeze.shape == (3, 2, 7)
        assert s0_diff.shape == s0_squeeze.shape
        assert np.all(s0_squeeze.lat == self.s0.lat.reshape(3, 2, 7))
        assert np.all(s0_diff.d_lon == self.diff.d_lon.reshape(3, 2, 7))
        assert np.may_share_memory(s0_squeeze.lat, self.s0.lat)

    def test_add_dimension(self):
        s0_adddim = self.s0[:, np.newaxis, :]
        s0_diff = s0_adddim.differentials["s"]
        assert s0_adddim.shape == (6, 1, 7)
        assert s0_diff.shape == s0_adddim.shape
        assert np.all(s0_adddim.lon == self.s0.lon[:, np.newaxis, :])
        assert np.all(s0_diff.d_lon == self.diff.d_lon[:, np.newaxis, :])
        assert np.may_share_memory(s0_adddim.lat, self.s0.lat)

    def test_take(self, method):
        if method:
            s0_take = self.s0.take((5, 2))
        else:
            s0_take = np.take(self.s0, (5, 2))
        s0_diff = s0_take.differentials["s"]
        assert s0_take.shape == (2,)
        assert s0_diff.shape == s0_take.shape
        assert np.all(s0_take.lon == self.s0.lon.take((5, 2)))
        assert np.all(s0_diff.d_lon == self.diff.d_lon.take((5, 2)))

    def test_broadcast_to_via_apply(self):
        s0_broadcast = self.s0._apply(np.broadcast_to, (3, 6, 7), subok=True)
        s0_diff = s0_broadcast.differentials["s"]
        assert type(s0_broadcast) is type(self.s0)
        assert s0_broadcast.shape == (3, 6, 7)
        assert s0_diff.shape == s0_broadcast.shape
        assert np.all(s0_broadcast.lon == self.s0.lon)
        assert np.all(s0_broadcast.lat == self.s0.lat)
        assert np.all(s0_broadcast.distance == self.s0.distance)
        assert np.may_share_memory(s0_broadcast.lon, self.s0.lon)
        assert np.may_share_memory(s0_broadcast.lat, self.s0.lat)
        assert np.may_share_memory(s0_broadcast.distance, self.s0.distance)


class TestSetShape(ShapeSetup):
    def test_shape_setting(self):
        # Shape-setting should be on the object itself, since copying removes
        # zero-strides due to broadcasting.  Hence, this should be the only
        # test in this class.
        self.s0.shape = (2, 3, 7)
        assert self.s0.shape == (2, 3, 7)
        assert self.s0.lon.shape == (2, 3, 7)
        assert self.s0.lat.shape == (2, 3, 7)
        assert self.s0.distance.shape == (2, 3, 7)
        assert self.diff.shape == (2, 3, 7)
        assert self.diff.d_lon.shape == (2, 3, 7)
        assert self.diff.d_lat.shape == (2, 3, 7)
        assert self.diff.d_distance.shape == (2, 3, 7)

        # this works with the broadcasting.
        self.s1.shape = (2, 3, 7)
        assert self.s1.shape == (2, 3, 7)
        assert self.s1.lon.shape == (2, 3, 7)
        assert self.s1.lat.shape == (2, 3, 7)
        assert self.s1.distance.shape == (2, 3, 7)
        assert self.s1.distance.strides == (0, 0, 0)

        # but this one does not.
        oldshape = self.s1.shape
        with pytest.raises(ValueError):
            self.s1.shape = (1,)
        with pytest.raises(AttributeError):
            self.s1.shape = (42,)
        assert self.s1.shape == oldshape
        assert self.s1.lon.shape == oldshape
        assert self.s1.lat.shape == oldshape
        assert self.s1.distance.shape == oldshape

        # Finally, a more complicated one that checks that things get reset
        # properly if it is not the first component that fails.
        s2 = SphericalRepresentation(
            self.s1.lon.copy(), self.s1.lat, self.s1.distance, copy=False
        )
        assert 0 not in s2.lon.strides
        assert 0 in s2.lat.strides
        with pytest.raises(AttributeError):
            s2.shape = (42,)
        assert s2.shape == oldshape
        assert s2.lon.shape == oldshape
        assert s2.lat.shape == oldshape
        assert s2.distance.shape == oldshape
        assert 0 not in s2.lon.strides
        assert 0 in s2.lat.strides


class TestShapeFunctions(ShapeSetup):
    @needs_array_function
    def test_broadcast_to(self):
        s0_broadcast = np.broadcast_to(self.s0, (3, 6, 7))
        s0_diff = s0_broadcast.differentials["s"]
        assert type(s0_broadcast) is type(self.s0)
        assert s0_broadcast.shape == (3, 6, 7)
        assert s0_diff.shape == s0_broadcast.shape
        assert np.all(s0_broadcast.lon == self.s0.lon)
        assert np.all(s0_broadcast.lat == self.s0.lat)
        assert np.all(s0_broadcast.distance == self.s0.distance)
        assert np.may_share_memory(s0_broadcast.lon, self.s0.lon)
        assert np.may_share_memory(s0_broadcast.lat, self.s0.lat)
        assert np.may_share_memory(s0_broadcast.distance, self.s0.distance)

        s1_broadcast = np.broadcast_to(self.s1, shape=(3, 6, 7))
        s1_diff = s1_broadcast.differentials["s"]
        assert s1_broadcast.shape == (3, 6, 7)
        assert s1_diff.shape == s1_broadcast.shape
        assert np.all(s1_broadcast.lat == self.s1.lat)
        assert np.all(s1_broadcast.lon == self.s1.lon)
        assert np.all(s1_broadcast.distance == self.s1.distance)
        assert s1_broadcast.distance.shape == (3, 6, 7)
        assert np.may_share_memory(s1_broadcast.lat, self.s1.lat)
        assert np.may_share_memory(s1_broadcast.lon, self.s1.lon)
        assert np.may_share_memory(s1_broadcast.distance, self.s1.distance)

        # A final test that "may_share_memory" equals "does_share_memory"
        # Do this on a copy, to keep self.s0 unchanged.
        sc = self.s0.copy()
        assert not np.may_share_memory(sc.lon, self.s0.lon)
        assert not np.may_share_memory(sc.lat, self.s0.lat)
        sc_broadcast = np.broadcast_to(sc, (3, 6, 7))
        assert np.may_share_memory(sc_broadcast.lon, sc.lon)
        # Can only write to copy, not to broadcast version.
        sc.lon[0, 0] = 22.0 * u.hourangle
        assert np.all(sc_broadcast.lon[:, 0, 0] == 22.0 * u.hourangle)

    @needs_array_function
    def test_atleast_1d(self):
        s00 = self.s0.ravel()[0]
        assert s00.ndim == 0
        s00_1d = np.atleast_1d(s00)
        assert s00_1d.ndim == 1
        assert np.all(representation_equal(s00[np.newaxis], s00_1d))
        assert np.may_share_memory(s00_1d.lon, s00.lon)

    @needs_array_function
    def test_atleast_2d(self):
        s0r = self.s0.ravel()
        assert s0r.ndim == 1
        s0r_2d = np.atleast_2d(s0r)
        assert s0r_2d.ndim == 2
        assert np.all(representation_equal(s0r[np.newaxis], s0r_2d))
        assert np.may_share_memory(s0r_2d.lon, s0r.lon)

    @needs_array_function
    def test_atleast_3d(self):
        assert self.s0.ndim == 2
        s0_3d, s1_3d = np.atleast_3d(self.s0, self.s1)
        assert s0_3d.ndim == s1_3d.ndim == 3
        assert np.all(representation_equal(self.s0[:, :, np.newaxis], s0_3d))
        assert np.all(representation_equal(self.s1[:, :, np.newaxis], s1_3d))
        assert np.may_share_memory(s0_3d.lon, self.s0.lon)

    def test_move_axis(self):
        # Goes via transpose so works without __array_function__ as well.
        s0_10 = np.moveaxis(self.s0, 0, 1)
        assert s0_10.shape == (self.s0.shape[1], self.s0.shape[0])
        assert np.all(representation_equal(self.s0.T, s0_10))
        assert np.may_share_memory(s0_10.lon, self.s0.lon)

    def test_roll_axis(self):
        # Goes via transpose so works without __array_function__ as well.
        s0_10 = np.rollaxis(self.s0, 1)
        assert s0_10.shape == (self.s0.shape[1], self.s0.shape[0])
        assert np.all(representation_equal(self.s0.T, s0_10))
        assert np.may_share_memory(s0_10.lon, self.s0.lon)

    @needs_array_function
    def test_fliplr(self):
        s0_lr = np.fliplr(self.s0)
        assert np.all(representation_equal(self.s0[:, ::-1], s0_lr))
        assert np.may_share_memory(s0_lr.lon, self.s0.lon)

    @needs_array_function
    def test_rot90(self):
        s0_270 = np.rot90(self.s0, 3)
        assert np.all(representation_equal(self.s0.T[:, ::-1], s0_270))
        assert np.may_share_memory(s0_270.lon, self.s0.lon)

    @needs_array_function
    def test_roll(self):
        s0r = np.roll(self.s0, 1, axis=0)
        assert np.all(representation_equal(s0r[1:], self.s0[:-1]))
        assert np.all(representation_equal(s0r[0], self.s0[-1]))

    @needs_array_function
    def test_delete(self):
        s0d = np.delete(self.s0, [2, 3], axis=0)
        assert np.all(representation_equal(s0d[:2], self.s0[:2]))
        assert np.all(representation_equal(s0d[2:], self.s0[4:]))

    @pytest.mark.parametrize("attribute", ["shape", "ndim", "size"])
    def test_shape_attribute_functions(self, attribute):
        function = getattr(np, attribute)
        result = function(self.s0)
        assert result == getattr(self.s0, attribute)