import numpy as np
import pytest

import pandas as pd
import pandas._testing as tm


def test_infer_objects(idx):
    with pytest.raises(NotImplementedError, match="to_frame"):
        idx.infer_objects()


def test_shift(idx):
    # GH8083 test the base class for shift
    msg = (
        "This method is only implemented for DatetimeIndex, PeriodIndex and "
        "TimedeltaIndex; Got type MultiIndex"
    )
    with pytest.raises(NotImplementedError, match=msg):
        idx.shift(1)
    with pytest.raises(NotImplementedError, match=msg):
        idx.shift(1, 2)


def test_groupby(idx):
    groups = idx.groupby(np.array([1, 1, 1, 2, 2, 2]))
    labels = idx.tolist()
    exp = {1: labels[:3], 2: labels[3:]}
    tm.assert_dict_equal(groups, exp)

    # GH5620
    groups = idx.groupby(idx)
    exp = {key: [key] for key in idx}
    tm.assert_dict_equal(groups, exp)


def test_truncate_multiindex():
    # GH 34564 for MultiIndex level names check
    major_axis = pd.Index(list(range(4)))
    minor_axis = pd.Index(list(range(2)))

    major_codes = np.array([0, 0, 1, 2, 3, 3])
    minor_codes = np.array([0, 1, 0, 1, 0, 1])

    index = pd.MultiIndex(
        levels=[major_axis, minor_axis],
        codes=[major_codes, minor_codes],
        names=["L1", "L2"],
    )

    result = index.truncate(before=1)
    assert "foo" not in result.levels[0]
    assert 1 in result.levels[0]
    assert index.names == result.names

    result = index.truncate(after=1)
    assert 2 not in result.levels[0]
    assert 1 in result.levels[0]
    assert index.names == result.names

    result = index.truncate(before=1, after=2)
    assert len(result.levels[0]) == 2
    assert index.names == result.names

    msg = "after < before"
    with pytest.raises(ValueError, match=msg):
        index.truncate(3, 1)


# TODO: reshape


def test_reorder_levels(idx):
    # this blows up
    with pytest.raises(IndexError, match="^Too many levels"):
        idx.reorder_levels([2, 1, 0])


def test_numpy_repeat():
    reps = 2
    numbers = [1, 2, 3]
    names = np.array(["foo", "bar"])

    m = pd.MultiIndex.from_product([numbers, names], names=names)
    expected = pd.MultiIndex.from_product([numbers, names.repeat(reps)], names=names)
    tm.assert_index_equal(np.repeat(m, reps), expected)

    msg = "the 'axis' parameter is not supported"
    with pytest.raises(ValueError, match=msg):
        np.repeat(m, reps, axis=1)


def test_append_mixed_dtypes():
    # GH 13660
    dti = pd.date_range("2011-01-01", freq="ME", periods=3)
    dti_tz = pd.date_range("2011-01-01", freq="ME", periods=3, tz="US/Eastern")
    pi = pd.period_range("2011-01", freq="M", periods=3)

    mi = pd.MultiIndex.from_arrays(
        [[1, 2, 3], [1.1, np.nan, 3.3], ["a", "b", "c"], dti, dti_tz, pi]
    )
    assert mi.nlevels == 6

    res = mi.append(mi)
    exp = pd.MultiIndex.from_arrays(
        [
            [1, 2, 3, 1, 2, 3],
            [1.1, np.nan, 3.3, 1.1, np.nan, 3.3],
            ["a", "b", "c", "a", "b", "c"],
            dti.append(dti),
            dti_tz.append(dti_tz),
            pi.append(pi),
        ]
    )
    tm.assert_index_equal(res, exp, check_freq=False)

    other = pd.MultiIndex.from_arrays(
        [
            ["x", "y", "z"],
            ["x", "y", "z"],
            ["x", "y", "z"],
            ["x", "y", "z"],
            ["x", "y", "z"],
            ["x", "y", "z"],
        ]
    )

    res = mi.append(other)
    exp = pd.MultiIndex.from_arrays(
        [
            [1, 2, 3, "x", "y", "z"],
            [1.1, np.nan, 3.3, "x", "y", "z"],
            ["a", "b", "c", "x", "y", "z"],
            dti.append(pd.Index(["x", "y", "z"])),
            dti_tz.append(pd.Index(["x", "y", "z"])),
            pi.append(pd.Index(["x", "y", "z"])),
        ]
    )
    tm.assert_index_equal(res, exp, check_freq=False)


def test_iter(idx):
    result = list(idx)
    expected = [
        ("foo", "one"),
        ("foo", "two"),
        ("bar", "one"),
        ("baz", "two"),
        ("qux", "one"),
        ("qux", "two"),
    ]
    assert result == expected


