#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import unittest

import numpy as np
import pandas as pd

from pyspark import pandas as ps
from pyspark.pandas.config import option_context
from pyspark.testing.pandasutils import PandasOnSparkTestCase
from pyspark.testing.sqlutils import SQLTestUtils


# This file contains test cases for 'Indexing, Iteration'
# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#indexing-iteration
class FrameIndexingMixin:
    @property
    def pdf(self):
        return pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]},
            index=np.random.rand(9),
        )

    @property
    def psdf(self):
        return ps.from_pandas(self.pdf)

    @property
    def df_pair(self):
        pdf = self.pdf
        psdf = ps.from_pandas(pdf)
        return pdf, psdf

    def test_head(self):
        pdf, psdf = self.df_pair

        self.assert_eq(psdf.head(2), pdf.head(2))
        self.assert_eq(psdf.head(3), pdf.head(3))
        self.assert_eq(psdf.head(0), pdf.head(0))
        self.assert_eq(psdf.head(-3), pdf.head(-3))
        self.assert_eq(psdf.head(-10), pdf.head(-10))
        with option_context("compute.ordered_head", True):
            self.assert_eq(psdf.head(), pdf.head())

    def test_items(self):
        pdf = pd.DataFrame(
            {"species": ["bear", "bear", "marsupial"], "population": [1864, 22000, 80000]},
            index=["panda", "polar", "koala"],
            columns=["species", "population"],
        )
        psdf = ps.from_pandas(pdf)

        for (p_name, p_items), (k_name, k_items) in zip(pdf.items(), psdf.items()):
            self.assert_eq(p_name, k_name)
            self.assert_eq(p_items, k_items)

    def test_keys(self):
        pdf = pd.DataFrame(
            [[1, 2], [4, 5], [7, 8]],
            index=["cobra", "viper", "sidewinder"],
            columns=["max_speed", "shield"],
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(psdf.keys(), pdf.keys())

    def test_tail(self):
        pdf = pd.DataFrame({"x": range(1000)})
        psdf = ps.from_pandas(pdf)

        self.assert_eq(pdf.tail(), psdf.tail())
        self.assert_eq(pdf.tail(10), psdf.tail(10))
        self.assert_eq(pdf.tail(-990), psdf.tail(-990))
        self.assert_eq(pdf.tail(0), psdf.tail(0))
        self.assert_eq(pdf.tail(-1001), psdf.tail(-1001))
        self.assert_eq(pdf.tail(1001), psdf.tail(1001))
        self.assert_eq((pdf + 1).tail(), (psdf + 1).tail())
        self.assert_eq((pdf + 1).tail(10), (psdf + 1).tail(10))
        self.assert_eq((pdf + 1).tail(-990), (psdf + 1).tail(-990))
        self.assert_eq((pdf + 1).tail(0), (psdf + 1).tail(0))
        self.assert_eq((pdf + 1).tail(-1001), (psdf + 1).tail(-1001))
        self.assert_eq((pdf + 1).tail(1001), (psdf + 1).tail(1001))
        with self.assertRaisesRegex(TypeError, "bad operand type for unary -: 'str'"):
            psdf.tail("10")

    def test_xs(self):
        d = {
            "num_legs": [4, 4, 2, 2],
            "num_wings": [0, 0, 2, 2],
            "class": ["mammal", "mammal", "mammal", "bird"],
            "animal": ["cat", "dog", "bat", "penguin"],
            "locomotion": ["walks", "walks", "flies", "walks"],
        }
        pdf = pd.DataFrame(data=d)
        pdf = pdf.set_index(["class", "animal", "locomotion"])
        psdf = ps.from_pandas(pdf)

        self.assert_eq(psdf.xs("mammal"), pdf.xs("mammal"))
        self.assert_eq(psdf.xs(("mammal",)), pdf.xs(("mammal",)))
        self.assert_eq(psdf.xs(("mammal", "dog", "walks")), pdf.xs(("mammal", "dog", "walks")))
        self.assert_eq(
            ps.concat([psdf, psdf]).xs(("mammal", "dog", "walks")),
            pd.concat([pdf, pdf]).xs(("mammal", "dog", "walks")),
        )
        self.assert_eq(psdf.xs("cat", level=1), pdf.xs("cat", level=1))
        self.assert_eq(psdf.xs("flies", level=2), pdf.xs("flies", level=2))
        self.assert_eq(psdf.xs("mammal", level=-3), pdf.xs("mammal", level=-3))

        msg = 'axis should be either 0 or "index" currently.'
        with self.assertRaisesRegex(NotImplementedError, msg):
            psdf.xs("num_wings", axis=1)
        with self.assertRaises(KeyError):
            psdf.xs(("mammal", "dog", "walk"))
        msg = r"'Key length \(4\) exceeds index depth \(3\)'"
        with self.assertRaisesRegex(KeyError, msg):
            psdf.xs(("mammal", "dog", "walks", "foo"))
        msg = "'key' should be a scalar value or tuple that contains scalar values"
        with self.assertRaisesRegex(TypeError, msg):
            psdf.xs(["mammal", "dog", "walks", "foo"])

        self.assertRaises(IndexError, lambda: psdf.xs("foo", level=-4))
        self.assertRaises(IndexError, lambda: psdf.xs("foo", level=3))

        self.assertRaises(KeyError, lambda: psdf.xs(("dog", "walks"), level=1))

        # non-string names
        pdf = pd.DataFrame(data=d)
        pdf = pdf.set_index(["class", "animal", "num_legs", "num_wings"])
        psdf = ps.from_pandas(pdf)

        self.assert_eq(psdf.xs(("mammal", "dog", 4)), pdf.xs(("mammal", "dog", 4)))
        self.assert_eq(psdf.xs(2, level=2), pdf.xs(2, level=2))

        self.assert_eq((psdf + "a").xs(("mammal", "dog", 4)), (pdf + "a").xs(("mammal", "dog", 4)))
        self.assert_eq((psdf + "a").xs(2, level=2), (pdf + "a").xs(2, level=2))

    def test_where(self):
        pdf, psdf = self.df_pair

        # pandas requires `axis` argument when the `other` is Series.
        # `axis` is not fully supported yet in pandas-on-Spark.
        self.assert_eq(
            psdf.where(psdf > 2, psdf.a + 10, axis=0), pdf.where(pdf > 2, pdf.a + 10, axis=0)
        )

        with self.assertRaisesRegex(TypeError, "type of cond must be a DataFrame or Series"):
            psdf.where(1)
        with self.assertRaisesRegex(
            NotImplementedError, 'axis should be either 0 or "index" currently.'
        ):
            psdf.where(psdf > 2, psdf.a + 10, axis=1)

    def test_mask(self):
        psdf = ps.from_pandas(self.pdf)

        with self.assertRaisesRegex(TypeError, "type of cond must be a DataFrame or Series"):
            psdf.mask(1)

    def test_query(self):
        pdf = pd.DataFrame({"A": range(1, 6), "B": range(10, 0, -2), "C": range(10, 5, -1)})
        psdf = ps.from_pandas(pdf)

        exprs = ("A > B", "A < C", "C == B")
        for expr in exprs:
            self.assert_eq(psdf.query(expr), pdf.query(expr))

        # test `inplace=True`
        for expr in exprs:
            dummy_psdf = psdf.copy()
            dummy_pdf = pdf.copy()

            pser = dummy_pdf.A
            psser = dummy_psdf.A
            dummy_pdf.query(expr, inplace=True)
            dummy_psdf.query(expr, inplace=True)

            self.assert_eq(dummy_psdf, dummy_pdf)
            self.assert_eq(psser, pser)

        # invalid values for `expr`
        invalid_exprs = (1, 1.0, (exprs[0],), [exprs[0]])
        for expr in invalid_exprs:
            with self.assertRaisesRegex(
                TypeError,
                "expr must be a string to be evaluated, {} given".format(type(expr).__name__),
            ):
                psdf.query(expr)

        # invalid values for `inplace`
        invalid_inplaces = (1, 0, "True", "False")
        for inplace in invalid_inplaces:
            with self.assertRaisesRegex(
                TypeError,
                'For argument "inplace" expected type bool, received type {}.'.format(
                    type(inplace).__name__
                ),
            ):
                psdf.query("a < b", inplace=inplace)

        # doesn't support for MultiIndex columns
        columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X"), ("C", "C")])
        psdf.columns = columns
        with self.assertRaisesRegex(TypeError, "Doesn't support for MultiIndex columns"):
            psdf.query("('A', 'Z') > ('B', 'X')")

    def test_insert(self):
        #
        # Basic DataFrame
        #
        pdf = pd.DataFrame([1, 2, 3])
        psdf = ps.from_pandas(pdf)

        psdf.insert(1, "b", 10)
        pdf.insert(1, "b", 10)
        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
        psdf.insert(2, "c", 0.1)
        pdf.insert(2, "c", 0.1)
        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
        psdf.insert(3, "d", psdf.b + 1)
        pdf.insert(3, "d", pdf.b + 1)
        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)

        psser = ps.Series([4, 5, 6])
        with ps.option_context("compute.ops_on_diff_frames", False):
            self.assertRaises(ValueError, lambda: psdf.insert(0, "y", psser))

        self.assertRaisesRegex(
            ValueError, "cannot insert b, already exists", lambda: psdf.insert(1, "b", 10)
        )
        self.assertRaisesRegex(
            TypeError,
            '"column" should be a scalar value or tuple that contains scalar values',
            lambda: psdf.insert(0, list("abc"), psser),
        )
        self.assertRaisesRegex(
            TypeError,
            "loc must be int",
            lambda: psdf.insert((1,), "b", 10),
        )
        self.assertRaisesRegex(
            NotImplementedError,
            "Assigning column name as tuple is only supported for MultiIndex columns for now.",
            lambda: psdf.insert(0, ("e",), 10),
        )

        self.assertRaises(ValueError, lambda: psdf.insert(0, "e", [7, 8, 9, 10]))
        with ps.option_context("compute.ops_on_diff_frames", False):
            self.assertRaises(ValueError, lambda: psdf.insert(0, "f", ps.Series([7, 8])))

        self.assertRaises(AssertionError, lambda: psdf.insert(100, "y", psser))
        self.assertRaises(AssertionError, lambda: psdf.insert(1, "y", psser, allow_duplicates=True))

        #
        # DataFrame with MultiIndex as columns
        #
        pdf = pd.DataFrame({("x", "a", "b"): [1, 2, 3]})
        psdf = ps.from_pandas(pdf)

        psdf.insert(1, "b", 10)
        pdf.insert(1, "b", 10)
        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
        psdf.insert(2, "c", 0.1)
        pdf.insert(2, "c", 0.1)
        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
        psdf.insert(3, "d", psdf.b + 1)
        pdf.insert(3, "d", pdf.b + 1)
        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)

        self.assertRaisesRegex(
            ValueError, "cannot insert d, already exists", lambda: psdf.insert(4, "d", 11)
        )
        self.assertRaisesRegex(
            ValueError,
            r"cannot insert \('x', 'a', 'b'\), already exists",
            lambda: psdf.insert(4, ("x", "a", "b"), 11),
        )
        self.assertRaisesRegex(
            ValueError,
            '"column" must have length equal to number of column levels.',
            lambda: psdf.insert(4, ("e",), 11),
        )

    def test_itertuples(self):
        pdf = pd.DataFrame({"num_legs": [4, 2], "num_wings": [0, 2]}, index=["dog", "hawk"])
        psdf = ps.from_pandas(pdf)

        for ptuple, ktuple in zip(
            pdf.itertuples(index=False, name="Animal"), psdf.itertuples(index=False, name="Animal")
        ):
            self.assert_eq(ptuple, ktuple)
        for ptuple, ktuple in zip(pdf.itertuples(name=None), psdf.itertuples(name=None)):
            self.assert_eq(ptuple, ktuple)
        for ptuple, ktuple in zip(
            pdf.itertuples(index=False, name=None), psdf.itertuples(index=False, name=None)
        ):
            self.assert_eq(ptuple, ktuple)

        pdf.index = pd.MultiIndex.from_arrays(
            [[1, 2], ["black", "brown"]], names=("count", "color")
        )
        psdf = ps.from_pandas(pdf)
        for ptuple, ktuple in zip(pdf.itertuples(name="Animal"), psdf.itertuples(name="Animal")):
            self.assert_eq(ptuple, ktuple)

        pdf.columns = pd.MultiIndex.from_arrays(
            [["CA", "WA"], ["age", "children"]], names=("origin", "info")
        )
        psdf = ps.from_pandas(pdf)
        for ptuple, ktuple in zip(pdf.itertuples(name="Animal"), psdf.itertuples(name="Animal")):
            self.assert_eq(ptuple, ktuple)

        pdf = pd.DataFrame([1, 2, 3])
        psdf = ps.from_pandas(pdf)
        for ptuple, ktuple in zip(
            (pdf + 1).itertuples(name="num"), (psdf + 1).itertuples(name="num")
        ):
            self.assert_eq(ptuple, ktuple)

        # DataFrames with a large number of columns (>254)
        pdf = pd.DataFrame(np.random.random((1, 255)))
        psdf = ps.from_pandas(pdf)
        for ptuple, ktuple in zip(pdf.itertuples(name="num"), psdf.itertuples(name="num")):
            self.assert_eq(ptuple, ktuple)

    def test_iterrows(self):
        pdf = pd.DataFrame(
            {
                ("x", "a", "1"): [1, 2, 3],
                ("x", "b", "2"): [4, 5, 6],
                ("y.z", "c.d", "3"): [7, 8, 9],
                ("x", "b", "4"): [10, 11, 12],
            },
            index=np.random.rand(3),
        )
        psdf = ps.from_pandas(pdf)

        for (pdf_k, pdf_v), (psdf_k, psdf_v) in zip(pdf.iterrows(), psdf.iterrows()):
            self.assert_eq(pdf_k, psdf_k)
            self.assert_eq(pdf_v, psdf_v)

        # MultiIndex
        pmidx = pd.Index([(1, 2), (3, 4), (5, 6)])
        pdf.index = pmidx
        psdf = ps.from_pandas(pdf)

        for (pdf_k, pdf_v), (psdf_k, psdf_v) in zip(pdf.iterrows(), psdf.iterrows()):
            self.assert_eq(pdf_k, psdf_k)
            self.assert_eq(pdf_v, psdf_v)

    def test_multiindex_column_access(self):
        columns = pd.MultiIndex.from_tuples(
            [
                ("a", "", "", "b"),
                ("c", "", "d", ""),
                ("e", "", "f", ""),
                ("e", "g", "", ""),
                ("", "", "", "h"),
                ("i", "", "", ""),
            ]
        )

        pdf = pd.DataFrame(
            [
                (1, "a", "x", 10, 100, 1000),
                (2, "b", "y", 20, 200, 2000),
                (3, "c", "z", 30, 300, 3000),
            ],
            columns=columns,
            index=np.random.rand(3),
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(psdf, pdf)
        self.assert_eq(psdf["a"], pdf["a"])
        self.assert_eq(psdf["a"]["b"], pdf["a"]["b"])
        self.assert_eq(psdf["c"], pdf["c"])
        self.assert_eq(psdf["c"]["d"], pdf["c"]["d"])
        self.assert_eq(psdf["e"], pdf["e"])
        self.assert_eq(psdf["e"][""]["f"], pdf["e"][""]["f"])
        self.assert_eq(psdf["e"]["g"], pdf["e"]["g"])
        self.assert_eq(psdf[""], pdf[""])
        self.assert_eq(psdf[""]["h"], pdf[""]["h"])
        self.assert_eq(psdf["i"], pdf["i"])

        self.assert_eq(psdf[["a", "e"]], pdf[["a", "e"]])
        self.assert_eq(psdf[["e", "a"]], pdf[["e", "a"]])

        self.assert_eq(psdf[("a",)], pdf[("a",)])
        self.assert_eq(psdf[("e", "g")], pdf[("e", "g")])
        # self.assert_eq(psdf[("i",)], pdf[("i",)])
        self.assert_eq(psdf[("i", "")], pdf[("i", "")])

        self.assertRaises(KeyError, lambda: psdf[("a", "b")])

    def test_getitem_with_none_key(self):
        psdf = self.psdf

        with self.assertRaisesRegex(KeyError, "none key"):
            psdf[None]

    def test_iter_dataframe(self):
        pdf, psdf = self.df_pair

        for value_psdf, value_pdf in zip(psdf, pdf):
            self.assert_eq(value_psdf, value_pdf)


class FrameIndexingTests(
    FrameIndexingMixin,
    PandasOnSparkTestCase,
    SQLTestUtils,
):
    pass


if __name__ == "__main__":
    from pyspark.pandas.tests.indexes.test_indexing import *  # noqa: F401

    try:
        import xmlrunner

        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
    except ImportError:
        testRunner = None
    unittest.main(testRunner=testRunner, verbosity=2)