#
# 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 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 'Reshaping, sorting, transposing'
# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#reshaping-sorting-transposing
class FrameReshapingMixin:
    @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 df_pair(self):
        pdf = self.pdf
        psdf = ps.from_pandas(pdf)
        return pdf, psdf

    def test_sort_values(self):
        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1]}, index=np.random.rand(7)
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(psdf.sort_values("b"), pdf.sort_values("b"))
        self.assert_eq(
            psdf.sort_values("b", ignore_index=True), pdf.sort_values("b", ignore_index=True)
        )

        for ascending in [True, False]:
            for na_position in ["first", "last"]:
                self.assert_eq(
                    psdf.sort_values("a", ascending=ascending, na_position=na_position),
                    pdf.sort_values("a", ascending=ascending, na_position=na_position),
                )

        self.assert_eq(psdf.sort_values(["a", "b"]), pdf.sort_values(["a", "b"]))
        self.assert_eq(
            psdf.sort_values(["a", "b"], ignore_index=True),
            pdf.sort_values(["a", "b"], ignore_index=True),
        )
        self.assert_eq(
            psdf.sort_values(["a", "b"], ascending=[False, True]),
            pdf.sort_values(["a", "b"], ascending=[False, True]),
        )

        self.assertRaises(ValueError, lambda: psdf.sort_values(["b", "a"], ascending=[False]))

        self.assert_eq(
            psdf.sort_values(["a", "b"], na_position="first"),
            pdf.sort_values(["a", "b"], na_position="first"),
        )

        self.assertRaises(ValueError, lambda: psdf.sort_values(["b", "a"], na_position="invalid"))

        pserA = pdf.a
        psserA = psdf.a
        self.assert_eq(psdf.sort_values("b", inplace=True), pdf.sort_values("b", inplace=True))
        self.assert_eq(psdf, pdf)
        self.assert_eq(psserA, pserA)

        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1]}, index=np.random.rand(7)
        )
        psdf = ps.from_pandas(pdf)
        pserA = pdf.a
        psserA = psdf.a
        self.assert_eq(
            psdf.sort_values("b", inplace=True, ignore_index=True),
            pdf.sort_values("b", inplace=True, ignore_index=True),
        )
        self.assert_eq(psdf, pdf)
        self.assert_eq(psserA, pserA)

        # multi-index indexes

        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1]},
            index=pd.MultiIndex.from_tuples(
                [
                    ("bar", "one"),
                    ("bar", "two"),
                    ("baz", "one"),
                    ("baz", "two"),
                    ("foo", "one"),
                    ("foo", "two"),
                    ("qux", "one"),
                ]
            ),
        )
        psdf = ps.from_pandas(pdf)
        self.assert_eq(psdf.sort_values("b"), pdf.sort_values("b"))
        self.assert_eq(
            psdf.sort_values("b", ignore_index=True), pdf.sort_values("b", ignore_index=True)
        )

        # multi-index columns
        pdf = pd.DataFrame(
            {("X", 10): [1, 2, 3, 4, 5, None, 7], ("X", 20): [7, 6, 5, 4, 3, 2, 1]},
            index=np.random.rand(7),
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(psdf.sort_values(("X", 20)), pdf.sort_values(("X", 20)))
        self.assert_eq(
            psdf.sort_values([("X", 20), ("X", 10)]), pdf.sort_values([("X", 20), ("X", 10)])
        )

        self.assertRaisesRegex(
            ValueError,
            "For a multi-index, the label must be a tuple with elements",
            lambda: psdf.sort_values(["X"]),
        )

        # non-string names
        pdf = pd.DataFrame(
            {10: [1, 2, 3, 4, 5, None, 7], 20: [7, 6, 5, 4, 3, 2, 1]}, index=np.random.rand(7)
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(psdf.sort_values(20), pdf.sort_values(20))
        self.assert_eq(psdf.sort_values([20, 10]), pdf.sort_values([20, 10]))

    def test_sort_index(self):
        pdf = pd.DataFrame(
            {"A": [2, 1, np.nan], "B": [np.nan, 0, np.nan]}, index=["b", "a", np.nan]
        )
        psdf = ps.from_pandas(pdf)

        # Assert invalid parameters
        self.assertRaises(NotImplementedError, lambda: psdf.sort_index(axis=1))
        self.assertRaises(NotImplementedError, lambda: psdf.sort_index(kind="mergesort"))
        self.assertRaises(ValueError, lambda: psdf.sort_index(na_position="invalid"))

        # Assert default behavior without parameters
        self.assert_eq(psdf.sort_index(), pdf.sort_index())
        # Assert ignoring index
        self.assert_eq(psdf.sort_index(ignore_index=True), pdf.sort_index(ignore_index=True))
        # Assert sorting descending
        self.assert_eq(psdf.sort_index(ascending=False), pdf.sort_index(ascending=False))
        # Assert sorting NA indices first
        self.assert_eq(psdf.sort_index(na_position="first"), pdf.sort_index(na_position="first"))
        # Assert sorting descending and NA indices first
        self.assert_eq(
            psdf.sort_index(ascending=False, na_position="first"),
            pdf.sort_index(ascending=False, na_position="first"),
        )

        # Assert sorting inplace
        pserA = pdf.A
        psserA = psdf.A
        self.assertEqual(psdf.sort_index(inplace=True), pdf.sort_index(inplace=True))
        self.assert_eq(psdf, pdf)
        self.assert_eq(psserA, pserA)
        pserA = pdf.A
        psserA = psdf.A
        self.assertEqual(
            psdf.sort_index(inplace=True, ascending=False, ignore_index=True),
            pdf.sort_index(inplace=True, ascending=False, ignore_index=True),
        )
        self.assert_eq(psdf, pdf)
        self.assert_eq(psserA, pserA)

        # Assert multi-indices
        pdf = pd.DataFrame(
            {"A": range(4), "B": range(4)[::-1]}, index=[["b", "b", "a", "a"], [1, 0, 1, 0]]
        )
        psdf = ps.from_pandas(pdf)
        self.assert_eq(psdf.sort_index(), pdf.sort_index())
        self.assert_eq(psdf.sort_index(level=[1, 0]), pdf.sort_index(level=[1, 0]))
        self.assert_eq(psdf.reset_index().sort_index(), pdf.reset_index().sort_index())
        # Assert ignoring index
        self.assert_eq(psdf.sort_index(ignore_index=True), pdf.sort_index(ignore_index=True))

        # Assert with multi-index columns
        columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B")])
        pdf.columns = columns
        psdf.columns = columns

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

    def test_nlargest(self):
        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1], "c": [1, 1, 2, 2, 3, 3, 3]},
            index=np.random.rand(7),
        )
        psdf = ps.from_pandas(pdf)
        self.assert_eq(psdf.nlargest(5, columns="a"), pdf.nlargest(5, columns="a"))
        self.assert_eq(psdf.nlargest(5, columns=["a", "b"]), pdf.nlargest(5, columns=["a", "b"]))
        self.assert_eq(psdf.nlargest(5, columns=["c"]), pdf.nlargest(5, columns=["c"]))
        self.assert_eq(
            psdf.nlargest(5, columns=["c"], keep="first"),
            pdf.nlargest(5, columns=["c"], keep="first"),
        )
        self.assert_eq(
            psdf.nlargest(5, columns=["c"], keep="last"),
            pdf.nlargest(5, columns=["c"], keep="last"),
        )
        msg = "`keep`=all is not implemented yet."
        with self.assertRaisesRegex(NotImplementedError, msg):
            psdf.nlargest(5, columns=["c"], keep="all")
        msg = 'keep must be either "first", "last" or "all".'
        with self.assertRaisesRegex(ValueError, msg):
            psdf.nlargest(5, columns=["c"], keep="xx")

    def test_nsmallest(self):
        pdf = pd.DataFrame(
            {"a": [1, 2, 3, 4, 5, None, 7], "b": [7, 6, 5, 4, 3, 2, 1], "c": [1, 1, 2, 2, 3, 3, 3]},
            index=np.random.rand(7),
        )
        psdf = ps.from_pandas(pdf)
        self.assert_eq(psdf.nsmallest(n=5, columns="a"), pdf.nsmallest(5, columns="a"))
        self.assert_eq(
            psdf.nsmallest(n=5, columns=["a", "b"]), pdf.nsmallest(5, columns=["a", "b"])
        )
        self.assert_eq(psdf.nsmallest(n=5, columns=["c"]), pdf.nsmallest(5, columns=["c"]))
        self.assert_eq(
            psdf.nsmallest(n=5, columns=["c"], keep="first"),
            pdf.nsmallest(5, columns=["c"], keep="first"),
        )
        self.assert_eq(
            psdf.nsmallest(n=5, columns=["c"], keep="last"),
            pdf.nsmallest(5, columns=["c"], keep="last"),
        )
        msg = "`keep`=all is not implemented yet."
        with self.assertRaisesRegex(NotImplementedError, msg):
            psdf.nlargest(5, columns=["c"], keep="all")
        msg = 'keep must be either "first", "last" or "all".'
        with self.assertRaisesRegex(ValueError, msg):
            psdf.nlargest(5, columns=["c"], keep="xx")

    def test_stack(self):
        pdf_single_level_cols = pd.DataFrame(
            [[0, 1], [2, 3]], index=["cat", "dog"], columns=["weight", "height"]
        )
        psdf_single_level_cols = ps.from_pandas(pdf_single_level_cols)

        self.assert_eq(
            psdf_single_level_cols.stack().sort_index(), pdf_single_level_cols.stack().sort_index()
        )

        multicol1 = pd.MultiIndex.from_tuples(
            [("weight", "kg"), ("weight", "pounds")], names=["x", "y"]
        )
        pdf_multi_level_cols1 = pd.DataFrame(
            [[1, 2], [2, 4]], index=["cat", "dog"], columns=multicol1
        )
        psdf_multi_level_cols1 = ps.from_pandas(pdf_multi_level_cols1)

        self.assert_eq(
            psdf_multi_level_cols1.stack().sort_index(), pdf_multi_level_cols1.stack().sort_index()
        )

        multicol2 = pd.MultiIndex.from_tuples([("weight", "kg"), ("height", "m")])
        pdf_multi_level_cols2 = pd.DataFrame(
            [[1.0, 2.0], [3.0, 4.0]], index=["cat", "dog"], columns=multicol2
        )
        psdf_multi_level_cols2 = ps.from_pandas(pdf_multi_level_cols2)

        self.assert_eq(
            psdf_multi_level_cols2.stack().sort_index(), pdf_multi_level_cols2.stack().sort_index()
        )

        pdf = pd.DataFrame(
            {
                ("y", "c"): [True, True],
                ("x", "b"): [False, False],
                ("x", "c"): [True, False],
                ("y", "a"): [False, True],
            }
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(psdf.stack().sort_index(), pdf.stack().sort_index())
        self.assert_eq(psdf[[]].stack().sort_index(), pdf[[]].stack().sort_index(), almost=True)

    def test_unstack(self):
        pdf = pd.DataFrame(
            np.random.randn(3, 3),
            index=pd.MultiIndex.from_tuples([("rg1", "x"), ("rg1", "y"), ("rg2", "z")]),
        )
        psdf = ps.from_pandas(pdf)

        self.assert_eq(psdf.unstack().sort_index(), pdf.unstack().sort_index(), almost=True)
        self.assert_eq(
            psdf.unstack().unstack().sort_index(), pdf.unstack().unstack().sort_index(), almost=True
        )

    def test_explode(self):
        pdf = pd.DataFrame(
            {"A": [[-1.0, np.nan], [0.0, np.inf], [1.0, -np.inf]], "B": 1}, index=["a", "b", "c"]
        )
        pdf.index.name = "index"
        pdf.columns.name = "columns"
        psdf = ps.from_pandas(pdf)

        expected_result1, result1 = pdf.explode("A"), psdf.explode("A")
        expected_result2, result2 = pdf.explode("B"), psdf.explode("B")
        expected_result3, result3 = (
            pdf.explode("A", ignore_index=True),
            psdf.explode("A", ignore_index=True),
        )

        self.assert_eq(result1, expected_result1, almost=True)
        self.assert_eq(result2, expected_result2, check_exact=False)
        self.assert_eq(result1.index.name, expected_result1.index.name)
        self.assert_eq(result1.columns.name, expected_result1.columns.name)
        self.assert_eq(result3, expected_result3, almost=True)
        self.assert_eq(result3.index, expected_result3.index)

        self.assertRaises(TypeError, lambda: psdf.explode(["A", "B"]))

        # MultiIndex
        midx = pd.MultiIndex.from_tuples(
            [("x", "a"), ("x", "b"), ("y", "c")], names=["index1", "index2"]
        )
        pdf.index = midx
        psdf = ps.from_pandas(pdf)

        expected_result1, result1 = pdf.explode("A"), psdf.explode("A")
        expected_result2, result2 = pdf.explode("B"), psdf.explode("B")
        expected_result3, result3 = (
            pdf.explode("A", ignore_index=True),
            psdf.explode("A", ignore_index=True),
        )

        self.assert_eq(result1, expected_result1, almost=True)
        self.assert_eq(result2, expected_result2, check_exact=False)
        self.assert_eq(result1.index.names, expected_result1.index.names)
        self.assert_eq(result1.columns.name, expected_result1.columns.name)
        self.assert_eq(result3, expected_result3, almost=True)
        self.assert_eq(result3.index, expected_result3.index)

        self.assertRaises(TypeError, lambda: psdf.explode(["A", "B"]))

        # MultiIndex columns
        columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X")], names=["column1", "column2"])
        pdf.columns = columns
        psdf.columns = columns

        expected_result1, result1 = pdf.explode(("A", "Z")), psdf.explode(("A", "Z"))
        expected_result2, result2 = pdf.explode(("B", "X")), psdf.explode(("B", "X"))
        expected_result3, result3 = pdf.A.explode("Z"), psdf.A.explode("Z")

        self.assert_eq(result1, expected_result1, almost=True)
        self.assert_eq(result2, expected_result2, check_exact=False)
        self.assert_eq(result1.index.names, expected_result1.index.names)
        self.assert_eq(result1.columns.names, expected_result1.columns.names)
        self.assert_eq(result3, expected_result3, almost=True)

        self.assertRaises(TypeError, lambda: psdf.explode(["A", "B"]))
        self.assertRaises(ValueError, lambda: psdf.explode("A"))

    def test_transpose(self):
        # TODO: what if with random index?
        pdf1 = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}, columns=["col1", "col2"])
        psdf1 = ps.from_pandas(pdf1)

        pdf2 = pd.DataFrame(
            data={"score": [9, 8], "kids": [0, 0], "age": [12, 22]},
            columns=["score", "kids", "age"],
        )
        psdf2 = ps.from_pandas(pdf2)

        self.assert_eq(pdf1.transpose().sort_index(), psdf1.transpose().sort_index())
        self.assert_eq(pdf2.transpose().sort_index(), psdf2.transpose().sort_index())

        with option_context("compute.max_rows", None):
            self.assert_eq(pdf1.transpose().sort_index(), psdf1.transpose().sort_index())

            self.assert_eq(pdf2.transpose().sort_index(), psdf2.transpose().sort_index())

        pdf3 = pd.DataFrame(
            {
                ("cg1", "a"): [1, 2, 3],
                ("cg1", "b"): [4, 5, 6],
                ("cg2", "c"): [7, 8, 9],
                ("cg3", "d"): [9, 9, 9],
            },
            index=pd.MultiIndex.from_tuples([("rg1", "x"), ("rg1", "y"), ("rg2", "z")]),
        )
        psdf3 = ps.from_pandas(pdf3)

        self.assert_eq(pdf3.transpose().sort_index(), psdf3.transpose().sort_index())

        with option_context("compute.max_rows", None):
            self.assert_eq(pdf3.transpose().sort_index(), psdf3.transpose().sort_index())

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

        pser = pdf.a
        psser = psdf.a

        pdf["x"] = [10, 20, 30, 40, 50, 60, 70, 80, 90]
        psdf["x"] = [10, 20, 30, 40, 50, 60, 70, 80, 90]

        self.assert_eq(psdf.sort_index(), pdf.sort_index())
        self.assert_eq(psser, pser)

        with self.assertRaisesRegex(ValueError, "Length of values does not match length of index"):
            psdf["z"] = [10, 20, 30, 40, 50, 60, 70, 80]

    def test_squeeze(self):
        axises = [None, 0, 1, "rows", "index", "columns"]

        # Multiple columns
        pdf = pd.DataFrame([[1, 2], [3, 4]], columns=["a", "b"], index=["x", "y"])
        psdf = ps.from_pandas(pdf)
        for axis in axises:
            self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis))
        # Multiple columns with MultiIndex columns
        columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X")])
        pdf.columns = columns
        psdf.columns = columns
        for axis in axises:
            self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis))

        # Single column with single value
        pdf = pd.DataFrame([[1]], columns=["a"], index=["x"])
        psdf = ps.from_pandas(pdf)
        for axis in axises:
            self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis))
        # Single column with single value with MultiIndex column
        columns = pd.MultiIndex.from_tuples([("A", "Z")])
        pdf.columns = columns
        psdf.columns = columns
        for axis in axises:
            self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis))

        # Single column with multiple values
        pdf = pd.DataFrame([1, 2, 3, 4], columns=["a"])
        psdf = ps.from_pandas(pdf)
        for axis in axises:
            self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis))
        # Single column with multiple values with MultiIndex column
        pdf.columns = columns
        psdf.columns = columns
        for axis in axises:
            self.assert_eq(pdf.squeeze(axis), psdf.squeeze(axis))


class FrameReshapingTests(
    FrameReshapingMixin,
    PandasOnSparkTestCase,
    SQLTestUtils,
):
    pass


if __name__ == "__main__":
    from pyspark.testing import main

    main()
