import numpy as np
from numpy.random import randn
import pytest

import pandas as pd
from pandas import DataFrame, Index, MultiIndex, Series
import pandas._testing as tm
from pandas.core.reshape.concat import concat
from pandas.core.reshape.merge import merge


@pytest.fixture
def left():
    """left dataframe (not multi-indexed) for multi-index join tests"""
    # a little relevant example with NAs
    key1 = ["bar", "bar", "bar", "foo", "foo", "baz", "baz", "qux", "qux", "snap"]
    key2 = ["two", "one", "three", "one", "two", "one", "two", "two", "three", "one"]

    data = np.random.randn(len(key1))
    return DataFrame({"key1": key1, "key2": key2, "data": data})


@pytest.fixture
def right():
    """right dataframe (multi-indexed) for multi-index join tests"""
    index = MultiIndex(
        levels=[["foo", "bar", "baz", "qux"], ["one", "two", "three"]],
        codes=[[0, 0, 0, 1, 1, 2, 2, 3, 3, 3], [0, 1, 2, 0, 1, 1, 2, 0, 1, 2]],
        names=["key1", "key2"],
    )

    return DataFrame(
        np.random.randn(10, 3), index=index, columns=["j_one", "j_two", "j_three"]
    )


@pytest.fixture
def left_multi():
    return DataFrame(
        dict(
            Origin=["A", "A", "B", "B", "C"],
            Destination=["A", "B", "A", "C", "A"],
            Period=["AM", "AM", "IP", "AM", "OP"],
            TripPurp=["hbw", "nhb", "hbo", "nhb", "hbw"],
            Trips=[1987, 3647, 2470, 4296, 4444],
        ),
        columns=["Origin", "Destination", "Period", "TripPurp", "Trips"],
    ).set_index(["Origin", "Destination", "Period", "TripPurp"])


@pytest.fixture
def right_multi():
    return DataFrame(
        dict(
            Origin=["A", "A", "B", "B", "C", "C", "E"],
            Destination=["A", "B", "A", "B", "A", "B", "F"],
            Period=["AM", "AM", "IP", "AM", "OP", "IP", "AM"],
            LinkType=["a", "b", "c", "b", "a", "b", "a"],
            Distance=[100, 80, 90, 80, 75, 35, 55],
        ),
        columns=["Origin", "Destination", "Period", "LinkType", "Distance"],
    ).set_index(["Origin", "Destination", "Period", "LinkType"])


@pytest.fixture
def on_cols_multi():
    return ["Origin", "Destination", "Period"]


@pytest.fixture
def idx_cols_multi():
    return ["Origin", "Destination", "Period", "TripPurp", "LinkType"]


class TestMergeMulti:
    def setup_method(self):
        self.index = MultiIndex(
            levels=[["foo", "bar", "baz", "qux"], ["one", "two", "three"]],
            codes=[[0, 0, 0, 1, 1, 2, 2, 3, 3, 3], [0, 1, 2, 0, 1, 1, 2, 0, 1, 2]],
            names=["first", "second"],
        )
        self.to_join = DataFrame(
            np.random.randn(10, 3),
            index=self.index,
            columns=["j_one", "j_two", "j_three"],
        )

        # a little rele
vant example with NAs
        key1 = ["bar", "bar", "bar", "foo", "foo", "baz", "baz", "qux", "qux", "snap"]
        key2 = [
            "two",
            "one",
            "three",
            "one",
            "two",
            "one",
            "two",
            "two",
            "three",
            "one",
        ]

        data = np.random.randn(len(key1))
        self.data = DataFrame({"key1": key1, "key2": key2, "data": data})

    def test_merge_on_mul
