import numpy as np
import pytest

import pandas as pd
import pandas._testing as tm
from pandas.core import algorithms
from pandas.core.arrays import PeriodArray


class TestSeriesIsIn:
    def test_isin(self):
        s = pd.Series(["A", "B", "C", "a", "B", "B", "A", "C"])

        result = s.isin(["A", "C"])
        expected = pd.Series([True, False, True, False, False, False, True, True])
        tm.assert_series_equal(result, expected)

        # GH#16012
        # This specific issue has to have a series over 1e6 in len, but the
        # comparison array (in_list) must be large enough so that numpy doesn't
        # do a manual masking trick that will avoid this issue altogether
        s = pd.Series(list("abcdefghijk" * 10**5))
        # If numpy doesn't do the manual comparison/mask, these
        # unorderable mixed types are what cause the exception in numpy
        in_list = [-1, "a", "b", "G", "Y", "Z", "E", "K", "E", "S", "I", "R", "R"] * 6

        assert s.isin(in_list).sum() == 200000

    def test_isin_with_string_scalar(self):
        # GH#4763
        s = pd.Series(["A", "B", "C", "a", "B", "B", "A", "C"])
        msg = (
            r"only list-like objects are allowed to be passed to isin\(\), "
            r"you passed a `str`"
        )
        with pytest.raises(TypeError, match=msg):
            s.isin("a")

        s = pd.Series(["aaa", "b", "c"])
        with pytest.raises(TypeError, match=msg):
            s.isin("aaa")

    def test_isin_datetimelike_mismatched_reso(self):
        expected = pd.Series([True, True, False, False, False])

        ser = pd.Series(pd.date_range("jan-01-2013", "jan-05-2013"))

        # fails on dtype conversion in the first place
        day_values = np.asarray(ser[0:2].values).astype("datetime64[D]")
        result = ser.isin(day_values)
        tm.assert_series_equal(result, expected)

        dta = ser[:2]._values.astype("M8[s]")
        result = ser.isin(dta)
        tm.assert_series_equal(result, expected)

    def test_isin_datetimelike_mismatched_reso_list(self):
        expected = pd.Series([True, True, False, False, False])

        ser = pd.Series(pd.date_range("jan-01-2013", "jan-05-2013"))

        dta = ser[:2]._values.astype("M8[s]")
        result = ser.isin(list(dta))
        tm.assert_series_equal(result, expected)

    def test_isin_with_i8(self):
        # GH#5021

        expected = pd.Series([True, True, False, False, False])
        expected2 = pd.Series([False, True, False, False, False])

        # datetime64[ns]
        s = pd.Series(pd.date_range("jan-01-2013", "jan-05-2013"))

        result = s.isin(s[0:2])
        tm.assert_series_equal(result, expected)

        result = s.isin(s[0:2].values)
        tm.assert_series_equal(result, expected)

        result = s.isin([s[1]])
        tm.assert_series_equal(result, expected2)

        result = s.isin([np.datetime64(s[1])])
        tm.assert_series_equal(result, expected2)

        result = s.isin(set(s[0:2]))
        tm.assert_series_equal(result, expected)

        # timedelta64[ns]
        s = pd.Series(pd.to_timedelta(range(5), unit="D"))
        result = s.isin(s[0:2])
        tm.assert_series_equal(result, expected)

    @pytest.mark.parametrize("empty", [[], pd.Series(dtype=object), np.array([])])
    def test_isin_empty(self, empty):
        # see GH#16991
        s = pd.Series(["a", "b"])
        expected = pd.Series([False, False])

        result = s.isin(empty)
        tm.assert_series_equal(expected, result)

    def test_isin_read_only(self):
        # https://github.com/pandas-dev/pandas/issues/37174
        arr = np.array([1, 2, 3])
        arr.setflags(write=False)
        s = pd.Series([1, 2, 3])
        result = s.isin(arr)
        expected = pd.Series([True, True, True])
        tm.assert_series_equal(result, expected)

    @pytest.mark.parametrize("dtype", [object, None])
    def test_isin_dt64_values_vs_ints(self, dtype):
        # GH#36621 dont cast integers to datetimes for isin
        dti = pd.date_range("2013-01-01", "2013-01-05")
        ser = pd.Series(dti)

        comps = np.asarray([1356998400000000000], dtype=dtype)

        res = dti.isin(comps)
        expected = np.array([False] * len(dti), dtype=bool)
        tm.assert_numpy_array_equal(res, expected)

        res = ser.isin(comps)
        tm.assert_series_equal(res, pd.Series(expected))

        res = pd.core.algorithms.isin(ser, comps)
        tm.assert_numpy_array_equal(res, expected)

    def test_isin_tzawareness_mismatch(self):
        dti = pd.date_range("2013-01-01", "2013-01-05")
        ser = pd.Series(dti)

        other = dti.tz_localize("UTC")

        res = dti.isin(other)
        expected = np.array([False] * len(dti), dtype=bool)
        tm.assert_numpy_array_equal(res, expected)

        res = ser.isin(other)
        tm.assert_series_equal(res, pd.Series(expected))

        res = pd.core.algorithms.isin(ser, other)
        tm.assert_numpy_array_equal(res, expected)

    def test_isin_period_freq_mismatch(self):
        dti = pd.date_range("2013-01-01", "2013-01-05")
        pi = dti.to_period("M")
        ser = pd.Series(pi)

        # We construct another PeriodIndex with the same i8 values
        #  but different dtype
        dtype = dti.to_period("Y").dtype
        other = PeriodArray._simple_new(pi.asi8, dtype=dtype)

        res = pi.isin(other)
        expected = np.array([False] * len(pi), dtype=bool)
        tm.assert_numpy_array_equal(res, expected)

        res = ser.isin(other)
        tm.assert_series_equal(res, pd.Series(expected))

        res = pd.core.algorithms.isin(ser, other)
        tm.assert_numpy_array_equal(res, expected)

    @pytest.mark.parametrize("values", [[-9.0, 0.0], [-9, 0]])
    def test_isin_float_in_int_series(self, values):
        # GH#19356 GH#21804
        ser = pd.Series(values)
        result = ser.isin([-9, -0.5])
        expected = pd.Series([True, False])
        tm.assert_series_equal(result, expected)

    @pytest.mark.parametrize("dtype", ["boolean", "Int64", "Float64"])
    @pytest.mark.parametrize(
        "data,values,expected",
        [
            ([0, 1, 0], [1], [False, True, False]),
            ([0, 1, 0], [1, pd.NA], [False, True, False]),
            ([0, pd.NA, 0], [1, 0], [True, False, True]),
            ([0, 1, pd.NA], [1, pd.NA], [False, True, True]),
            ([0, 1, pd.NA], [1, np.nan], [False, True, False]),
            ([0, pd.NA, pd.NA], [np.nan, pd.NaT, None], [False, False, False]),
        ],
    )
    def test_isin_masked_types(self, dtype, data, values, expected):
        # GH#42405
        ser = pd.Series(data, dtype=dtype)

        result = ser.isin(values)
        expected = pd.Series(expected, dtype="boolean")

        tm.assert_series_equal(result, expected)


