import numpy as np
import pytest

from pandas import DataFrame, MultiIndex, Series
import pandas._testing as tm


@pytest.fixture
def simple_multiindex_dataframe():
    """
    Factory function to create simple 3 x 3 dataframe with
    both columns and row MultiIndex using supplied data or
    random data by default.
    """

    def _simple_multiindex_dataframe(data=None):
        if data is None:
            data = np.random.randn(3, 3)
        return DataFrame(
            data, columns=[[2, 2, 4], [6, 8, 10]], index=[[4, 4, 8], [8, 10, 12]]
        )

    return _simple_multiindex_dataframe


@pytest.mark.parametrize(
    "indexer, expected",
    [
        (
            lambda df: df.iloc[0],
            lambda arr: Series(arr[0], index=[[2, 2, 4], [6, 8, 10]], name=(4, 8)),
        ),
        (
            lambda df: df.iloc[2],
            lambda arr: Series(arr[2], index=[[2, 2, 4], [6, 8, 10]], name=(8, 12)),
        ),
        (
            lambda df: df.iloc[:, 2],
            lambda arr: Series(arr[:, 2], index=[[4, 4, 8], [8, 10, 12]], name=(4, 10)),
        ),
    ],
)
def test_iloc_returns_series(indexer, expected, simple_multiindex_dataframe):
    arr = np.random.randn(3, 3)
    df = simple_multiindex_dataframe(arr)
    result = indexer(df)
    expected = expected(arr)
    tm.assert_series_equal(result, expected)


def test_iloc_returns_dataframe(simple_multiindex_dataframe):
    df = simple_multiindex_dataframe()
    result = df.iloc[[0, 1]]
    expected = df.xs(4, drop_level=False)
    tm.assert_frame_equal(result, expected)


def test_iloc_returns_scalar(simple_multiindex_dataframe):
    arr = np.random.randn(3, 3)
    df = simple_multiindex_dataframe(arr)
    result = df.iloc[2, 2]
    expected = arr[2, 2]
    assert result == expected


def test_iloc_getitem_multiple_items():
    # GH 5528
    tup = zip(*[["a", "a", "b", "b"], ["x", "y", "x", "y"]])
    index = MultiIndex.from_tuples(tup)
    df = DataFrame(np.random.randn(4, 4), index=index)
    result = df.iloc[[2, 3]]
    expected = df.xs("b", drop_level=False)
    tm.assert_frame_equal(result, expected)


def test_iloc_getitem_labels():
    # this is basically regular indexing
    arr = np.random.randn(4, 3)
    df = DataFrame(
        arr,
        columns=[["i", "i", "j"], ["A", "A", "B"]],
        index=[["i", "i", "j", "k"], ["X", "X", "Y", "Y"]],
    )
    result = df.iloc[2, 2]
    expected = arr[2, 2]
    assert result == expected


def test_frame_getitem_slice(multiindex_dataframe_random_data):
    df = multiindex_dataframe_random_data
    result = df.iloc[:4]
    expected = df[:4]
    tm.assert_frame_equal(result, expected)


def test_frame_setitem_slice(multiindex_dataframe_random_data):
    df = multiindex_dataframe_random_data
    df.iloc[:4] = 0

    assert (df.values[:4] == 0).all()
    assert (df.values[4:] != 0).all()


def test_indexing_ambiguity_bug_1678():
    # GH 1678
    columns = MultiIndex.from_tuples(
        [("Ohio", "Green"), ("Ohio", "Red"), ("Colorado", "Green")]
    )
    index = MultiIndex.from_tuples([("a", 1), ("a", 2), ("b", 1), ("b", 2)])

    df = DataFrame(np.arange(12).reshape((4, 3)), index=index, columns=columns)

    result = df.iloc[:, 1]
    expected = df.loc[:, ("Ohio", "Red")]
    tm.assert_series_equal(result, expected)


def test_iloc_integer_locations():
    # GH 13797
    data = [
        ["str00", "str01"],
        ["str10", "str11"],
        ["str20", "srt21"],
        ["str30", "str31"],
        ["str40", "str41"],
    ]

    index = MultiIndex.from_tuples(
        [("CC", "A"), ("CC", "B"), ("CC", "B"), ("BB", "a"), ("BB", "b")]
    )

    expected = DataFrame(data)
    df = DataFrame(data, index=index)

    result = DataFrame([[df.iloc[r, c] for c in range(2)] for r in range(5)])

    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize(
    "data, indexes, values, expected_k",
    [
        # test without indexer value in first level of MultiIndex
        ([[2, 22, 5], [2, 33, 6]], [0, -1, 1], [2, 3, 1], [7, 10]),
        # test like code sample 1 in the issue
        ([[1, 22, 555], [1, 33, 666]], [0, -1, 1], [200, 300, 100], [755, 1066]),
        # test like code sample 2 in the issue
        ([[1, 3, 7], [2, 4, 8]], [0, -1, 1], [10, 10, 1000], [17, 1018]),
        # test like code sample 3 in the issue
        ([[1, 11, 4], [2, 22, 5], [3, 33, 6]], [0, -1, 1], [4, 7, 10], [8, 15, 13]),
    ],
)
def test_iloc_setitem_int_multiindex_series(data, indexes, values, expected_k):
    # GH17148
    df = DataFrame(data=data, columns=["i", "j", "k"])
    df = df.set_index(["i", "j"])

    series = df.k.copy()
    for i, v in zip(indexes, values):
        series.iloc[i] += v

    df["k"] = expected_k
    expected = df.k
    tm.assert_series_equal(series, expected)


def test_getitem_iloc(multiindex_dataframe_random_data):
    df = multiindex_dataframe_random_data
    result = df.iloc[2]
    expected = df.xs(df.index[2])
    tm.assert_series_equal(result, expected)