tikey(self, left, right, join_type):
        on_cols = ["key1", "key2"]
        result = left.join(right, on=on_cols, how=join_type).reset_index(drop=True)

        expected = pd.merge(left, right.reset_index(), on=on_cols, how=join_type)

        tm.assert_frame_equal(result, expected)

        result = left.join(right, on=on_cols, how=join_type, sort=True).reset_index(
            drop=True
        )

        expected = pd.merge(
            left, right.reset_index(), on=on_cols, how=join_type, sort=True
        )

        tm.assert_frame_equal(result, expected)

    @pytest.mark.parametrize("sort", [False, True])
    def test_left_join_multi_index(self, left, right, sort):
        icols = ["1st", "2nd", "3rd"]

        def bind_cols(df):
            iord = lambda a: 0 if a != a else ord(a)
            f = lambda ts: ts.map(iord) - ord("a")
            return f(df["1st"]) + f(df["3rd"]) * 1e2 + df["2nd"].fillna(0) * 1e4

        def run_asserts(left, right, sort):
            res = left.join(right, on=icols, how="left", sort=sort)

            assert len(left) < len(res) + 1
            assert not res["4th"].isna().any()
            assert not res["5th"].isna().any()

            tm.assert_series_equal(res["4th"], -res["5th"], check_names=False)
            result = bind_cols(res.iloc[:, :-2])
            tm.assert_series_equal(res["4th"], result, check_names=False)
            assert result.name is None

            if sort:
                tm.assert_frame_equal(res, res.sort_values(icols, kind="mergesort"))

            out = merge(left, right.reset_index(), on=icols, sort=sort, how="left")

            res.index = np.arange(len(res))
            tm.assert_frame_equal(out, res)

        lc = list(map(chr, np.arange(ord("a"), ord("z") + 1)))
        left = DataFrame(np.random.choice(lc, (5000, 2)), columns=["1st", "3rd"])
        left.insert(1, "2nd", np.random.randint(0, 1000, len(left)))

        i = np.random.permutation(len(left))
        right = left.iloc[i].copy()

        left["4th"] = bind_cols(left)
        right["5th"] = -bind_cols(right)
        right.set_index(icols, inplace=True)

        run_asserts(left, right, sort)

        # inject some nulls
        left.loc[1::23, "1st"] = np.nan
        left.loc[2::37, "2nd"] = np.nan
        left.loc[3::43, "3rd"] = np.nan
        left["4th"] = bind_cols(left)

        i = np.random.permutation(len(left))
        right = left.iloc[i, :-1]
        right["5th"] = -bind_cols(right)
        right.set_index(icols, inplace=True)

        run_asserts(left, right, sort)

    @pytest.mark.parametrize("sort", [False, True])
    def test_merge_right_vs_left(self, left, right, sort):
        # compare left vs right merge with multikey
        on_cols = ["key1", "key2"]
        merged_left_right = left.merge(
            right, left_on=on_cols, right_index=True, how="left", sort=sort
        )

        merge_right_left = right.merge(
            left, right_on=on_cols, left_index=True, how="right", sort=sort
        )

        # Reorder columns
        merge_right_left = merge_right_left[merged_left_right.columns]

        tm.assert_frame_equal(merged_left_right, merge_right_left)

    def test_merge_multiple_cols_with_mixed_cols_index(self):
        # GH29522
        s = pd.Series(
            range(6),
            pd.MultiIndex.from_product([["A", "B"], [1, 2, 3]], names=["lev1", "lev2"]),
            name="Amount",
        )
        df = pd.DataFrame(
            {"lev1": list("AAABBB"), "lev2": [1, 2, 3, 1, 2, 3], "col": 0}
        )
        result = pd.merge(df, s.reset_index(), on=["lev1", "lev2"])
        expected = pd.DataFrame(
            {
                "lev1": list("AAABBB"),
                "lev2": [1, 2, 3, 1, 2, 3],
                "col": [0] * 6,
                "Amount": range(6),
            }
        )
        tm.assert_frame_equal(result, expected)

    def test_compress_group_combinations(self):

        # ~ 40000000 possible unique groups
        key1 = tm.rands_array(10, 10000)
        key1 = np.tile(key1, 2)
        key2 = key1[::-1]

        df = DataFrame({"key1": key1, "key2": key2, "value1": np.random.randn(20000)})

        df2 = DataFrame(
            {"key1": key1[::2], "key2": key2[::2], "value2": np.random.randn(10000)}
        )

        # just to hit the label compression code path
        merge(df, df2, how="outer")

    def test_left_join_index_preserve_order(self):

        on_cols = ["k1", "k2"]
        left = DataFrame(
            {
                "k1": [0, 1, 2] * 8,
                "k2": ["foo", "bar"] * 12,
                "v": np.array(np.arange(24), dtype=np.int64),
            }
        )

        index = MultiIndex.from_tuples([(2, "bar"), (1, "foo")])
        right = DataFrame({"v2": [5, 7]}, index=index)

        result = left.join(right, on=on_cols)

        expected = left.copy()
        expected["v2"] = np.nan
        expected.loc[(expected.k1