import numpy as np
import pytest

import pandas as pd
from pandas import DataFrame, MultiIndex, Series, Timestamp, date_range, isna, notna
import pandas._testing as tm
import pandas.core.common as com


class TestMultiIndexSetItem:
    def test_setitem_multiindex(self):
        for index_fn in ("loc",):

            def assert_equal(a, b):
                assert a == b

            def check(target, indexers, value, compare_fn, expected=None):
                fn = getattr(target, index_fn)
                fn.__setitem__(indexers, value)
                result = fn.__getitem__(indexers)
                if expected is None:
                    expected = value
                compare_fn(result, expected)

            # GH7190
            index = MultiIndex.from_product(
                [np.arange(0, 100), np.arange(0, 80)], names=["time", "firm"]
            )
            t, n = 0, 2
            df = DataFrame(
                np.nan,
                columns=["A", "w", "l", "a", "x", "X", "d", "profit"],
                index=index,
            )
            check(target=df, indexers=((t, n), "X"), value=0, compare_fn=assert_equal)

            df = DataFrame(
                -999, columns=["A", "w", "l", "a", "x", "X", "d", "profit"], index=index
            )
            check(target=df, indexers=((t, n), "X"), value=1, compare_fn=assert_equal)

            df = DataFrame(
                columns=["A", "w", "l", "a", "x", "X", "d", "profit"], index=index
            )
            check(target=df, indexers=((t, n), "X"), value=2, compare_fn=assert_equal)

            # gh-7218: assigning with 0-dim arrays
            df = DataFrame(
                -999, columns=["A", "w", "l", "a", "x", "X", "d", "profit"], index=index
            )
            check(
                target=df,
                indexers=((t, n), "X"),
                value=np.array(3),
                compare_fn=assert_equal,
                expected=3,
            )

            # GH5206
            df = DataFrame(
                np.arange(25).reshape(5, 5), columns="A,B,C,D,E".split(","), dtype=float
            )
            df["F"] = 99
            row_selection = df["A"] % 2 == 0
            col_selection = ["B", "C"]
            df.loc[row_selection, col_selection] = df["F"]
            output = DataFrame(99.0, index=[0, 2, 4], columns=["B", "C"])
            tm.assert_frame_equal(df.loc[row_selection, col_selection], output)
            check(
                target=df,
                indexers=(row_selection, col_selection),
                value=df["F"],
                compare_fn=tm.assert_frame_equal,
                expected=output,
            )

            # GH11372
            idx = MultiIndex.from_product(
                [["A", "B", "C"], date_range("2015-01-01", "2015-04-01", freq="MS")]
            )
            cols = MultiIndex.from_product(
                [["foo", "bar"], date_range("2016-01-01", "2016-02-01", freq="MS")]
            )

            df = DataFrame(np.random.random((12, 4)), index=idx, columns=cols)

            subidx = MultiIndex.from_tuples(
                [("A", Timestamp("2015-01-01")), ("A", Timestamp("2015-02-01"))]
            )
            subcols = MultiIndex.from_tuples(
                [("foo", Timestamp("2016-01-01")), ("foo", Timestamp("2016-02-01"))]
            )

            vals = DataFrame(np.random.random((2, 2)), index=subidx, columns=subcols)
            check(
                target=df,
                indexers=(subidx, subcols),
                value=vals,
                compare_fn=tm.assert_frame_equal,
            )
            # set all columns
            vals = DataFrame(np.random.random((2, 4)), index=subidx, columns=cols)
            check(
                target=df,
                indexers=(subidx, slice(None, None, None)),
                value=vals,
                compare_fn=tm.assert_frame_equal,
            )
            # identity
            copy = df.copy()
            check(
                target=df,
                indexers=(df.index, df.columns),
                value=df,
                compare_fn=tm.assert_frame_equal,
                expected=copy,
            )

    def test_multiindex_setitem(self):

        # GH 3738
        # setting with a multi-index right hand side
        arrays = [
            np.array(["bar", "bar", "baz", "qux", "qux", "bar"]),
            np.array(["one", "two", "one", "one", "two", "one"]),
            np.arange(0, 6, 1),
        ]

        df_orig = DataFrame(
            np.random.randn(6, 3), index=arrays, columns=["A", "B", "C"]
        ).sort_index()

        expected = df_orig.loc[["bar"]] * 2
        df = df_orig.copy()
        df.loc[["bar"]] *= 2
        tm.assert_frame_equal(df.loc[["bar"]], expected)

        # raise because these have differing levels
        msg = "cannot align on a multi-index with out specifying the join levels"
        with pytest.raises(TypeError, match=msg):
            df.loc["bar"] *= 2

        # from SO
        # https://stackoverflow.com/questions/24572040/pandas-access-the-level-of-multiindex-for-inplace-operation
        df_orig = DataFrame.from_dict(
            {
                "price": {
                    ("DE", "Coal", "Stock"): 2,
                    ("DE", "Gas", "Stock"): 4,
                    ("DE", "Elec", "Demand"): 1,
                    ("FR", "Gas", "Stock"): 5,
                    ("FR", "Solar", "SupIm"): 0,
                    ("FR", "Wind", "SupIm"): 0,
                }
            }
        )
        df_orig.index = MultiIndex.from_tuples(
            df_orig.index, names=["Sit", "Com", "Type"]
        )

        expected = df_orig.copy()
        expected.iloc[[0, 2, 3]] *= 2

        idx = pd.IndexSlice
        df = df_orig.copy()
        df.loc[idx[:, :, "Stock"], :] *= 2
        tm.assert_frame_equal(df, expected)

        df = df_orig.copy()
        df.loc[idx[:, :, "Stock"], "price"] *= 2
        tm.assert_frame_equal(df, expected)

    def test_multiindex_assignment(self):

        # GH3777 part 2

        # mixed dtype
        df = DataFrame(
            np.random.randint(5, 10, size=9).reshape(3, 3),
            columns=list("abc"),
            index=[[4, 4, 8], [8, 10, 12]],
        )
        df["d"] = np.nan
        arr = np.array([0.0, 1.0])

        df.loc[4, "d"] = arr
        tm.assert_series_equal(df.loc[4, "d"], Series(arr, index=[8, 10], name="d"))

        # single dtype
        df = DataFrame(
            np.random.randint(5, 10, size=9).reshape(3, 3),
            columns=list("abc"),
            index=[[4, 4, 8], [8, 10, 12]],
        )

        df.loc[4, "c"] = arr
        exp = Series(arr, index=[8, 10], name="c", dtype="float64")
        tm.assert_series_equal(df.loc[4, "c"], exp)

        # scalar ok
        df.loc[4, "c"] = 10
        exp = Series(10, index=[8, 10], name="c", dtype="float64")
        tm.assert_series_equal(df.loc[4, "c"], exp)

        # invalid assignments
        msg = "Must have equal len keys and value when setting with an iterable"
        with pytest.raises(ValueError, match=msg):
            df.loc[4, "c"] = [0, 1, 2, 3]

        with pytest.raises(ValueError, match=msg):
            df.loc[4, "c"] = [0]

        # But with a length-1 listlike column indexer this behaves like
        #  `df.loc[4, "c"] = 0
        df.loc[4, ["c"]] = [0]
        assert (df.loc[4, "c"] == 0).all()

    def test_groupby_example(self):
        # groupby example
        NUM_ROWS = 100
        NUM_COLS = 10
        col_names = ["A" + num for num in map(str, np.arange(NUM_COLS).tolist())]
        index_cols = col_names[:5]

        df = DataFrame(
            np.random.randint(5, size=(NUM_ROWS, NUM_COLS)),
            dtype=np.int64,
            columns=col_names,
        )
        df = df.set_index(index_cols).sort_index()
        grp = df.groupby(level=index_cols[:4])
        df["new_col"] = np.nan

        f_index = np.arange(5)

        def f(name, df2):
            return Series(np.arange(df2.shape[0]), name=df2.index.values[0]).reindex(
                f_index
            )

        # FIXME: dont leave commented-out
        # TODO(wesm): unused?
        # new_df = pd.concat([f(name, df2) for name, df2 in grp], axis=1).T

        # we are actually operating on a copy here
        # but in this case, that's ok
        for name, df2 in grp:
            new_vals = np.arange(df2.shape[0])
            df.loc[name, "new_col"] = new_vals

    def test_series_setitem(self, multiindex_year_month_day_dataframe_random_data):
        ymd = multiindex_year_month_day_dataframe_random_data
        s = ymd["A"]

        s[2000, 3] = np.nan
        assert isna(s.values[42:65]).all()
        assert notna(s.values[:42]).all()
        assert notna(s.values[65:]).all()

        s[2000, 3, 10] = np.nan
        assert isna(s.iloc[49])

        with pytest.raises(KeyError, match="49"):
            # GH#33355 dont fall-back to positional when leading level is int
            s[49]

    def test_frame_getitem_setitem_boolean(self, multiindex_dataframe_random_data):
        frame = multiindex_dataframe_random_data
        df = frame.T.copy()
        values = df.values

        result = df[df > 0]
        expected = df.where(df > 0)
        tm.assert_frame_equal(result, expected)

        df[df > 0] = 5
        values[values > 0] = 5
        tm.assert_almost_equal(df.values, values)

        df[df == 5] = 0
        values[values == 5] = 0
        tm.assert_almost_equal(df.values, values)

        # a df that needs alignment first
        df[df[:-1] < 0] = 2
        np.putmask(values[:-1], values[:-1] < 0, 2)
        tm.assert_almost_equal(df.values, values)

        with pytest.raises(TypeError, match="boolean values only"):
            df[df * 0] = 2

    def test_frame_getitem_setitem_multislice(self):
        levels = [["t1", "t2"], ["a", "b", "c"]]
        codes = [[0, 0, 0, 1, 1], [0, 1, 2, 0, 1]]
        midx = MultiIndex(codes=codes, levels=levels, names=[None, "id"])
        df = DataFrame({"value": [1, 2, 3, 7, 8]}, index=midx)

        result = df.loc[:, "value"]
        tm.assert_series_equal(df["value"], result)

        result = df.loc[df.index[1:3], "value"]
        tm.assert_series_equal(df["value"][1:3], result)

        result = df.loc[:, :]
        tm.assert_frame_equal(df, result)

        result = df
        df.loc[:, "value"] = 10
        result["value"] = 10
        tm.assert_frame_equal(df, result)

        df.loc[:, :] = 10
        tm.assert_frame_equal(df, result)

    def test_frame_setitem_multi_column(self):
        df = DataFrame(
            np.random.randn(10, 4), columns=[["a", "a", "b", "b"], [0, 1, 0, 1]]
        )

        cp = df.copy()
        cp["a"] = cp["b"]
        tm.assert_frame_equal(cp["a"], cp["b"])

        # set with ndarray
        cp = df.copy()
        cp["a"] = cp["b"].values
        tm.assert_frame_equal(cp["a"], cp["b"])

        # ---------------------------------------
        # #1803
        columns = MultiIndex.from_tuples([("A", "1"), ("A", "2"), ("B", "1")])
        df = DataFrame(index=[1, 3, 5], columns=columns)

        # Works, but adds a column instead of updating the two existing ones
        df["A"] = 0.0  # Doesn't work
        assert (df["A"].values == 0).all()

        # it broadcasts
        df["B", "1"] = [1, 2, 3]
        df["A"] = df["B", "1"]

        sliced_a1 = df["A", "1"]
        sliced_a2 = df["A", "2"]
        sliced_b1 = df["B", "1"]
        tm.assert_series_equal(sliced_a1, sliced_b1, check_names=False)
        tm.assert_series_equal(sliced_a2, sliced_b1, check_names=False)
        assert sliced_a1.name == ("A", "1")
        assert sliced_a2.name == ("A", "2")
        assert sliced_b1.name == ("B", "1")

    def test_getitem_setitem_tuple_plus_columns(
        self, multiindex_year_month_day_dataframe_random_data
    ):
        # GH #1013
        ymd = multiindex_year_month_day_dataframe_random_data
        df = ymd[:5]

        result = df.loc[(2000, 1, 6), ["A", "B", "C"]]
        expected = df.loc[2000, 1, 6][["A", "B", "C"]]
        tm.assert_series_equal(result, expected)

    def test_getitem_setitem_slice_integers(self):
        index = MultiIndex(
            levels=[[0, 1, 2], [0, 2]], codes=[[0, 0, 1, 1, 2, 2], [0, 1, 0, 1, 0, 1]]
        )

        frame = DataFrame(
            np.random.randn(len(index), 4), index=index, columns=["a", "b", "c", "d"]
        )
        res = frame.loc[1:2]
        exp = frame.reindex(frame.index[2:])
        tm.assert_frame_equal(res, exp)

        frame.loc[1:2] = 7
        assert (frame.loc[1:2] == 7).values.all()

        series = Series(np.random.randn(len(index)), index=index)

        res = series.loc[1:2]
        exp = series.reindex(series.index[2:])
        tm.assert_series_equal(res, exp)

        series.loc[1:2] = 7
        assert (series.loc[1:2] == 7).values.all()

    def test_setitem_change_dtype(self, multiindex_dataframe_random_data):
        frame = multiindex_dataframe_random_data
        dft = frame.T
        s = dft["foo", "two"]
        dft["foo", "two"] = s > s.median()
        tm.assert_series_equal(dft["foo", "two"], s > s.median())
        # assert isinstance(dft._data.blocks[1].items, MultiIndex)

        reindexed = dft.reindex(columns=[("foo", "two")])
        tm.assert_series_equal(reindexed["foo", "two"], s > s.median())

    def test_set_column_scalar_with_loc(self, multiindex_dataframe_random_data):
        frame = multiindex_dataframe_random_data
        subset = frame.index[[1, 4, 5]]

        frame.loc[subset] = 99
        assert (frame.loc[subset].values == 99).all()

        col = frame["B"]
        col[subset] = 97
        assert (frame.loc[subset, "B"] == 97).all()

    def test_nonunique_assignment_1750(self):
        df = DataFrame(
            [[1, 1, "x", "X"], [1, 1, "y", "Y"], [1, 2, "z", "Z"]], columns=list("ABCD")
        )

        df = df.set_index(["A", "B"])
        ix = MultiIndex.from_tuples([(1, 1)])

        df.loc[ix, "C"] = "_"

        assert (df.xs((1, 1))["C"] == "_").all()

    def test_astype_assignment_with_dups(self):

        # GH 4686
        # assignment with dups that has a dtype change
        cols = MultiIndex.from_tuples([("A", "1"), ("B", "1"), ("A", "2")])
        df = DataFrame(np.arange(3).reshape((1, 3)), columns=cols, dtype=object)
        index = df.index.copy()

        df["A"] = df["A"].astype(np.float64)
        tm.assert_index_equal(df.index, index)

    def test_setitem_nonmonotonic(self):
        # https://github.com/pandas-dev/pandas/issues/31449
        index = MultiIndex.from_tuples(
            [("a", "c"), ("b", "x"), ("a", "d")], names=["l1", "l2"]
        )
        df = DataFrame(data=[0, 1, 2], index=index, columns=["e"])
        df.loc["a", "e"] = np.arange(99, 101, dtype="int64")
        expected = DataFrame({"e": [99, 1, 100]}, index=index)
        tm.assert_frame_equal(df, expected)


def test_frame_setitem_view_direct(multiindex_dataframe_random_data):
    # this works because we are modifying the underlying array
    # really a no-no
    df = multiindex_dataframe_random_data.T
    df["foo"].values[:] = 0
    assert (df["foo"].values == 0).all()


def test_frame_setitem_copy_raises(multiindex_dataframe_random_data):
    # will raise/warn as its chained assignment
    df = multiindex_dataframe_random_data.T
    msg = "A value is trying to be set on a copy of a slice from a DataFrame"
    with pytest.raises(com.SettingWithCopyError, match=msg):
        df["foo"]["one"] = 2


def test_frame_setitem_copy_no_write(multiindex_dataframe_random_data):
    frame = multiindex_dataframe_random_data.T
    expected = frame
    df = frame.copy()
    msg = "A value is trying to be set on a copy of a slice from a DataFrame"
    with pytest.raises(com.SettingWithCopyError, match=msg):
        df["foo"]["one"] = 2

    result = df
    tm.assert_frame_equal(result, expected)