from collections import defaultdict
from datetime import datetime
from itertools import product

import numpy as np
import pytest

import pandas as pd
import pandas._testing as tm
from pandas.core import algorithms
from pandas.core.algorithms import safe_sort
import pandas.core.common as com
from pandas.core.sorting import (
    _decons_group_index,
    get_group_index,
    is_int64_overflow_possible,
    lexsort_indexer,
    nargsort,
)


@pytest.fixture
def left_right():
    low, high, n = -1 << 10, 1 << 10, 600
    left = pd.DataFrame(
        np.random.default_rng(2).integers(low, high, (n, 7)), columns=list("ABCDEFG")
    )
    shape = left.nunique().to_numpy()
    assert is_int64_overflow_possible(shape)
    left["left"] = left.sum(axis=1)
    right = left.sample(
        frac=1, random_state=np.random.default_rng(2), ignore_index=True
    )
    right.columns = [*right.columns[:-1].tolist(), "right"]
    right["right"] *= -1
    return left, right


class TestSorting:
    @pytest.mark.slow
    def test_int64_overflow(self):
        B = np.concatenate((np.arange(149), np.arange(149), np.arange(75)))
        A = np.arange(373)
        df = pd.DataFrame(
            {
                "A": A,
                "B": B,
                "C": A,
                "D": B,
                "E": A,
                "F": B,
                "G": A,
                "H": B,
                "values": np.random.default_rng(2).standard_normal(373),
            }
        )
        assert is_int64_overflow_possible(df[list("ABCDEFGH")].nunique().to_numpy())

        lg = df.groupby(["A", "B", "C", "D", "E", "F", "G", "H"])
        rg = df.groupby(["H", "G", "F", "E", "D", "C", "B", "A"])

        left = lg.sum()["values"]
        right = rg.sum()["values"]

        exp_index, _ = left.index.sortlevel()
        tm.assert_index_equal(left.index, exp_index)

        exp_index, _ = right.index.sortlevel(0)
        tm.assert_index_equal(right.index, exp_index)

        tups = list(map(tuple, df[["A", "B", "C", "D", "E", "F", "G", "H"]].values))
        tups = com.asarray_tuplesafe(tups)

        expected = df.groupby(tups).sum()["values"]

        for k, v in expected.items():
            assert left[k] == right[k[::-1]]
            assert left[k] == v
        assert len(left) == len(right)

    def test_int64_overflow_groupby_large_range(self):
        # GH9096
        values = range(55109)
        data = pd.DataFrame.from_dict(
            {"a": values, "b": values, "c": values, "d": values}
        )
        grouped = data.groupby(["a", "b", "c", "d"])
        assert len(grouped) == len(values)

    @pytest.mark.slow
    @pytest.mark.parametrize("agg", ["mean", "median"])
    def test_int64_overflow_groupby_large_df_shuffled(self, agg):
        rs = np.random.default_rng(2)
        arr = rs.integers(-1 << 12, 1 << 12, (11595, 5))
        i = rs.choice(len(arr), len(arr))
        arr = np.vstack((arr, arr[i]))  # add some duplicate rows

        i = rs.permutation(len(arr))
        arr = arr[i]  # shuffle rows

        df = pd.DataFrame(arr, columns=list("abcde"))
        df["jim"], df["joe"] = np.zeros((2, len(df)))
        gr = df.groupby(list("abcde"))

        # verify this is testing what it is supposed to test!
        assert is_int64_overflow_possible(
            tuple(ping.ngroups for ping in gr._grouper.groupings)
        )

        mi = pd.MultiIndex.from_arrays(
            [ar.ravel() for ar in np.array_split(np.unique(arr, axis=0), 5, axis=1)],
            names=list("abcde"),
        )

        res = pd.DataFrame(
            np.zeros((len(mi), 2)), columns=["jim", "joe"], index=mi
        ).sort_index()

        tm.assert_frame_equal(getattr(gr, agg)(), res)

    @pytest.mark.parametrize(
        "order, na_position, exp",
        [
            [
                True,
                "last",
                list(range(5, 105)) + list(range(5)) + list(range(105, 110)),
            ],
            [
                True,
                "first",
                list(range(5)) + list(range(105, 110)) + list(range(5, 105)),
            ],
            [
                False,
                "last",
                list(range(104, 4, -1)) + list(range(5)) + list(range(105, 110)),
            ],
            [
                False,
                "first",
                list(range(5)) + list(range(105, 110)) + list(range(104, 4, -1)),
            ],
        ],
    )
    def test_lexsort_indexer(self, order, na_position, exp):
        keys = [[np.nan] * 5 + list(range(100)) + [np.nan] * 5]
        result = lexsort_indexer(keys, orders=order, na_position=na_position)
        tm.assert_numpy_array_equal(result, np.array(exp, dtype=np.intp))

    @pytest.mark.parametrize("order", [True, False])
    @pytest.mark.parametrize("na_position", ["last", "first"])
    @pytest.mark.parametrize(
        "dtype",
        [
            np.float64,
            np.float32,
            np.int64,
            np.uint64,
            np.bool_,
        ],
    )
    def test_lexsort_indexer_numeric_dtypes(self, dtype, order, na_position):
        # GH#15389 - fast path for numeric dtypes should match Categorical path

        if dtype == np.bool_:
            key1 = np.array([True, False, True, False, True])
        else:
            key1 = np.array([3, 1, 5, 2, 4], dtype=dtype)
        key2 = np.arange(len(key1), dtype=np.int64)
        keys = [key1, key2]

        result = lexsort_indexer(keys, orders=order, na_position=na_position)

        # Compare with Categorical-based reference implementation
        labels = []
        for key, ord_ in zip(reversed(keys), [order, order], strict=True):
            cat = pd.Categorical(key, ordered=True)
            codes = cat.codes
            num_categories = len(cat.categories)
            mask = codes == -1
            if na_position == "last" and mask.any():
                codes = np.where(mask, num_categories, codes)
            if not ord_:
                codes = np.where(mask, codes, num_categories - codes - 1)
            labels.append(codes)
        expected = np.lexsort(labels)

        tm.assert_numpy_array_equal(result, expected)

    @pytest.mark.parametrize("order", [True, False])
    @pytest.mark.parametrize("na_position", ["last", "first"])
    def test_lexsort_indexer_float_with_nan(self, order, na_position):
        # GH#15389 - float fast path must handle NaN placement correctly

        key1 = np.array([3.0, np.nan, 1.0, np.nan, 2.0])
        key2 = np.arange(len(key1), dtype=np.float64)
        keys = [key1, key2]

        result = lexsort_indexer(keys, orders=order, na_position=na_position)

        labels = []
        for key, ord_ in zip(reversed(keys), [order, order], strict=True):
            cat = pd.Categorical(key, ordered=True)
            codes = cat.codes
            num_categories = len(cat.categories)
            mask = codes == -1
            if na_position == "last" and mask.any():
                codes = np.where(mask, num_categories, codes)
            if not ord_:
                codes = np.where(mask, codes, num_categories - codes - 1)
            labels.append(codes)
        expected = np.lexsort(labels)

        tm.assert_numpy_array_equal(result, expected)

    def test_lexsort_indexer_int_descending_overflow(self):
        # GH#15389 - int64 min value must not overflow during descending sort
        key1 = np.array([2, np.iinfo(np.int64).min, 1, 0], dtype=np.int64)
        key2 = np.arange(len(key1), dtype=np.int64)

        result = lexsort_indexer([key1, key2], orders=False, na_position="last")

        # Descending: 2, 1, 0, int64_min
        expected = np.array([0, 2, 3, 1], dtype=np.intp)
        tm.assert_numpy_array_equal(result, expected)

    @pytest.mark.parametrize(
        "ascending, na_position, exp",
        [
            [
                True,
                "last",
                list(range(5, 105)) + list(range(5)) + list(range(105, 110)),
            ],
            [
                True,
                "first",
                list(range(5)) + list(range(105, 110)) + list(range(5, 105)),
            ],
            [
                False,
                "last",
                list(range(104, 4, -1)) + list(range(5)) + list(range(105, 110)),
            ],
            [
                False,
                "first",
                list(range(5)) + list(range(105, 110)) + list(range(104, 4, -1)),
            ],
        ],
    )
    def test_nargsort(self, ascending, na_position, exp):
        # list places NaNs last, np.array(..., dtype="O") may not place NaNs first
        items = np.array([np.nan] * 5 + list(range(100)) + [np.nan] * 5, dtype="O")

        # mergesort is the most difficult to get right because we want it to be
        # stable.

        # According to numpy/core/tests/test_multiarray, """The number of
        # sorted items must be greater than ~50 to check the actual algorithm
        # because quick and merge sort fall over to insertion sort for small
        # arrays."""

        result = nargsort(
            items, kind="mergesort", ascending=ascending, na_position=na_position
        )
        tm.assert_numpy_array_equal(result, np.array(exp), check_dtype=False)


