import numpy as np
import pytest

from pandas.core.dtypes.common import is_categorical_dtype
from pandas.core.dtypes.dtypes import CategoricalDtype

import pandas as pd
from pandas import (
    Categorical,
    CategoricalIndex,
    DataFrame,
    Index,
    Interval,
    Series,
    Timestamp,
)
from pandas.api.types import CategoricalDtype as CDT
from pandas.util import testing as tm
from pandas.util.testing import assert_frame_equal, assert_series_equal


class TestCategoricalIndex:
    def setup_method(self, method):

        self.df = DataFrame(
            {
                "A": np.arange(6, dtype="int64"),
                "B": Series(list("aabbca")).astype(CDT(list("cab"))),
            }
        ).set_index("B")
        self.df2 = DataFrame(
            {
                "A": np.arange(6, dtype="int64"),
                "B": Series(list("aabbca")).astype(CDT(list("cabe"))),
            }
        ).set_index("B")
        self.df3 = DataFrame(
            {
                "A": np.arange(6, dtype="int64"),
                "B": (Series([1, 1, 2, 1, 3, 2]).astype(CDT([3, 2, 1], ordered=True))),
            }
        ).set_index("B")
        self.df4 = DataFrame(
            {
                "A": np.arange(6, dtype="int64"),
                "B": (Series([1, 1, 2, 1, 3, 2]).astype(CDT([3, 2, 1], ordered=False))),
            }
        ).set_index("B")

    def test_loc_scalar(self):
        result = self.df.loc["a"]
        expected = DataFrame(
            {"A": [0, 1, 5], "B": (Series(list("aaa")).astype(CDT(list("cab"))))}
        ).set_index("B")
        assert_frame_equal(result, expected)

        df = self.df.copy()
        df.loc["a"] = 20
        expected = DataFrame(
            {
                "A": [20, 20, 2, 3, 4, 20],
                "B": (Series(list("aabbca")).astype(CDT(list("cab")))),
            }
        ).set_index("B")
        assert_frame_equal(df, expected)

        # value not in the categories
        with pytest.raises(KeyError, match=r"^'d'$"):
            df.loc["d"]

        msg = "cannot append a non-category item to a CategoricalIndex"
        with pytest.raises(TypeError, match=msg):
            df.loc["d"] = 10

        msg = (
            "cannot insert an item into a CategoricalIndex that is not"
            " already an existing category"
        )
        with pytest.raises(TypeError, match=msg):
            df.loc["d", "A"] = 10
        with pytest.raises(TypeError, match=msg):
            df.loc["d", "C"] = 10

    def test_getitem_scalar(self):

        cats = Categorical([Timestamp("12-31-1999"), Timestamp("12-31-2000")])

        s = Series([1, 2], index=cats)

        expected = s.iloc[0]
        result = s[cats[0]]
        assert result == expected

    def test_slicing_directly(self):
        cat = Categorical(["a", "b", "c", "d", "a", "b", "c"])
        sliced = cat[3]
        assert sliced == "d"
        sliced = cat[3:5]
        expected = Categorical(["d", "a"], categories=["a", "b", "c", "d"])
        tm.assert_numpy_array_equal(sliced._codes, expected._codes)
        tm.assert_index_equal(sliced.categories, expected.categories)

    def test_slicing(self):
        cat = Series(Categorical([1, 2, 3, 4]))
        reversed = cat[::-1]
        exp = np.array([4, 3, 2, 1], dtype=np.int64)
        tm.assert_numpy_array_equal(reversed.__array__(), exp)

        df = DataFrame({"value": (np.arange(100) + 1).astype("int64")})
        df["D"] = pd.cut(df.value, bins=[0, 25, 50, 75, 100])

        expected = Series([11, Interval(0, 25)], index=["value", "D"], name=10)
        result = df.iloc[10]