from datetime import timedelta

import numpy as np
import pytest

import pandas as pd
from pandas import (
    Categorical,
    CategoricalIndex,
    Index,
    IntervalIndex,
    MultiIndex,
    date_range,
)
from pandas.core.indexes.base import InvalidIndexError
import pandas.util.testing as tm
from pandas.util.testing import assert_almost_equal


def test_slice_locs_partial(idx):
    sorted_idx, _ = idx.sortlevel(0)

    result = sorted_idx.slice_locs(("foo", "two"), ("qux", "one"))
    assert result == (1, 5)

    result = sorted_idx.slice_locs(None, ("qux", "one"))
    assert result == (0, 5)

    result = sorted_idx.slice_locs(("foo", "two"), None)
    assert result == (1, len(sorted_idx))

    result = sorted_idx.slice_locs("bar", "baz")
    assert result == (2, 4)


def test_slice_locs():
    df = tm.makeTimeDataFrame()
    stacked = df.stack()
    idx = stacked.index

    slob = slice(*idx.slice_locs(df.index[5], df.index[15]))
    sliced = stacked[slob]
    expected = df[5:16].stack()
    tm.assert_almost_equal(sliced.values, expected.values)

    slob = slice(
        *idx.slice_locs(
            df.index[5] + timedelta(seconds=30), df.index[15] - timedelta(seconds=30)
        )
    )
    sliced = stacked[slob]
    expected = df[6:15].stack()
    tm.assert_almost_equal(sliced.values, expected.values)


def test_slice_locs_with_type_mismatch():
    df = tm.makeTimeDataFrame()
    stacked = df.stack()
    idx = stacked.index
    with pytest.raises(TypeError, match="^Level type mismatch"):
        idx.slice_locs((1, 3))
    with pytest.raises(TypeError, match="^Level type mismatch"):
        idx.slice_locs(df.index[5] + timedelta(seconds=30), (5, 2))
    df = tm.makeCustomDataframe(5, 5)
    stacked = df.stack()
    idx = stacked.index
    with pytest.raises(TypeError, match="^Level type mismatch"):
        idx.slice_locs(timedelta(seconds=30))
    # TODO: Try creating a UnicodeDecodeError in exception message
    with pytest.raises(TypeError, match="^Level type mismatch"):
        idx.slice_locs(df.index[1], (16, "a"))


def test_slice_locs_not_sorted():
    index = MultiIndex(
        levels=[Index(np.arange(4)), Index(np.arange(4)), Index(np.arange(4))],
        codes=[
            np.array([0, 0, 1, 2, 2, 2, 3, 3]),
            np.array([0, 1, 0, 0, 0, 1, 0, 1]),
            np.array([1, 0, 1, 1, 0, 0, 1, 0]),
        ],
    )
    msg = "[Kk]ey length.*greater than MultiIndex lexsort depth"
    with pytest.raises(KeyError, match=msg):
        index.slice_locs((1, 0, 1), (2, 1, 0))

    # works
    sorted_index, _ = index.sortlevel(0)
    # should there be a test case here???
    sorted_index.slice_locs((1, 0, 1), (2, 1, 0))


def test_slice_locs_not_contained():
    # some searchsorted action

    index = MultiIndex(
        levels=[[0, 2, 4, 6], [0, 2, 4]],
        codes=[[0, 0, 0, 1, 1, 2, 3, 3, 3], [0, 1, 2, 1, 2, 2, 0, 1, 2]],
        sortorder=0,
    )

    result = index.slice_locs((1, 0), (5, 2))
    assert result == (3, 6)

    result = index.slice_locs(1, 5)
    assert result == (3, 6)

    result = index.slice_locs((2, 2), (5, 2))
    assert result == (3, 6)

    result = index.slice_locs(2, 5)
    assert result == (3, 6)

    result = index.slice_locs((1, 0), (6, 3))
    assert result == (3, 8)

    result = index.slice_locs(-1, 10)
    assert result == (0, len(index))


def test_putmask_with_wrong_mask(idx):
    # GH18368

    msg = "putmask: mask and data must be the same size"
    with pytest.raises(ValueError, match=msg):
        idx.putmask(np.ones(len(idx) + 1, np.bool), 1)

    with pytest.raises(ValueError, match=msg):
        idx.putmask(np.ones(len(idx) - 1, np.bool), 1)

    with pytest.raises(ValueError, match=msg):
        idx.putmask("foo", 1)


def test_get_indexer():
    major_axis = Index(np.arange(4))
    minor_axis = Index(np.arange(2))

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

    index = MultiIndex(
        levels=[major_axis, minor_axis], codes=[major_codes, minor_codes]
    )
    idx1 = index[:5]
    idx2 = index[[1, 3, 5]]

    r1 = idx1.get_indexer(idx2)
    assert_almost_equal(r1, np.array([1, 3, -1], dtype=np.intp))

    r1 = idx2.get_indexer(idx1, method="pad")
    e1 = np.array([-1, 0, 0, 1, 1], dtype=np.intp)
    assert_almost_equal(r1, e1)

    r2 = idx2.get_indexer(idx1[::-1], method="pad")
    assert_almost_equal(r2, e1[::-1])

    rffill1 = idx2.get_indexer(idx1, method="ffill")
    assert_almost_equal(r1, rffill1)

    r1 = idx2.get_indexer(idx1, method="backfill")
    e1 = np.array([0, 0, 1, 1, 2], dtype=np.intp)
    assert_almost_equal(r1, e1)

    r2 = idx2.get_indexer(idx1[::-1], method="backfill")
    assert_almost_equal(r2, e1[::-1])

    rbfill1 = idx2.get_indexer(idx1, method="bfill")
    assert_almost_equal(r1, rbfill1)

    # pass non-MultiIndex
    r1 = idx1.get_indexer(idx2.values)
    rexp1 = idx1.get_indexer(idx2)
    assert_almost_equal(r1, rexp1)

    r1 = idx1.get_indexer([1, 2, 3])
    assert (r1 == [-1, -1, -1]).all()

    # create index with duplicates
    idx1 = Index(list(range(10)) + list(range(10)))
    idx2 = Index(list(range(20)))

    msg = "Reindexing only valid with uniquely valued Index objects"
    with pytest.raises(InvalidIndexError, match=msg):
        idx1.get_indexer(idx2)


def test_get_indexer_nearest():
    midx = MultiIndex.from_tuples([("a", 1), ("b", 2)])
    msg = "method='nearest' not implemented yet for MultiIndex; see GitHub issue 9365"
    with pytest.raises(NotImplementedError, match=msg):
        midx.get_indexer(["a"], method="nearest")
    msg = "tolerance not implemented yet for MultiIndex"
    with pytest.raises(NotImplementedError, match=msg):
        midx.get_indexer(["a"], method="pad", tolerance=2)


def test_getitem(idx):
    # scalar
    assert idx[2] == ("bar", "one")

    # slice
    result = idx[2:5]
    expected = idx[[2, 3, 4]]
    assert result.equals(expected)