#
# 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
from datetime import datetime

import pandas as pd

from pyspark import pandas as ps
from pyspark.testing.pandasutils import PandasOnSparkTestCase
from pyspark.pandas.utils import (
    SPARK_CONF_ARROW_ENABLED,
    SPARK_CONF_PANDAS_STRUCT_MODE,
)
from pyspark.testing.sqlutils import SQLTestUtils
from pyspark.testing.utils import is_ansi_mode_test


class ConversionMixin:
    @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=[0, 1, 3, 5, 6, 8, 9, 9, 9],
        )

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

    def test_index_from_series(self):
        pser = pd.Series([1, 2, 3], name="a", index=[10, 20, 30])
        psser = ps.from_pandas(pser)

        self.assert_eq(ps.Index(psser), pd.Index(pser))
        self.assert_eq(ps.Index(psser, dtype="float"), pd.Index(pser, dtype="float"))
        self.assert_eq(ps.Index(psser, name="x"), pd.Index(pser, name="x"))

        self.assert_eq(ps.Index(psser, dtype="int64"), pd.Index(pser, dtype="int64"))
        self.assert_eq(ps.Index(psser, dtype="float64"), pd.Index(pser, dtype="float64"))

        pser = pd.Series([datetime(2021, 3, 1), datetime(2021, 3, 2)], name="x", index=[10, 20])
        psser = ps.from_pandas(pser)

        self.assert_eq(ps.Index(psser), pd.Index(pser))
        self.assert_eq(ps.DatetimeIndex(psser), pd.DatetimeIndex(pser))

    def test_index_from_index(self):
        pidx = pd.Index([1, 2, 3], name="a")
        psidx = ps.from_pandas(pidx)

        self.assert_eq(ps.Index(psidx), pd.Index(pidx))
        self.assert_eq(ps.Index(psidx, dtype="float"), pd.Index(pidx, dtype="float"))
        self.assert_eq(ps.Index(psidx, name="x"), pd.Index(pidx, name="x"))
        self.assert_eq(ps.Index(psidx, copy=True), pd.Index(pidx, copy=True))

        self.assert_eq(ps.Index(psidx, dtype="int64"), pd.Index(pidx, dtype="int64"))
        self.assert_eq(ps.Index(psidx, dtype="float64"), pd.Index(pidx, dtype="float64"))

        pidx = pd.DatetimeIndex(["2021-03-01", "2021-03-02"])
        psidx = ps.from_pandas(pidx)

        self.assert_eq(ps.Index(psidx), pd.Index(pidx))
        self.assert_eq(ps.DatetimeIndex(psidx), pd.DatetimeIndex(pidx))

    def test_multiindex_from_arrays(self):
        arrays = [["a", "a", "b", "b"], ["red", "blue", "red", "blue"]]
        pidx = pd.MultiIndex.from_arrays(arrays)
        psidx = ps.MultiIndex.from_arrays(arrays)

        self.assert_eq(pidx, psidx)

    def test_multiindex_from_tuples(self):
        tuples = [(1, "red"), (1, "blue"), (2, "red"), (2, "blue")]
        pidx = pd.MultiIndex.from_tuples(tuples)
        psidx = ps.MultiIndex.from_tuples(tuples)

        self.assert_eq(pidx, psidx)

    def test_multiindex_from_product(self):
        iterables = [[0, 1, 2], ["green", "purple"]]
        pidx = pd.MultiIndex.from_product(iterables)
        psidx = ps.MultiIndex.from_product(iterables)

        self.assert_eq(pidx, psidx)

    def test_multi_index_from_index(self):
        tuples = [(1, "red"), (1, "blue"), (2, "red"), (2, "blue")]
        pmidx = pd.Index(tuples)
        psmidx = ps.Index(tuples)

        self.assertTrue(isinstance(psmidx, ps.MultiIndex))
        self.assert_eq(pmidx, psmidx)

        # Specify the `names`
        pmidx = pd.Index(tuples)
        pmidx.names = ["Hello", "Koalas"]
        psmidx = ps.Index(tuples)
        psmidx.names = ["Hello", "Koalas"]

        self.assertTrue(isinstance(psmidx, ps.MultiIndex))
        self.assert_eq(pmidx, psmidx)

    def test_multiindex_from_frame(self):
        pdf = pd.DataFrame(
            [["HI", "Temp"], ["HI", "Precip"], ["NJ", "Temp"], ["NJ", "Precip"]], columns=["a", "b"]
        )
        psdf = ps.from_pandas(pdf)
        pidx = pd.MultiIndex.from_frame(pdf)
        psidx = ps.MultiIndex.from_frame(psdf)

        self.assert_eq(pidx, psidx)

        # Specify `names`
        pidx = pd.MultiIndex.from_frame(pdf, names=["state", "observation"])
        psidx = ps.MultiIndex.from_frame(psdf, names=["state", "observation"])
        self.assert_eq(pidx, psidx)

        pidx = pd.MultiIndex.from_frame(pdf, names=("state", "observation"))
        psidx = ps.MultiIndex.from_frame(psdf, names=("state", "observation"))
        self.assert_eq(pidx, psidx)

        # MultiIndex columns
        pidx = pd.MultiIndex.from_tuples([("a", "w"), ("b", "x")])
        pdf.columns = pidx
        psdf = ps.from_pandas(pdf)

        pidx = pd.MultiIndex.from_frame(pdf)
        psidx = ps.MultiIndex.from_frame(psdf)

        self.assert_eq(pidx, psidx)

        # tuples for names
        pidx = pd.MultiIndex.from_frame(pdf, names=[("a", "w"), ("b", "x")])
        psidx = ps.MultiIndex.from_frame(psdf, names=[("a", "w"), ("b", "x")])

        self.assert_eq(pidx, psidx)

        err_msg = "Input must be a DataFrame"
        with self.assertRaisesRegex(TypeError, err_msg):
            ps.MultiIndex.from_frame({"a": [1, 2, 3], "b": [4, 5, 6]})

        self.assertRaises(TypeError, lambda: ps.MultiIndex.from_frame(psdf, names="ab"))

        # non-string names
        self.assert_eq(
            ps.MultiIndex.from_frame(psdf, names=[0, 1]),
            pd.MultiIndex.from_frame(pdf, names=[0, 1]),
        )
        self.assert_eq(
            ps.MultiIndex.from_frame(psdf, names=[("x", 0), ("y", 1)]),
            pd.MultiIndex.from_frame(pdf, names=[("x", 0), ("y", 1)]),
        )

        pdf = pd.DataFrame([["HI", "Temp"], ["HI", "Precip"], ["NJ", "Temp"], ["NJ", "Precip"]])
        psdf = ps.from_pandas(pdf)
        self.assert_eq(ps.MultiIndex.from_frame(psdf), pd.MultiIndex.from_frame(pdf))

    def test_to_series(self):
        pidx = self.pdf.index
        psidx = self.psdf.index

        self.assert_eq(psidx.to_series(), pidx.to_series())
        self.assert_eq(psidx.to_series(name="a"), pidx.to_series(name="a"))

        # With name
        pidx.name = "Koalas"
        psidx.name = "Koalas"
        self.assert_eq(psidx.to_series(), pidx.to_series())
        self.assert_eq(psidx.to_series(name=("x", "a")), pidx.to_series(name=("x", "a")))

        # With tupled name
        pidx.name = ("x", "a")
        psidx.name = ("x", "a")
        self.assert_eq(psidx.to_series(), pidx.to_series())
        self.assert_eq(psidx.to_series(name="a"), pidx.to_series(name="a"))

        self.assert_eq((psidx + 1).to_series(), (pidx + 1).to_series())

        # Multiindex
        arrays = [[1, 2], ["red", "blue"]]
        pidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color"))
        psidx = ps.from_pandas(pidx)

        if is_ansi_mode_test:
            with self.sql_conf(
                {
                    SPARK_CONF_PANDAS_STRUCT_MODE: "row",
                }
            ):
                self.assert_eq(
                    list(psidx.to_series().values),
                    list(pidx.to_series().values),
                )
        else:
            self.assert_eq(list(psidx.to_series().values), [["1", "red"], ["2", "blue"]])

        pidx = self.pdf.set_index("b", append=True).index
        psidx = self.psdf.set_index("b", append=True).index

        with self.sql_conf({SPARK_CONF_ARROW_ENABLED: False, SPARK_CONF_PANDAS_STRUCT_MODE: "row"}):
            self.assert_eq(psidx.to_series(), pidx.to_series(), check_exact=False)
            self.assert_eq(psidx.to_series(name="a"), pidx.to_series(name="a"), check_exact=False)

        expected_error_message = "Series.name must be a hashable type"
        with self.assertRaisesRegex(TypeError, expected_error_message):
            psidx.to_series(name=["x", "a"])

    def test_to_frame(self):
        pidx = self.pdf.index
        psidx = self.psdf.index

        self.assert_eq(psidx.to_frame(), pidx.to_frame())
        self.assert_eq(psidx.to_frame(index=False), pidx.to_frame(index=False))

        pidx.name = "a"
        psidx.name = "a"

        self.assert_eq(psidx.to_frame(), pidx.to_frame())
        self.assert_eq(psidx.to_frame(index=False), pidx.to_frame(index=False))

        self.assert_eq(psidx.to_frame(name="x"), pidx.to_frame(name="x"))
        self.assert_eq(psidx.to_frame(index=False, name="x"), pidx.to_frame(index=False, name="x"))

        self.assertRaises(TypeError, lambda: psidx.to_frame(name=["x"]))

        # non-string name
        self.assert_eq(psidx.to_frame(name=10), pidx.to_frame(name=10))
        self.assert_eq(psidx.to_frame(name=("x", 10)), pidx.to_frame(name=("x", 10)))

        pidx = self.pdf.set_index("b", append=True).index
        psidx = self.psdf.set_index("b", append=True).index

        self.assert_eq(psidx.to_frame(), pidx.to_frame())
        self.assert_eq(psidx.to_frame(index=False), pidx.to_frame(index=False))

        self.assert_eq(psidx.to_frame(name=["x", "y"]), pidx.to_frame(name=["x", "y"]))
        self.assert_eq(psidx.to_frame(name=("x", "y")), pidx.to_frame(name=("x", "y")))
        self.assert_eq(
            psidx.to_frame(index=False, name=["x", "y"]),
            pidx.to_frame(index=False, name=["x", "y"]),
        )

        self.assertRaises(TypeError, lambda: psidx.to_frame(name="x"))
        self.assertRaises(ValueError, lambda: psidx.to_frame(name=["x"]))

        # non-string names
        self.assert_eq(psidx.to_frame(name=[10, 20]), pidx.to_frame(name=[10, 20]))
        self.assert_eq(psidx.to_frame(name=("x", 10)), pidx.to_frame(name=("x", 10)))
        # Since pandas 1.5.0, the result is changed as below:
        #      (x, 10)  (y, 20)
        #   b
        # 0 4        0        4
        # 1 5        1        5
        # 3 6        3        6
        # 5 3        5        3
        # 6 2        6        2
        # 8 1        8        1
        # 9 0        9        0
        #   0        9        0
        #   0        9        0
        #
        # The columns should be `Index([('x', 20), ('y', 20)], dtype='object')`,
        # but pandas API on Spark doesn't support such a way for creating Index.
        # So, we currently cannot follow the behavior of pandas.
        expected_result = ps.DataFrame(
            {("x", 10): [0, 1, 3, 5, 6, 8, 9, 9, 9], ("y", 20): [4, 5, 6, 3, 2, 1, 0, 0, 0]},
            index=ps.MultiIndex.from_tuples(
                [(0, 4), (1, 5), (3, 6), (5, 3), (6, 2), (8, 1), (9, 0), (9, 0), (9, 0)],
                names=[None, "b"],
            ),
        )
        self.assert_eq(psidx.to_frame(name=[("x", 10), ("y", 20)]), expected_result)

    def test_to_list(self):
        # Index
        pidx = pd.Index([1, 2, 3, 4, 5])
        psidx = ps.from_pandas(pidx)
        # MultiIndex
        tuples = [(1, "red"), (1, "blue"), (2, "red"), (2, "green")]
        pmidx = pd.MultiIndex.from_tuples(tuples)
        psmidx = ps.from_pandas(pmidx)

        self.assert_eq(psidx.tolist(), pidx.tolist())
        self.assert_eq(psmidx.tolist(), pmidx.tolist())

    def test_to_numpy(self):
        pidx = pd.Index([1, 2, 3, 4])
        psidx = ps.from_pandas(pidx)

        self.assert_eq(pidx.to_numpy(copy=True), psidx.to_numpy(copy=True))


class ConversionTests(
    ConversionMixin,
    PandasOnSparkTestCase,
    SQLTestUtils,
):
    pass


if __name__ == "__main__":
    from pyspark.pandas.tests.indexes.test_conversion 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)
