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

import io
import warnings

import numpy as np
import numpy.testing as npt
import pytest

from astropy import units as u
from astropy.table import Column, QTable, Row, Table, hstack
from astropy.table.bst import BST
from astropy.table.column import BaseColumn
from astropy.table.index import SlicedIndex, get_index
from astropy.table.soco import SCEngine
from astropy.table.sorted_array import SortedArray
from astropy.time import Time
from astropy.utils.compat.optional_deps import HAS_H5PY, HAS_SORTEDCONTAINERS
from astropy.utils.exceptions import AstropyDeprecationWarning, AstropyWarning

from .test_table import SetupData

available_engines = [BST, SortedArray]
NATIVE_INT_NAME = np.array(0).dtype.name

if HAS_SORTEDCONTAINERS:
    available_engines.append(SCEngine)


@pytest.fixture(params=available_engines)
def engine(request):
    return request.param


_col = [1, 2, 3, 4, 5]


@pytest.fixture(
    params=[
        _col,
        u.Quantity(_col),
        Time(_col, format="jyear"),
    ]
)
def main_col(request):
    return request.param


def assert_col_equal(col, array):
    __tracebackhide__ = True
    if isinstance(col, Time):
        assert np.all(col == Time(array, format="jyear"))
    else:
        assert np.all(col == col.__class__(array))


def assert_tables_equal(t1: Table, t2: Table) -> None:
    # Check table colnames equal and values equal
    vals_eq = t1.values_equal(t2)  # this raises if colnames not equal
    for col_eq in vals_eq.itercols():
        npt.assert_equal(col_eq, True)

    assert t1.meta == t2.meta