def test_sub(idx):
    first = idx

    # - now raises (previously was set op difference)
    msg = "cannot perform __sub__ with this index type: MultiIndex"
    with pytest.raises(TypeError, match=msg):
        first - idx[-3:]
    with pytest.raises(TypeError, match=msg):
        idx[-3:] - first
    with pytest.raises(TypeError, match=msg):
        idx[-3:] - first.tolist()
    msg = "cannot perform __rsub__ with this index type: MultiIndex"
    with pytest.raises(TypeError, match=msg):
        first.tolist() - idx[-3:]


def test_map(idx):
    # callable
    index = idx

    result = index.map(lambda x: x)
    tm.assert_index_equal(result, index)


@pytest.mark.parametrize(
    "mapper",
    [
        lambda values, idx: {i: e for e, i in zip(values, idx, strict=True)},
        lambda values, idx: pd.Series(values, idx),
    ],
)
def test_map_dictlike(idx, mapper):
    identity = mapper(idx.values, idx)

    # we don't infer to uint64 dtype for a dict
    if idx.dtype == np.uint64 and isinstance(identity, dict):
        expected = idx.astype("int64")
    else:
        expected = idx

    result = idx.map(identity)
    tm.assert_index_equal(result, expected)

    # empty mappable
    expected = pd.Index([np.nan] * len(idx))
    result = idx.map(mapper(expected, idx))
    tm.assert_index_equal(result, expected)


@pytest.mark.parametrize(
    "func",
    [
        np.exp,
        np.exp2,
        np.expm1,
        np.log,
        np.log2,
        np.log10,
        np.log1p,
        np.sqrt,
        np.sin,
        np.cos,
        np.tan,
        np.arcsin,
        np.arccos,
        np.arctan,
        np.sinh,
        np.cosh,
        np.tanh,
        np.arcsinh,
        np.arccosh,
        np.arctanh,
        np.deg2rad,
        np.rad2deg,
    ],
    ids=lambda func: func.__name__,
)
def test_numpy_ufuncs(idx, func):
    # test ufuncs of numpy. see:
    # https://numpy.org/doc/stable/reference/ufuncs.html

    expected_exception = TypeError
    msg = (
        "loop of ufunc does not support argument 0 of type tuple which "
        f"has no callable {func.__name__} method"
    )
    with pytest.raises(expected_exception, match=msg):
        func(idx)


@pytest.mark.parametrize(
    "func",
    [np.isfinite, np.isinf, np.isnan, np.signbit],
    ids=lambda func: func.__name__,
)
def test_numpy_type_funcs(idx, func):
    msg = (
        f"ufunc '{func.__name__}' not supported for the input types, and the inputs "
        "could not be safely coerced to any supported types according to "
        "the casting rule ''safe''"
    )
    with pytest.raises(TypeError, match=msg):
        func(idx)


def test_categorical_multiindex_preserved_after_arithmetic():
    # GH#42785 - CategoricalIndex levels should be preserved after
    # division when the divisor has fewer levels (alignment/broadcasting).
    # The bug specifically lost CategoricalIndex on leftmost levels.
    v0 = pd.CategoricalIndex([10, 20, 30], name="V0")
    v1 = pd.CategoricalIndex(["B", "A"], name="V1")
    v2 = pd.CategoricalIndex(["X", "Y"], name="V2")
    mi = pd.MultiIndex.from_product([v0, v1, v2])

    ser = pd.Series(range(12), index=mi)
    norm = ser.groupby(level=["V1", "V2"]).sum()
    result = ser.div(norm)

    for level_idx in range(result.index.nlevels):
        assert isinstance(result.index.levels[level_idx], pd.CategoricalIndex)


def test_map_multiindex_mixed_length_tuple_dict_keys():
    # GH#40115 mapping a MultiIndex with a dict whose keys are *mixed-length*
    # tuples must not raise InvalidIndexError: the mapper key index has to stay
    # a flat object Index rather than being tupleized into a MultiIndex
    mi = pd.MultiIndex.from_tuples([("a", "a1"), ("c", "c1"), ("f", "f1")])
    mapping = {("a",): "xA", ("c", "c1"): "C1", ("f",): "xF"}
    result = mi.map(mapping)
    expected = pd.Index([np.nan, "C1", np.nan])
    tm.assert_index_equal(result, expected)

    # a dict that matches none of the (2-tuple) entries yields all-NaN, not a
    # raise -- the minimal repro reduced by the maintainer
    result_none = mi.map({("a",): "xA", ("z", "z1"): "zz"})
    tm.assert_index_equal(result_none, pd.Index([np.nan, np.nan, np.nan]))