import numpy as np
import pytest

from pandas import DataFrame, MultiIndex, Series
from pandas.util import 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)