@pytest.mark.usefixtures("table_types")
class TestIndex(SetupData):
    def _setup(self, main_col, table_types):
        super()._setup(table_types)
        self.main_col = main_col
        if isinstance(main_col, u.Quantity):
            self._table_type = QTable
        if not isinstance(main_col, list):
            self._column_type = lambda x: x  # don't change mixin type
        self.mutable = isinstance(main_col, (list, u.Quantity))

    def make_col(self, name, lst):
        return self._column_type(lst, name=name)

    def make_val(self, val):
        if isinstance(self.main_col, Time):
            return Time(val, format="jyear")
        return val

    @property
    def t(self):
        if not hasattr(self, "_t"):
            # Note that order of columns is important, and the 'a' column is
            # last to ensure that the index column does not need to be the first
            # column (as was discovered in #10025).  Most testing uses 'a' and
            # ('a', 'b') for the columns.
            self._t = self._table_type()
            self._t["b"] = self._column_type([4.0, 5.1, 6.2, 7.0, 1.1])
            self._t["c"] = self._column_type(["7", "8", "9", "10", "11"])
            self._t["a"] = self._column_type(self.main_col)
        return self._t

    @pytest.mark.parametrize("composite", [False, True])
    def test_table_index(self, main_col, table_types, composite, engine):
        self._setup(main_col, table_types)
        t = self.t
        t.add_index(("a", "b") if composite else "a", engine=engine)
        assert np.all(t.indices[0].sorted_data() == [0, 1, 2, 3, 4])

        if not self.mutable:
            return

        # test altering table columns
        t["a"][0] = 4
        t.add_row((6.0, "7", 6))
        t["a"][3] = 10
        t.remove_row(2)
        t.add_row((5.0, "9", 4))

        assert_col_equal(t["a"], np.array([4, 2, 10, 5, 6, 4]))
        assert np.allclose(t["b"], np.array([4.0, 5.1, 7.0, 1.1, 6.0, 5.0]))
        assert np.all(t["c"].data == np.array(["7", "8", "10", "11", "7", "9"]))
        index = t.indices[0]
        ll = list(index.data.items())

        if composite:
            assert np.all(
                ll
                == [
                    ((2, 5.1), [1]),
                    ((4, 4.0), [0]),
                    ((4, 5.0), [5]),
                    ((5, 1.1), [3]),
                    ((6, 6.0), [4]),
                    ((10, 7.0), [2]),
                ]
            )
        else:
            assert np.all(
                ll
                == [((2,), [1]), ((4,), [0, 5]), ((5,), [3]), ((6,), [4]), ((10,), [2])]
            )
        t.remove_indices("a")
        assert len(t.indices) == 0

    def test_table_slicing(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        t = self.t
        t.add_index("a", engine=engine)
        assert np.all(t.indices[0].sorted_data() == [0, 1, 2, 3, 4])

        for slice_ in ([0, 2], np.array([0, 2])):
            t2 = t[slice_]
            # t2 should retain an index on column 'a'
            assert len(t2.indices) == 1
            assert_col_equal(t2["a"], [1, 3])

            # the index in t2 should reorder row numbers after slicing
            assert np.all(t2.indices[0].sorted_data() == [0, 1])
            # however, this index should be a deep copy of t1's index
            assert np.all(t.indices[0].sorted_data() == [0, 1, 2, 3, 4])

    def test_remove_rows(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        if not self.mutable:
            return
        t = self.t
        t.add_index("a", engine=engine)

        # remove individual row
        t2 = t.copy()
        t2.remove_rows(2)
        assert_col_equal(t2["a"], [1, 2, 4, 5])
        assert np.all(t2.indices[0].sorted_data() == [0, 1, 2, 3])

        # remove by list, ndarray, or slice
        for cut in ([0, 2, 4], np.array([0, 2, 4]), slice(0, 5, 2)):
            t2 = t.copy()
            t2.remove_rows(cut)
            assert_col_equal(t2["a"], [2, 4])
            assert np.all(t2.indices[0].sorted_data() == [0, 1])

        with pytest.raises(ValueError):
            t.remove_rows((0, 2, 4))

    def test_col_get_slice(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        t = self.t
        t.add_index("a", engine=engine)

        # get slice
        t2 = t[1:3]  # table slice
        assert_col_equal(t2["a"], [2, 3])
        assert np.all(t2.indices[0].sorted_data() == [0, 1])

        col_slice = t["a"][1:3]
        assert_col_equal(col_slice, [2, 3])
        # true column slices discard indices
        if isinstance(t["a"], BaseColumn):
            assert len(col_slice.info.indices) == 0

        # take slice of slice
        t2 = t[::2]
        assert_col_equal(t2["a"], np.array([1, 3, 5]))
        t3 = t2[::-1]
        assert_col_equal(t3["a"], np.array([5, 3, 1]))
        assert np.all(t3.indices[0].sorted_data() == [2, 1, 0])
        t3 = t2[:2]
        assert_col_equal(t3["a"], np.array([1, 3]))
        assert np.all(t3.indices[0].sorted_data() == [0, 1])
        # out-of-bound slices
        for t_empty in (t2[3:], t2[2:1], t3[2:]):
            assert len(t_empty["a"]) == 0
            assert np.all(t_empty.indices[0].sorted_data() == [])

        if self.mutable:
            # get boolean mask
            mask = t["a"] % 2 == 1
            t2 = t[mask]
            assert_col_equal(t2["a"], [1, 3, 5])
            assert np.all(t2.indices[0].sorted_data() == [0, 1, 2])

    def test_col_set_slice(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        if not self.mutable:
            return
        t = self.t
        t.add_index("a", engine=engine)

        # set slice
        t2 = t.copy()
        t2["a"][1:3] = np.array([6, 7])
        assert_col_equal(t2["a"], np.array([1, 6, 7, 4, 5]))
        assert np.all(t2.indices[0].sorted_data() == [0, 3, 4, 1, 2])

        # change original table via slice reference
        t2 = t.copy()
        t3 = t2[1:3]
        assert_col_equal(t3["a"], np.array([2, 3]))
        assert np.all(t3.indices[0].sorted_data() == [0, 1])
        t3["a"][0] = 5
        assert_col_equal(t3["a"], np.array([5, 3]))
        assert_col_equal(t2["a"], np.array([1, 5, 3, 4, 5]))
        assert np.all(t3.indices[0].sorted_data() == [1, 0])
        assert np.all(t2.indices[0].sorted_data() == [0, 2, 3, 1, 4])

        # set boolean mask
        t2 = t.copy()
        mask = t["a"] % 2 == 1
        t2["a"][mask] = 0.0
        assert_col_equal(t2["a"], [0, 2, 0, 4, 0])
        assert np.all(t2.indices[0].sorted_data() == [0, 2, 4, 1, 3])

    def test_multiple_slices(self, main_col, table_types, engine):
        self._setup(main_col, table_types)

        if not self.mutable:
            return

        t = self.t
        t.add_index("a", engine=engine)

        for i in range(6, 51):
            t.add_row((1.0, "A", i))

        assert_col_equal(t["a"], list(range(1, 51)))
        assert np.all(t.indices[0].sorted_data() == list(range(50)))

        evens = t[::2]
        assert np.all(evens.indices[0].sorted_data() == list(range(25)))
        reverse = evens[::-1]
        index = reverse.indices[0]
        assert (index.start, index.stop, index.step) == (48, -2, -2)
        assert np.all(index.sorted_data() == list(range(24, -1, -1)))

        # modify slice of slice
        reverse[-10:] = 0
        expected = np.array(list(range(1, 51)))
        expected[:20][expected[:20] % 2 == 1] = 0
        assert_col_equal(t["a"], expected)
        assert_col_equal(evens["a"], expected[::2])
        assert_col_equal(reverse["a"], expected[::2][::-1])
        # first ten evens are now zero
        assert np.all(
            t.indices[0].sorted_data()
            == (
                [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 1, 3, 5, 7, 9, 11, 13, 15, 17, 19]
                + list(range(20, 50))
            )
        )
        assert np.all(evens.indices[0].sorted_data() == list(range(25)))
        assert np.all(reverse.indices[0].sorted_data() == list(range(24, -1, -1)))

        # try different step sizes of slice
        t2 = t[1:20:2]
        assert_col_equal(t2["a"], [2, 4, 6, 8, 10, 12, 14, 16, 18, 20])
        assert np.all(t2.indices[0].sorted_data() == list(range(10)))
        t3 = t2[::3]
        assert_col_equal(t3["a"], [2, 8, 14, 20])
        assert np.all(t3.indices[0].sorted_data() == [0, 1, 2, 3])
        t4 = t3[2::-1]
        assert_col_equal(t4["a"], [14, 8, 2])
        assert np.all(t4.indices[0].sorted_data() == [2, 1, 0])

    def test_sort(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        t = self.t[::-1]  # reverse table
        assert_col_equal(t["a"], [5, 4, 3, 2, 1])
        t.add_index("a", engine=engine)
        assert np.all(t.indices[0].sorted_data() == [4, 3, 2, 1, 0])

        if not self.mutable:
            return

        # sort table by column a
        t2 = t.copy()
        t2.sort("a")
        assert_col_equal(t2["a"], [1, 2, 3, 4, 5])
        assert np.all(t2.indices[0].sorted_data() == [0, 1, 2, 3, 4])

        # sort table by primary key
        t2 = t.copy()
        t2.sort()
        assert_col_equal(t2["a"], [1, 2, 3, 4, 5])
        assert np.all(t2.indices[0].sorted_data() == [0, 1, 2, 3, 4])

    def test_insert_row(self, main_col, table_types, engine):
        self._setup(main_col, table_types)

        if not self.mutable:
            return

        t = self.t
        t.add_index("a", engine=engine)
        t.insert_row(2, (1.0, "12", 6))
        assert_col_equal(t["a"], [1, 2, 6, 3, 4, 5])
        assert np.all(t.indices[0].sorted_data() == [0, 1, 3, 4, 5, 2])
        t.insert_row(1, (4.0, "13", 0))
        assert_col_equal(t["a"], [1, 0, 2, 6, 3, 4, 5])
        assert np.all(t.indices[0].sorted_data() == [1, 0, 2, 4, 5, 6, 3])

    def test_index_modes(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        t = self.t
        t.add_index("a", engine=engine)

        # first, no special mode
        assert len(t[[1, 3]].indices) == 1
        assert len(t[::-1].indices) == 1
        assert len(self._table_type(t).indices) == 1
        assert np.all(t.indices[0].sorted_data() == [0, 1, 2, 3, 4])
        t2 = t.copy()

        # non-copy mode
        with t.index_mode("discard_on_copy"):
            assert len(t[[1, 3]].indices) == 0
            assert len(t[::-1].indices) == 0
            assert len(self._table_type(t).indices) == 0
            assert len(t2.copy().indices) == 1  # mode should only affect t

        # make sure non-copy mode is exited correctly
        assert len(t[[1, 3]].indices) == 1

        if not self.mutable:
            return

        # non-modify mode
        with t.index_mode("freeze"):
            assert np.all(t.indices[0].sorted_data() == [0, 1, 2, 3, 4])
            t["a"][0] = 6
            assert np.all(t.indices[0].sorted_data() == [0, 1, 2, 3, 4])
            t.add_row((1.5, "12", 2))
            assert np.all(t.indices[0].sorted_data() == [0, 1, 2, 3, 4])
            t.remove_rows([1, 3])
            assert np.all(t.indices[0].sorted_data() == [0, 1, 2, 3, 4])
            assert_col_equal(t["a"], [6, 3, 5, 2])
            # mode should only affect t
            assert np.all(t2.indices[0].sorted_data() == [0, 1, 2, 3, 4])
            t2["a"][0] = 6
            assert np.all(t2.indices[0].sorted_data() == [1, 2, 3, 4, 0])

        # make sure non-modify mode is exited correctly
        assert np.all(t.indices[0].sorted_data() == [3, 1, 2, 0])

        if isinstance(t["a"], BaseColumn):
            assert len(t["a"][::-1].info.indices) == 0
            with t.index_mode("copy_on_getitem"):
                assert len(t["a"][[1, 2]].info.indices) == 1
                # mode should only affect t
                assert len(t2["a"][[1, 2]].info.indices) == 0

            assert len(t["a"][::-1].info.indices) == 0
            assert len(t2["a"][::-1].info.indices) == 0

    def test_index_retrieval(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        t = self.t
        t.add_index("a", engine=engine)
        t.add_index(["a", "c"], engine=engine)
        assert len(t.indices) == 2
        assert len(t.indices["a"].columns) == 1
        assert len(t.indices["a", "c"].columns) == 2

        with pytest.raises(IndexError):
            t.indices["b"]

    def test_col_rename(self, main_col, table_types, engine):
        """
        Checks for a previous bug in which copying a Table
        with different column names raised an exception.
        """
        self._setup(main_col, table_types)
        t = self.t
        t.add_index("a", engine=engine)
        t2 = self._table_type(self.t, names=["d", "e", "f"])
        assert len(t2.indices) == 1

    def test_table_loc(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        t = self.t

        t.add_index("a", engine=engine)
        t.add_index("b", engine=engine)

        t2 = t.loc[self.make_val(3)]  # single label, with primary key 'a'
        assert_col_equal(t2["a"], [3])
        assert isinstance(t2, Row)

        # list search
        t2 = t.loc[[self.make_val(1), self.make_val(4), self.make_val(2)]]
        assert_col_equal(t2["a"], [1, 4, 2])  # same order as input list
        if not isinstance(main_col, Time):
            # ndarray search
            t2 = t.loc[np.array([1, 4, 2])]
            assert_col_equal(t2["a"], [1, 4, 2])
        assert_col_equal(t2["a"], [1, 4, 2])
        t2 = t.loc[self.make_val(3) : self.make_val(5)]  # range search
        assert_col_equal(t2["a"], [3, 4, 5])
        t2 = t.loc.with_index("b")[5.0:7.0]
        assert_col_equal(t2["b"], [5.1, 6.2, 7.0])
        # search by sorted index
        t2 = t.iloc[0:2]  # two smallest rows by column 'a'
        assert_col_equal(t2["a"], [1, 2])
        t2 = t.iloc.with_index("b")[2:]  # exclude two smallest rows in column 'b'
        assert_col_equal(t2["b"], [5.1, 6.2, 7.0])

        for t2 in (t.loc[:], t.iloc[:]):
            assert_col_equal(t2["a"], [1, 2, 3, 4, 5])

    def test_table_loc_indices(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        t = self.t

        t.add_index("a", engine=engine)
        t.add_index("b", engine=engine)

        t2 = t.loc_indices[self.make_val(3)]  # single label, with primary key 'a'
        assert t2 == 2

        # list search
        t2 = t.loc_indices[[self.make_val(1), self.make_val(4), self.make_val(2)]]
        for i, p in zip(t2, [1, 4, 2]):  # same order as input list
            assert i == p - 1

    def test_invalid_search(self, main_col, table_types, engine):
        # using .loc and .loc_indices with a value not present should raise an exception
        self._setup(main_col, table_types)
        t = self.t

        t.add_index("a")
        with pytest.raises(KeyError):
            t.loc[self.make_val(6)]
        with pytest.raises(KeyError):
            t.loc_indices[self.make_val(6)]

    def test_copy_index_references(self, main_col, table_types, engine):
        # check against a bug in which indices were given an incorrect
        # column reference when copied
        self._setup(main_col, table_types)
        t = self.t

        t.add_index("a")
        t.add_index("b")
        t2 = t.copy()
        assert t2.indices["a"].columns[0] is t2["a"]
        assert t2.indices["b"].columns[0] is t2["b"]

    def test_unique_index(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        t = self.t

        t.add_index("a", engine=engine, unique=True)
        assert np.all(t.indices["a"].sorted_data() == [0, 1, 2, 3, 4])

        if self.mutable:
            with pytest.raises(ValueError):
                t.add_row((5.0, "9", 5))

    def test_copy_indexed_table(self, table_types):
        self._setup(_col, table_types)
        t = self.t
        t.add_index("a")
        t.add_index(["a", "b"])
        for tp in (self._table_type(t), t.copy()):
            assert len(t.indices) == len(tp.indices)
            for index, indexp in zip(t.indices, tp.indices):
                assert np.all(index.data.data == indexp.data.data)
                assert index.data.data.colnames == indexp.data.data.colnames

    def test_updating_row_byindex(self, main_col, table_types, engine):
        self._setup(main_col, table_types)
        t = Table(
            [["a", "b", "c", "d"], [2, 3, 4, 5], [3, 4, 5, 6]],
            names=("a", "b", "c"),
            meta={"name": "first table"},
        )

        t.add_index("a", engine=engine)
        t.add_index("b", engine=engine)

        t.loc["c"] = ["g", 40, 50]  # single label, with primary key 'a'
        t2 = t[2]
        assert list(t2) == ["g", 40, 50]

        # list search
        t.loc[["a", "d", "b"]] = [["a", 20, 30], ["d", 50, 60], ["b", 30, 40]]
        t2 = [["a", 20, 30], ["d", 50, 60], ["b", 30, 40]]
        for i, p in zip(t2, [1, 4, 2]):  # same order as input list
            assert list(t[p - 1]) == i

    def test_invalid_updates(self, main_col, table_types, engine):
        # using .loc and .loc_indices with a value not present should raise an exception
        self._setup(main_col, table_types)
        t = Table(
            [[1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6]],
            names=("a", "b", "c"),
            meta={"name": "first table"},
        )

        t.add_index("a")
        with pytest.raises(ValueError):
            t.loc[3] = [[1, 2, 3]]
        with pytest.raises(ValueError):
            t.loc[[1, 4, 2]] = [[1, 2, 3], [4, 5, 6]]
        with pytest.raises(ValueError):
            t.loc[[1, 4, 2]] = [[1, 2, 3], [4, 5, 6], [2, 3]]
        with pytest.raises(ValueError):
            t.loc[[1, 4, 2]] = [[1, 2, 3], [4, 5], [2, 3]]


def test_get_index():
    a = [1, 4, 5, 2, 7, 4, 45]
    b = [2.0, 5.0, 8.2, 3.7, 4.3, 6.5, 3.3]
    t = Table([a, b], names=("a", "b"), meta={"name": "first table"})
    t.add_index(["a"])
    # Getting the values of index using names
    x1 = get_index(t, names=["a"])

    assert isinstance(x1, SlicedIndex)
    assert len(x1.columns) == 1
    assert len(x1.columns[0]) == 7
    assert x1.columns[0].info.name == "a"
    # Getting the vales of index using table_copy
    x2 = get_index(t, table_copy=t[["a"]])

    assert isinstance(x2, SlicedIndex)
    assert len(x2.columns) == 1
    assert len(x2.columns[0]) == 7
    assert x2.columns[0].info.name == "a"

    with pytest.raises(ValueError):
        get_index(t, names=["a"], table_copy=t[["a"]])
    with pytest.raises(ValueError):
        get_index(t, names=None, table_copy=None)


@pytest.mark.parametrize("table_type", [Table, QTable])
def test_index_loc_with_quantity(engine, table_type):
    t = table_type()
    t["a"] = [3, 1, 2] * u.m
    t["b"] = [1, 2, 3]
    t.add_index("a", engine=engine)

    unit = u.m if table_type is QTable else 1
    assert tuple(t.loc[1 * unit]) == (1 * unit, 2)
    assert np.all(t.loc_indices[[1 * unit, 3 * unit]] == [1, 0])
    assert tuple(t.iloc[1]) == (2 * unit, 3)
    for loc in (t.loc, t.iloc):
        t_loc = loc[:]
        assert len(t_loc) == 3
        assert np.all(t_loc["a"] == [1, 2, 3] * unit)
        assert np.all(t_loc["b"] == [2, 3, 1])


def test_index_loc_with_string(engine):
    t = Table()
    t["a"] = ["z", "a", "m"]
    t["b"] = [1, 2, 3]
    t.add_index("a", engine=engine)

    assert tuple(t.loc["a"]) == ("a", 2)
    assert np.all(t.loc_indices[["a", "z"]] == [1, 0])
    assert tuple(t.iloc[1]) == ("m", 3)
    for loc in (t.loc, t.iloc):
        t_loc = loc[:]
        assert len(t_loc) == 3
        assert np.all(t_loc["a"] == ["a", "m", "z"])
        assert np.all(t_loc["b"] == [2, 3, 1])


def test_table_index_time_warning(engine):
    # Make sure that no ERFA warnings are emitted when indexing a table by
    # a Time column with a non-default time scale
    tab = Table()
    tab["a"] = Time([1, 2, 3], format="jyear", scale="tai")
    tab["b"] = [4, 3, 2]
    with warnings.catch_warnings(record=True) as wlist:
        tab.add_index(("a", "b"), engine=engine)
    assert len(wlist) == 0


@pytest.mark.parametrize(
    "col",
    [
        Column(np.arange(50000, 50005)),
        np.arange(50000, 50005) * u.m,
        Time(np.arange(50000, 50005), format="mjd"),
    ],
)
def test_table_index_does_not_propagate_to_column_slices(col):
    # They lost contact to the parent table, so they should also not have
    # information on the indices; this helps prevent large memory usage if,
    # e.g., a large time column is turned into an object array; see gh-10688.
    tab = QTable()
    tab["t"] = col
    tab.add_index("t")
    t = tab["t"]
    assert t.info.indices
    tx = t[1:]
    assert not tx.info.indices
    tabx = tab[1:]
    t = tabx["t"]
    assert t.info.indices


def test_hstack_qtable_table():
    # Check in particular that indices are initialized or copied correctly
    # for a Column that is being converted to a Quantity.
    qtab = QTable([np.arange(5.0) * u.m], names=["s"])
    qtab.add_index("s")
    tab = Table([Column(np.arange(5.0), unit=u.s)], names=["t"])
    qstack = hstack([qtab, tab])
    assert qstack["t"].info.indices == []
    assert qstack.indices == []


def test_index_slice_exception():
    with pytest.raises(TypeError, match="index_slice must be tuple or slice"):
        SlicedIndex(None, None)


@pytest.fixture(scope="module")
def simple_table():
    """Simple table with an index on column 'a'."""
    t = Table()
    t["a"] = [3, 1, 2, 3]
    t["b"] = ["x", "y", "z", "w"]
    t.add_index("a")
    return t


@pytest.mark.parametrize("key", [None, "a"])
@pytest.mark.parametrize(
    "item,length,cls",
    [
        (slice(0, 0), 0, Table),
        ([], 0, Table),
        ([1], 1, Table),
        ([1, 3], 3, Table),
        (np.array([]), 0, Table),
        (np.array([1]), 1, Table),
        (3, 2, Table),  # scalar index with multiple rows
        (1, None, Row),  # scalar index with single row
    ],
)
def test_index_zero_slice_or_sequence_or_scalar(simple_table, key, item, length, cls):
    """Test that indexing with various types gives the expected result.

    Tests fix for #18037.
    """
    loc = simple_table.loc.with_index(key) if key is not None else simple_table.loc
    tloc = loc[item]
    assert isinstance(tloc, cls)
    assert tloc.colnames == simple_table.colnames

    rows = simple_table.loc_indices[item]
    if cls is Table:
        assert len(tloc) == length
        assert len(rows) == length


@pytest.mark.parametrize(
    "method,item",
    [
        ("loc", (2, 5)),
        ("iloc", 1),
        ("loc_indices", (2, 5)),
    ],
)
def test_index_id_item_deprecation_and_with_index(method, item):
    """t.loc/iloc/loc_indices[index_id, item] raises a deprecation warning.

    Also test that these methods
    """
    t = Table()
    t["a"] = [1, 2, 3]
    t["b"] = [4, 5, 6]
    t["c"] = ["x", "y", "z"]
    index_id = ("a", "b")
    t.add_index(index_id)
    prop = getattr(t, method)
    # Test calling like t.loc.with_index("a", "b") and t.loc.with_index(("a", "b")).
    out_call_1 = prop.with_index(*index_id)[item]
    out_call_2 = prop.with_index(index_id)[item]
    with pytest.warns(
        AstropyDeprecationWarning,
        match=r"Calling `Table.loc/iloc/loc_indices\[index_id, item\]`",
    ):
        out_depr = prop[index_id, item]
    assert type(out_depr) is type(out_call_1)
    assert out_depr == out_call_1
    assert type(out_call_1) is type(out_call_2)
    assert out_call_1 == out_call_2


def test_engine_type_error():
    t = Table()
    t["a"] = [1, 2]
    t["b"] = [3, 4]
    with pytest.raises(
        TypeError,
        match=r"engine must be an Engine class or instance, got 'b' instead.",
    ):
        t.add_index("a", "b")  # Easy mistake, too bad engine= is not keyword-only


@pytest.mark.parametrize(
    "masked",
    [pytest.param(False, id="raw-array"), pytest.param(True, id="masked array")],
)
def test_nd_columun_as_index(masked):
    # see https://github.com/astropy/astropy/issues/13292
    # and https://github.com/astropy/astropy/pull/16360
    t = Table()
    data = np.arange(0, 6)
    if masked:
        data = np.ma.masked_inside(data, 2, 4)
    t.add_column(data.reshape(3, -1), name="arr")
    with pytest.raises(
        ValueError, match="Multi-dimensional column 'arr' cannot be used as an index."
    ):
        t.add_index("arr")


def test_indices_read_unknown_engine():
    lines = [
        "# %ECSV 1.0",
        "# ---",
        "# datatype:",
        f"# - {{name: a, datatype: {NATIVE_INT_NAME}}}",
        f"# - {{name: __index__, datatype: {NATIVE_INT_NAME}}}",
        "# meta: !!omap",
        "# - __table_indices__:",
        "#     indices:",
        "#     - colnames: [a]",
        "#       engine: Foo",
        "#       index_colname: __index__",
        "#       unique: true",
        "#     primary_key: [a]",
        "# schema: astropy-2.0",
        "a __index__",
        "1 0",
        "3 2",
        "2 1",
    ]
    text = "\n".join(lines)

    with pytest.warns(
        AstropyWarning,
        match=r"Unknown index engine 'Foo', creating index using SortedArray engine",
    ):
        t = Table.read(text, format="ecsv")
    # a==3 at row 1
    assert t.loc_indices[1] == 0
    assert t.loc_indices[3] == 1
    assert t.loc_indices[2] == 2


def test_indices_serialization_unique_representation():
    t = Table()
    t["a"] = [1, 3, 2]
    t.add_index("a", unique=True)
    out = io.StringIO()
    t.write(out, format="ecsv", write_indices=True)
    assert out.getvalue().splitlines() == [
        "# %ECSV 1.0",
        "# ---",
        "# datatype:",
        f"# - {{name: a, datatype: {NATIVE_INT_NAME}}}",
        f"# - {{name: __index__, datatype: {NATIVE_INT_NAME}}}",
        "# meta: !!omap",
        "# - __table_indices__:",
        "#     indices:",
        "#     - colnames: [a]",
        "#       index_colname: __index__",
        "#       unique: true",
        "#     primary_key: [a]",
        "# schema: astropy-2.0",
        "a __index__",
        "1 0",
        "3 2",
        "2 1",
    ]
    t2 = Table.read(out.getvalue(), format="ecsv")
    assert t2.indices[0].data.unique is True


@pytest.mark.parametrize("engine", [SortedArray, SCEngine])
def test_indices_serialization_representation_single(engine):
    """Add explicit test of serialization representation for single-index case.

    The `primary` key is not included in this case.
    """
    t = Table()
    t["a"] = [1, 3, 2]
    t.add_index("a", engine=engine)
    out = io.StringIO()
    t.write(out, format="ecsv", write_indices=True)
    exp = [
        "# %ECSV 1.0",
        "# ---",
        "# datatype:",
        f"# - {{name: a, datatype: {NATIVE_INT_NAME}}}",
        f"# - {{name: __index__, datatype: {NATIVE_INT_NAME}}}",
        "# meta: !!omap",
        "# - __table_indices__:",
        "#     indices:",
        "#     - colnames: [a]",
        "#       index_colname: __index__",
        "#     primary_key: [a]",
        "# schema: astropy-2.0",
        "a __index__",
        "1 0",
        "3 2",
        "2 1",
    ]

    if engine is SCEngine:
        exp.insert(9, "#       engine: SCEngine")

    assert out.getvalue().splitlines() == exp


def test_indices_serialization_representation_multiple():
    """Add explicit test of serialization representation for single-index case.

    This includes the `primary` key and a collision.
    """
    t = Table()
    t["a"] = [1, 3, 2]
    t["__index__1"] = [5, 4, 3]
    t.add_index(["a", "__index__1"])
    t.add_index("a")
    out = io.StringIO()
    t.write(out, format="ecsv", write_indices=True)

    exp = [
        "# %ECSV 1.0",
        "# ---",
        "# datatype:",
        f"# - {{name: a, datatype: {NATIVE_INT_NAME}}}",
        f"# - {{name: __index__1, datatype: {NATIVE_INT_NAME}}}",
        f"# - {{name: __index__, datatype: {NATIVE_INT_NAME}}}",
        f"# - {{name: __index__2, datatype: {NATIVE_INT_NAME}}}",
        "# meta: !!omap",
        "# - __table_indices__:",
        "#     indices:",
        "#     - colnames: [a, __index__1]",
        "#       index_colname: __index__",
        "#     - colnames: [a]",
        "#       index_colname: __index__2",
        "#     primary_key: [a, __index__1]",
        "# schema: astropy-2.0",
        "a __index__1 __index__ __index__2",
        "1 5 0 0",
        "3 4 2 2",
        "2 3 1 1",
    ]
    assert out.getvalue().splitlines() == exp


@pytest.mark.parametrize("dtype", [np.int16, np.float32, np.int64, np.float64])
def test_indices_roundtrip_various_dtypes(dtype):
    """Test that serialization round-trip works for various index dtypes."""
    t = Table()
    t["a"] = np.array([1, 3, 2], dtype=dtype)
    t["b"] = np.array([5, 6, 7], dtype=dtype)
    t.add_index("a")
    t.add_index(["a", "b"])
    out = io.StringIO()
    t.write(out, format="ecsv", write_indices=True)
    t2 = Table.read(out.getvalue(), format="ecsv")

    assert_tables_equal(t, t2)
    assert_indices_equal(t, t2, [("a",), ("a", "b")])


@pytest.mark.parametrize("single_index", [True, False])
@pytest.mark.parametrize("engine", [SortedArray, SCEngine])
@pytest.mark.parametrize("fmt", ["fits", "ecsv", "hdf5"])
def test_indices_roundtrip_through_file(single_index, fmt, engine, tmp_path):
    if single_index and fmt != "ecsv":
        # Save a few compute cycles, since single_index is really impacting just the
        # serialization data and the engine and fmt don't matter.
        pytest.skip()

    if not HAS_H5PY and fmt == "hdf5":
        pytest.skip("hdf5 tests require h5py")

    t = QTable()
    t["a"] = Time([1, 3, 2, 2], format="cxcsec")
    t["b"] = [3, 2, 2, 1]
    t["__index__"] = [3, 1, 4, 2]  # Force a collision
    indices_colnames = [
        ["a"],
        ["b", "a"],
        ["a", "b", "__index__"],
        ["__index__"],
    ]
    if single_index:
        indices_colnames = indices_colnames[:1]

    for colnames in indices_colnames:
        t.add_index(colnames, engine=engine)

    path = tmp_path / f"out.{fmt}"
    kwargs = {"serialize_meta": True, "path": "root"} if fmt == "hdf5" else {}
    t.write(path, format=fmt, write_indices=True, **kwargs)

    kwargs = {"astropy_native": True} if fmt == "fits" else {}
    t2 = QTable.read(path, format=fmt, **kwargs)
    if fmt == "fits":
        # FITS does not round-trip the format
        t2["a"].format = "cxcsec"

    assert_tables_equal(t, t2)
    assert_indices_equal(t, t2, indices_colnames)


def assert_indices_equal(t, t2, indices_colnames):
    assert len(t.indices) == len(t2.indices)
    assert t.primary_key == t2.primary_key

    for colnames in indices_colnames:
        index = t.indices[colnames]
        index2 = t2.indices[colnames]
        assert index.id == index2.id
        # Table rows sorted in index order
        assert_tables_equal(
            t.iloc.with_index(colnames)[:], t2.iloc.with_index(colnames)[:]
        )
        # Check that the engine row_index column/list is identical
        assert np.all(index.data.sorted_data() == index2.data.sorted_data())
        # Check engine items as a list of pairs of the form
        # [(key, [row 1, row 2, ...]), ...].
        assert index.data.items() == index2.data.items()

        key0 = tuple(t.iloc.with_index(colnames)[0][colnames])
        key1 = tuple(t.iloc.with_index(colnames)[-1][colnames])
        assert t.loc.with_index(colnames)[key0] == t2.loc.with_index(colnames)[key0]
        assert t.loc.with_index(colnames)[key1] == t2.loc.with_index(colnames)[key1]
        assert (
            t.loc_indices.with_index(colnames)[key0]
            == t2.loc_indices.with_index(colnames)[key0]
        )
        assert (
            t.loc_indices.with_index(colnames)[key1]
            == t2.loc_indices.with_index(colnames)[key1]
        )


@pytest.mark.parametrize("index_first", [True, False])
def test_slice_an_indexed_table(index_first):
    """Test slicing a table that is already indexed.

    Test of fix for https://github.com/astropy/astropy/issues/10732.

    #10732 is the case index_first=True, but also test slicing first (index_first=False)
    since we're at it.
    """
    t = Table()
    t["a"] = [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
    t["b"] = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
    t["c"] = ["e", "f", "g", "h", "i", "j", "k", "a", "b", "c"]

    if index_first:
        t.add_index("a")
        t.add_index(["b", "c"])
        ts = t[::2]
    else:
        ts = t[::2]
        ts.add_index("a")
        ts.add_index(["b", "c"])

    assert ts.pformat() == [
        " a   b   c ",
        "--- --- ---",
        "  9   0   e",
        "  7   0   g",
        "  5   0   i",
        "  3   1   k",
        "  1   1   b",
    ]
    # Index access works
    assert str(ts.loc[5]).splitlines() == [
        " a   b   c ",
        "--- --- ---",
        "  5   0   i",
    ]

    # Remove row 2 (a==5), check index access still works
    ts.remove_row(2)
    assert ts.pformat() == [
        " a   b   c ",
        "--- --- ---",
        "  9   0   e",
        "  7   0   g",
        "  3   1   k",
        "  1   1   b",
    ]
    assert str(ts.loc[1]).splitlines() == [
        " a   b   c ",
        "--- --- ---",
        "  1   1   b",
    ]

    # Remove row 2 (now a==3), check index access still works
    ts.remove_row(2)
    assert ts.pformat() == [
        " a   b   c ",
        "--- --- ---",
        "  9   0   e",
        "  7   0   g",
        "  1   1   b",
    ]
    assert str(ts.loc[7]).splitlines() == [
        " a   b   c ",
        "--- --- ---",
        "  7   0   g",
    ]

    # Make sure primary index and secondary index look right (with original=True)
    assert str(ts.indices[0]).splitlines() == [
        "<SlicedIndex original=True index=<Index columns=('a',) data=<SortedArray length=3>",
        " a  rows",
        "--- ----",
        "  1    2",
        "  7    1",
        "  9    0>>",
    ]
    assert str(ts.indices[1]).splitlines() == [
        "<SlicedIndex original=True index=<Index columns=('b', 'c') data=<SortedArray length=3>",
        " b   c  rows",
        "--- --- ----",
        "  0   e    0",
        "  0   g    1",
        "  1   b    2>>",
    ]


def test_unique_indices_after_multicol_index_slice():
    """Test that table indices after slicing are correct.

    This tests code in Table._new_from_slice() that ensures uniqueness of table index
    objects when slicing (via slice, ndarray, list etc) a table with a multi-column
    index.
    """
    t = Table()
    t["a"] = [2, 3]
    t["b"] = [3, 5]
    t.add_index(["a", "b"])
    t2 = t[:1]
    assert len(t2.indices) == 1  # without fix would be 2, both with id ("a", "b").
    assert t2.indices[0].id == ("a", "b")


def test_index_not_corrupted_on_failed_row_assignment(engine):
    """Regression test: index must survive a failed row assignment.

    When ``table[row] = values`` raises because one of the values is
    incompatible with its column dtype, the table index was left in an
    inconsistent state.  Specifically, ``Index.replace`` removed the
    existing key from the sorted array *before* trying to insert the new
    one; if the insert failed the old key was permanently gone even though
    the column data was never changed.

    After the fix the index must round-trip correctly: the original key is
    still findable and no ghost key is present.  The test is run for all
    three available index engines (BST, SortedArray, SCEngine).
    """
    t = Table({"x": [1, 2, 3], "y": [4, 5, 6]})
    t.add_index("y", engine=engine)

    with pytest.raises(ValueError):
        # "bad" is not convertible to the int64 dtype of column y
        t[0] = (99, "bad")

    # Data must be unchanged
    assert t[0]["x"] == 1
    assert t[0]["y"] == 4

    # Index must still find the original key
    result = t.loc[4]
    assert result["x"] == 1
    assert result["y"] == 4

    # No ghost entry for the attempted new value
    with pytest.raises(KeyError):
        t.loc[99]