def test_isin_large_series_mixed_dtypes_and_nan(monkeypatch):
    # https://github.com/pandas-dev/pandas/issues/37094
    # combination of object dtype for the values
    # and > _MINIMUM_COMP_ARR_LEN elements
    min_isin_comp = 5
    ser = pd.Series([1, 2, np.nan] * min_isin_comp)
    with monkeypatch.context() as m:
        m.setattr(algorithms, "_MINIMUM_COMP_ARR_LEN", min_isin_comp)
        result = ser.isin({"foo", "bar"})
    expected = pd.Series([False] * 3 * min_isin_comp)
    tm.assert_series_equal(result, expected)


@pytest.mark.parametrize(
    "dtype, data, values, expected",
    [
        ("boolean", [pd.NA, False, True], [False, pd.NA], [True, True, False]),
        ("Int64", [pd.NA, 2, 1], [1, pd.NA], [True, False, True]),
        ("boolean", [pd.NA, False, True], [pd.NA, True, "a", 20], [True, False, True]),
        ("boolean", [pd.NA, False, True], [], [False, False, False]),
        ("Float64", [20.0, 30.0, pd.NA], [pd.NA], [False, False, True]),
    ],
)
def test_isin_large_series_and_pdNA(dtype, data, values, expected, monkeypatch):
    # https://github.com/pandas-dev/pandas/issues/60678
    # combination of  large series (> _MINIMUM_COMP_ARR_LEN elements) and
    # values contains pdNA
    min_isin_comp = 2
    ser = pd.Series(data, dtype=dtype)
    expected = pd.Series(expected, dtype="boolean")

    with monkeypatch.context() as m:
        m.setattr(algorithms, "_MINIMUM_COMP_ARR_LEN", min_isin_comp)
        result = ser.isin(values)
    tm.assert_series_equal(result, expected)