class TestMerge:
    def test_int64_overflow_outer_merge(self):
        # #2690, combinatorial explosion
        df1 = pd.DataFrame(
            np.random.default_rng(2).standard_normal((1000, 7)),
            columns=[*list("ABCDEF"), "G1"],
        )
        df2 = pd.DataFrame(
            np.random.default_rng(3).standard_normal((1000, 7)),
            columns=[*list("ABCDEF"), "G2"],
        )
        result = pd.merge(df1, df2, how="outer")
        assert len(result) == 2000

    @pytest.mark.slow
    def test_int64_overflow_check_sum_col(self, left_right):
        left, right = left_right

        out = pd.merge(left, right, how="outer")
        assert len(out) == len(left)
        tm.assert_series_equal(out["left"], -out["right"], check_names=False)
        result = out.iloc[:, :-2].sum(axis=1)
        tm.assert_series_equal(out["left"], result, check_names=False)
        assert result.name is None

    @pytest.mark.slow
    @pytest.mark.filterwarnings("ignore:The inplace keyword in DataFrame.sort_values")
    def test_int64_overflow_how_merge(self, left_right, join_type):
        left, right = left_right

        out = pd.merge(left, right, how="outer")
        out.sort_values(out.columns.tolist(), inplace=True)
        tm.assert_frame_equal(out, pd.merge(left, right, how=join_type, sort=True))

    @pytest.mark.slow
    def test_int64_overflow_sort_false_order(self, left_right):
        left, right = left_right

        # check that left merge w/ sort=False maintains left frame order
        out = pd.merge(left, right, how="left", sort=False)
        tm.assert_frame_equal(left, out[left.columns.tolist()])

        out = pd.merge(right, left, how="left", sort=False)
        tm.assert_frame_equal(right, out[right.columns.tolist()])

    @pytest.mark.slow
    def test_int64_overflow_one_to_many_none_match(self, join_type, sort):
        # one-2-many/none match
        how = join_type
        low, high, n = -1 << 10, 1 << 10, 589
        left = pd.DataFrame(
            np.random.default_rng(2).integers(low, high, (n, 7)).astype("int64"),
            columns=list("ABCDEFG"),
        )

        # confirm that this is checking what it is supposed to check
        shape = left.nunique().to_numpy()
        assert is_int64_overflow_possible(shape)

        # add duplicates to left frame
        left = pd.concat([left, left], ignore_index=True)

        right = pd.DataFrame(
            np.random.default_rng(3).integers(low, high, (n // 2, 7)).astype("int64"),
            columns=list("ABCDEFG"),
        )

        # add duplicates & overlap with left to the right frame
        i = np.random.default_rng(4).choice(len(left), n)
        right = pd.concat([right, right, left.iloc[i]], ignore_index=True)

        left["left"] = np.random.default_rng(2).standard_normal(len(left))
        right["right"] = np.random.default_rng(2).standard_normal(len(right))

        # shuffle left & right frames
        left = left.sample(
            frac=1, ignore_index=True, random_state=np.random.default_rng(5)
        )
        right = right.sample(
            frac=1, ignore_index=True, random_state=np.random.default_rng(6)
        )

        # manually compute outer merge
        ldict, rdict = defaultdict(list), defaultdict(list)

        for idx, row in left.set_index(list("ABCDEFG")).iterrows():
            ldict[idx].append(row["left"])

        for idx, row in right.set_index(list("ABCDEFG")).iterrows():
            rdict[idx].append(row["right"])

        vals = []
        for k, lval in ldict.items():
            rval = rdict.get(k, [np.nan])
            for lv, rv in product(lval, rval):
                vals.append((*k, lv, rv))

        for k, rval in rdict.items():
            if k not in ldict:
                vals.extend((*k, np.nan, rv) for rv in rval)

        out = pd.DataFrame(vals, columns=[*list("ABCDEFG"), "left", "right"])
        out = out.sort_values(out.columns.to_list(), ignore_index=True)

        jmask = {
            "left": out["left"].notna(),
            "right": out["right"].notna(),
            "inner": out["left"].notna() & out["right"].notna(),
            "outer": np.ones(len(out), dtype="bool"),
        }

        mask = jmask[how]
        frame = out[mask].sort_values(out.columns.to_list(), ignore_index=True)
        assert mask.all() ^ mask.any() or how == "outer"

        res = pd.merge(left, right, how=how, sort=sort)
        if sort:
            kcols = list("ABCDEFG")
            tm.assert_frame_equal(
                res[kcols], res[kcols].sort_values(kcols, kind="mergesort")
            )

        # as in GH9092 dtypes break with outer/right join
        # 2021-12-18: dtype does not break anymore
        tm.assert_frame_equal(
            frame, res.sort_values(res.columns.to_list(), ignore_index=True)
        )


@pytest.mark.parametrize(
    "codes_list, shape",
    [
        [
            [
                np.tile([0, 1, 2, 3, 0, 1, 2, 3], 100).astype(np.int64),
                np.tile([0, 2, 4, 3, 0, 1, 2, 3], 100).astype(np.int64),
                np.tile([5, 1, 0, 2, 3, 0, 5, 4], 100).astype(np.int64),
            ],
            (4, 5, 6),
        ],
        [
            [
                np.tile(np.arange(10000, dtype=np.int64), 5),
                np.tile(np.arange(10000, dtype=np.int64), 5),
            ],
            (10000, 10000),
        ],
    ],
)
def test_decons(codes_list, shape):
    group_index = get_group_index(codes_list, shape, sort=True, xnull=True)
    codes_list2 = _decons_group_index(group_index, shape)

    for a, b in zip(codes_list, codes_list2, strict=True):
        tm.assert_numpy_array_equal(a, b)


@pytest.fixture(params=[5_000, 1], ids=["argsort", "counting"])
def switch_counting_sort_min_len(request, monkeypatch):
    # GH#66129 safe_sort ranks integer values spanning a modest range with a
    #  counting sort rather than an argsort, but only above a length gate.
    #  Lowering the gate runs the shared tests below down both paths, which
    #  have to agree.
    with monkeypatch.context() as m:
        m.setattr(algorithms, "_MIN_COUNTING_SORT_LEN", request.param)
        yield request.param


class TestSafeSort:
    @pytest.mark.parametrize(
        "arg, exp",
        [
            [[3, 1, 2, 0, 4], [0, 1, 2, 3, 4]],
            [
                np.array(list("baaacb"), dtype=object),
                np.array(list("aaabbc"), dtype=object),
            ],
            [[], []],
        ],
    )
    def test_basic_sort(self, arg, exp):
        result = safe_sort(np.array(arg))
        expected = np.array(exp)
        tm.assert_numpy_array_equal(result, expected)

    @pytest.mark.parametrize("dtype", ["int64", "uint64", "float64", "object"])
    @pytest.mark.parametrize("verify", [True, False])
    @pytest.mark.parametrize(
        "codes, exp_codes",
        [
            [[0, 1, 1, 2, 3, 0, -1, 4], [3, 1, 1, 2, 0, 3, -1, 4]],
            [[], []],
        ],
    )
    def test_codes(self, verify, codes, exp_codes, dtype, switch_counting_sort_min_len):
        values = np.array([3, 1, 2, 0, 4], dtype=dtype)
        expected = np.array([0, 1, 2, 3, 4], dtype=dtype)

        result, result_codes = safe_sort(
            values, codes, use_na_sentinel=True, verify=verify
        )
        expected_codes = np.array(exp_codes, dtype=np.intp)
        tm.assert_numpy_array_equal(result, expected)
        tm.assert_numpy_array_equal(result_codes, expected_codes)

    @pytest.mark.parametrize("dtype", ["int64", "uint64", "float64", "object"])
    def test_codes_out_of_bound(self, dtype, switch_counting_sort_min_len):
        values = np.array([3, 1, 2, 0, 4], dtype=dtype)
        expected = np.array([0, 1, 2, 3, 4], dtype=dtype)

        # out of bound indices
        codes = [0, 101, 102, 2, 3, 0, 99, 4]
        result, result_codes = safe_sort(values, codes, use_na_sentinel=True)
        expected_codes = np.array([3, -1, -1, 2, 0, 3, -1, 4], dtype=np.intp)
        tm.assert_numpy_array_equal(result, expected)
        tm.assert_numpy_array_equal(result_codes, expected_codes)

    @pytest.mark.parametrize("dtype", ["int64", "uint64"])
    def test_codes_out_of_bound_wrap(self, dtype, switch_counting_sort_min_len):
        # out of bound indices wrap when there is no sentinel to mask them with
        values = np.array([3, 1, 2, 0, 4], dtype=dtype)
        # wrapping these lands on 0, 1, 3, 2, 4 respectively
        codes = [0, 6, 8, 2, -1]

        _, result_codes = safe_sort(values, codes, use_na_sentinel=False)
        expected_codes = np.array([3, 1, 0, 2, 4], dtype=np.intp)
        tm.assert_numpy_array_equal(result_codes, expected_codes)

    @pytest.mark.parametrize("dtype", ["int64", "int32", "int16", "uint64"])
    @pytest.mark.parametrize("anchor", ["min", "max"])
    @pytest.mark.parametrize("assume_unique", [True, False])
    def test_codes_counting_sort_dtype_bounds(
        self, dtype, anchor, assume_unique, switch_counting_sort_min_len
    ):
        # GH#66129 anchor the values at either end of the dtype so that shifting
        #  by the minimum is exercised for large-negative and large-positive bases
        info = np.iinfo(dtype)
        num = 5
        base = info.min if anchor == "min" else info.max - num + 1
        values = np.arange(base, base + num, dtype=dtype)[::-1].copy()
        codes = np.array([0, 3, 4, 1], dtype=np.intp)

        result, result_codes = safe_sort(values, codes, assume_unique=assume_unique)
        tm.assert_numpy_array_equal(result, np.sort(values))
        tm.assert_numpy_array_equal(result_codes, np.array([4, 1, 0, 3], dtype=np.intp))

    def test_codes_counting_sort_duplicates(self, switch_counting_sort_min_len):
        # GH#66129 the counting pass doubles as the uniqueness check
        values = np.array([3, 1, 2, 0, 3], dtype="int64")
        codes = np.arange(len(values), dtype=np.intp)
        with pytest.raises(ValueError, match="values should be unique"):
            safe_sort(values, codes)

    def test_codes_counting_sort_wide_range(self, switch_counting_sort_min_len):
        # GH#66129 too wide a range declines the counting path; the results have
        #  to match the argsort path it falls back to
        values = np.array([300, 1, 20000, 0, 4], dtype="int64")
        codes = np.array([0, 2, 4, -1, 1], dtype=np.intp)

        result, result_codes = safe_sort(values, codes)
        tm.assert_numpy_array_equal(result, np.sort(values))
        tm.assert_numpy_array_equal(
            result_codes, np.array([3, 4, 2, -1, 1], dtype=np.intp)
        )

    def test_codes_counting_sort_above_real_min_len(self):
        # GH#66129 the unpatched cutoff has to actually reach the counting path
        num = algorithms._MIN_COUNTING_SORT_LEN + 1
        rng = np.random.default_rng(2)
        values = rng.permutation(np.arange(num, dtype="int64"))
        codes = rng.integers(0, num, size=500).astype(np.intp)

        result, result_codes = safe_sort(values, codes.copy())

        # ranks[i] is the position of values[i] once sorted
        ranks = np.empty(num, dtype=np.intp)
        ranks[values.argsort()] = np.arange(num, dtype=np.intp)
        tm.assert_numpy_array_equal(result, np.sort(values))
        tm.assert_numpy_array_equal(result_codes, ranks.take(codes))

    @pytest.mark.parametrize("codes", [[-1, -1], [2, -1], [2, 2]])
    def test_codes_empty_array_out_of_bound(self, codes):
        empty_values = np.array([])
        expected_codes = -np.ones_like(codes, dtype=np.intp)
        _, result_codes = safe_sort(empty_values, codes)
        tm.assert_numpy_array_equal(result_codes, expected_codes)

    def test_mixed_integer(self):
        values = np.array(["b", 1, 0, "a", 0, "b"], dtype=object)
        result = safe_sort(values)
        expected = np.array([0, 0, 1, "a", "b", "b"], dtype=object)
        tm.assert_numpy_array_equal(result, expected)

    def test_mixed_integer_with_codes(self):
        values = np.array(["b", 1, 0, "a"], dtype=object)
        codes = [0, 1, 2, 3, 0, -1, 1]
        result, result_codes = safe_sort(values, codes)
        expected = np.array([0, 1, "a", "b"], dtype=object)
        expected_codes = np.array([3, 1, 0, 2, 3, -1, 1], dtype=np.intp)
        tm.assert_numpy_array_equal(result, expected)
        tm.assert_numpy_array_equal(result_codes, expected_codes)

    def test_unsortable(self):
        # GH 13714
        arr = np.array([1, 2, datetime(2011, 1, 1), 0, 3], dtype=object)
        msg = "'[<>]' not supported between instances of .*"
        with pytest.raises(TypeError, match=msg):
            safe_sort(arr)

    @pytest.mark.parametrize(
        "arg, codes, err, msg",
        [
            [1, None, TypeError, "Only np.ndarray, ExtensionArray, and Index"],
            [np.array([0, 1, 2]), 1, TypeError, "Only list-like objects or None"],
            [np.array([0, 1, 2, 1]), [0, 1], ValueError, "values should be unique"],
        ],
    )
    def test_exceptions(self, arg, codes, err, msg):
        with pytest.raises(err, match=msg):
            safe_sort(values=arg, codes=codes)

    @pytest.mark.parametrize(
        "arg, exp", [[[1, 3, 2], [1, 2, 3]], [[1, 3, pd.NA, 2], [1, 2, 3, pd.NA]]]
    )
    def test_extension_array(self, arg, exp):
        a = pd.array(arg, dtype="Int64")
        result = safe_sort(a)
        expected = pd.array(exp, dtype="Int64")
        tm.assert_extension_array_equal(result, expected)

    @pytest.mark.parametrize("verify", [True, False])
    def test_extension_array_codes(self, verify):
        a = pd.array([1, 3, 2], dtype="Int64")
        result, codes = safe_sort(a, [0, 1, -1, 2], use_na_sentinel=True, verify=verify)
        expected_values = pd.array([1, 2, 3], dtype="Int64")
        expected_codes = np.array([0, 2, -1, 1], dtype=np.intp)
        tm.assert_extension_array_equal(result, expected_values)
        tm.assert_numpy_array_equal(codes, expected_codes)


def test_mixed_str_null(nulls_fixture):
    values = np.array(["b", nulls_fixture, "a", "b"], dtype=object)
    result = safe_sort(values)
    expected = np.array(["a", "b", "b", nulls_fixture], dtype=object)
    tm.assert_numpy_array_equal(result, expected)


def test_safe_sort_multiindex():
    # GH#48412
    arr1 = pd.Series([2, 1, pd.NA, pd.NA], dtype="Int64")
    arr2 = [2, 1, 3, 3]
    midx = pd.MultiIndex.from_arrays([arr1, arr2])
    result = safe_sort(midx)
    expected = pd.MultiIndex.from_arrays(
        [pd.Series([1, 2, pd.NA, pd.NA], dtype="Int64"), [1, 2, 3, 3]]
    )
    tm.assert_index_equal(result, expected)