@pytest.mark.parametrize(
    "dtype, data, values, expected",
    [
        ("int64[pyarrow]", [1, None], [1, pd.NA], [True, True]),
        ("binary[pyarrow]", [None, b"\xe3"], [pd.NA, b"\xe3"], [True, True]),
    ],
)
def test_isin_arrow_dtype_with_na_value(dtype, data, values, expected):
    # GH#63304
    pytest.importorskip("pyarrow")

    result = pd.Series(data, dtype=dtype).isin(values)
    expected = pd.Series(expected)

    tm.assert_series_equal(result, expected)


def test_isin_complex_numbers():
    # GH 17927
    array = [0, 1j, 1j, 1, 1 + 1j, 1 + 2j, 1 + 1j]
    result = pd.Series(array).isin([1j, 1 + 1j, 1 + 2j])
    expected = pd.Series([False, True, True, False, True, True, True], dtype=bool)
    tm.assert_series_equal(result, expected)


@pytest.mark.parametrize(
    "data,is_in",
    [([1, [2]], [1]), (["simple str", [{"values": 3}]], ["simple str"])],
)
def test_isin_filtering_with_mixed_object_types(data, is_in):
    # GH 20883

    ser = pd.Series(data)
    result = ser.isin(is_in)
    expected = pd.Series([True, False])

    tm.assert_series_equal(result, expected)


@pytest.mark.parametrize("data", [[1, 2, 3], [1.0, 2.0, 3.0]])
@pytest.mark.parametrize("isin", [[1, 2], [1.0, 2.0]])
def test_isin_filtering_on_iterable(data, isin):
    # GH 50234

    ser = pd.Series(data)
    result = ser.isin(i for i in isin)
    expected_result = pd.Series([True, True, False])

    tm.assert_series_equal(result, expected_result)


@pytest.mark.parametrize("set_cls", [set, frozenset])
@pytest.mark.parametrize("dtype", ["int64", "int32", "uint8", "uint64"])
@pytest.mark.parametrize("n_comps", [2, 20])
def test_isin_set_matches_list(set_cls, dtype, n_comps):
    # GH#25507: passing a set must match the list-based path. With a small
    # comps (n_comps=2 < len(targets)) this exercises the set-membership fast
    # path; with a larger comps it exercises the materialize-to-list fallback.
    ser = pd.Series(range(n_comps), dtype=dtype)
    targets = [1, 3, 5, 7, 9, 11, 13]
    expected = ser.isin(list(targets))
    result = ser.isin(set_cls(targets))
    tm.assert_series_equal(result, expected)


@pytest.mark.parametrize("set_cls", [set, frozenset])
def test_isin_set_matches_list_bool(set_cls):
    # GH#25507: boolean comps go through the same set fast path (taken here
    # since the single-element comps is smaller than the two-element set).
    ser = pd.Series([True])
    expected = ser.isin([True, False])
    result = ser.isin(set_cls([True, False]))
    tm.assert_series_equal(result, expected)


def test_isin_empty_set():
    # GH#25507 set fast path must handle empty set
    ser = pd.Series([1, 2, 3])
    result = ser.isin(set())
    expected = pd.Series([False, False, False])
    tm.assert_series_equal(result, expected)


@pytest.mark.parametrize("n_comps", [1, 2])
@pytest.mark.parametrize("magnitude", [2**53, 2**53 + 1])
@pytest.mark.parametrize("values", [{2.0**53, 0.5}, {2**53 + 1, 0.5}])
def test_isin_set_large_int_comps_matches_list(n_comps, magnitude, values):
    # GH#25507: comps at or beyond float64's exact-int range can match
    # lossily under the materialized path (which casts through float64) but
    # not under exact set membership; the result must not depend on the
    # comps size.
    ser = pd.Series([magnitude] * n_comps, dtype="int64")
    result = ser.isin(values)
    expected = ser.isin(list(values))
    tm.assert_series_equal(result, expected)