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

from pyspark.sql.functions import udf, pandas_udf, PandasUDFType, assert_true, lit
from pyspark.sql.types import (
    DoubleType,
    StructType,
    StructField,
    LongType,
    DayTimeIntervalType,
    VariantType,
)
from pyspark.errors import ParseException, PythonException, PySparkTypeError
from pyspark.util import PythonEvalType
from pyspark.testing.sqlutils import (
    ReusedSQLTestCase,
    have_pandas,
    have_pyarrow,
    pandas_requirement_message,
    pyarrow_requirement_message,
)


@unittest.skipIf(
    not have_pandas or not have_pyarrow,
    cast(str, pandas_requirement_message or pyarrow_requirement_message),
)
class PandasUDFTestsMixin:
    def test_pandas_udf_basic(self):
        udf = pandas_udf(lambda x: x, DoubleType())
        self.assertEqual(udf.returnType, DoubleType())
        self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, VariantType())
        self.assertEqual(udf.returnType, VariantType())
        self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, DoubleType(), PandasUDFType.SCALAR)
        self.assertEqual(udf.returnType, DoubleType())
        self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, VariantType(), PandasUDFType.SCALAR)
        self.assertEqual(udf.returnType, VariantType())
        self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        udf = pandas_udf(
            lambda x: x, StructType([StructField("v", DoubleType())]), PandasUDFType.GROUPED_MAP
        )
        self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        udf = pandas_udf(
            lambda x: x, StructType([StructField("v", VariantType())]), PandasUDFType.GROUPED_MAP
        )
        self.assertEqual(udf.returnType, StructType([StructField("v", VariantType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

    def test_pandas_udf_basic_with_return_type_string(self):
        udf = pandas_udf(lambda x: x, "double", PandasUDFType.SCALAR)
        self.assertEqual(udf.returnType, DoubleType())
        self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, "variant", PandasUDFType.SCALAR)
        self.assertEqual(udf.returnType, VariantType())
        self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, "v double", PandasUDFType.GROUPED_MAP)
        self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, "v variant", PandasUDFType.GROUPED_MAP)
        self.assertEqual(udf.returnType, StructType([StructField("v", VariantType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, "v double", functionType=PandasUDFType.GROUPED_MAP)
        self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, "v variant", functionType=PandasUDFType.GROUPED_MAP)
        self.assertEqual(udf.returnType, StructType([StructField("v", VariantType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        udf = pandas_udf(lambda x: x, returnType="v double", functionType=PandasUDFType.GROUPED_MAP)
        self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        udf = pandas_udf(
            lambda x: x, returnType="v variant", functionType=PandasUDFType.GROUPED_MAP
        )
        self.assertEqual(udf.returnType, StructType([StructField("v", VariantType())]))
        self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

    def test_pandas_udf_decorator(self):
        @pandas_udf(DoubleType())
        def foo(x):
            return x

        self.assertEqual(foo.returnType, DoubleType())
        self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        @pandas_udf(returnType=DoubleType())
        def foo(x):
            return x

        self.assertEqual(foo.returnType, DoubleType())
        self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

        schema = StructType([StructField("v", DoubleType())])

        @pandas_udf(schema, PandasUDFType.GROUPED_MAP)
        def foo(x):
            return x

        self.assertEqual(foo.returnType, schema)
        self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        @pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP)
        def foo(x):
            return x

        self.assertEqual(foo.returnType, schema)
        self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        @pandas_udf(returnType=schema, functionType=PandasUDFType.GROUPED_MAP)
        def foo(x):
            return x

        self.assertEqual(foo.returnType, schema)
        self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

    def test_pandas_udf_decorator_with_return_type_string(self):
        schema = StructType([StructField("v", DoubleType())])

        @pandas_udf("v double", PandasUDFType.GROUPED_MAP)
        def foo(x):
            return x

        self.assertEqual(foo.returnType, schema)
        self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)

        @pandas_udf(returnType="double", functionType=PandasUDFType.SCALAR)
        def foo(x):
            return x

        self.assertEqual(foo.returnType, DoubleType())
        self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)

    def test_udf_wrong_arg(self):
        with self.quiet():
            self.check_udf_wrong_arg()

            with self.assertRaises(ParseException):

                @pandas_udf("blah")
                def foo(x):
                    return x

            with self.assertRaises(PySparkTypeError) as pe:

                @pandas_udf(returnType="double", functionType=PandasUDFType.GROUPED_MAP)
                def foo(df):
                    return df

            self.check_error(
                exception=pe.exception,
                errorClass="INVALID_RETURN_TYPE_FOR_PANDAS_UDF",
                messageParameters={
                    "eval_type": "SQL_GROUPED_MAP_PANDAS_UDF "
                    "or SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE",
                    "return_type": "DoubleType()",
                },
            )

            with self.assertRaisesRegex(ValueError, "Invalid function"):

                @pandas_udf(returnType="k int, v double", functionType=PandasUDFType.GROUPED_MAP)
                def foo(k, v, w):
                    return k

    def check_udf_wrong_arg(self):
        with self.assertRaises(PySparkTypeError) as pe:

            @pandas_udf(functionType=PandasUDFType.SCALAR)
            def foo(x):
                return x

        self.check_error(
            exception=pe.exception,
            errorClass="CANNOT_BE_NONE",
            messageParameters={"arg_name": "returnType"},
        )

        with self.assertRaises(PySparkTypeError) as pe:

            @pandas_udf("double", 100)
            def foo(x):
                return x

        self.check_error(
            exception=pe.exception,
            errorClass="INVALID_PANDAS_UDF_TYPE",
            messageParameters={"arg_name": "functionType", "arg_type": "100"},
        )

        with self.assertRaisesRegex(ValueError, "0-arg pandas_udfs.*not.*supported"):
            pandas_udf(lambda: 1, LongType(), PandasUDFType.SCALAR)

        with self.assertRaisesRegex(ValueError, "0-arg pandas_udfs.*not.*supported"):

            @pandas_udf(LongType(), PandasUDFType.SCALAR)
            def zero_with_type():
                return 1

        with self.assertRaisesRegex(ValueError, "0-arg pandas_udfs.*not.*supported"):

            @pandas_udf(LongType(), PandasUDFType.SCALAR_ITER)
            def zero_with_type():
                yield 1
                yield 2

        with self.assertRaises(PySparkTypeError) as pe:

            @pandas_udf(returnType=PandasUDFType.GROUPED_MAP)
            def foo(df):
                return df

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_DATATYPE_OR_STR",
            messageParameters={"arg_name": "returnType", "arg_type": "int"},
        )

    def test_stopiteration_in_udf(self):
        def foo(x):
            raise StopIteration()

        exc_message = "StopIteration"
        df = self.spark.range(0, 100)

        # plain udf (test for SPARK-23754)
        self.assertRaisesRegex(
            PythonException, exc_message, df.withColumn("v", udf(foo)("id")).collect
        )

        # pandas scalar udf
        self.assertRaisesRegex(
            PythonException,
            exc_message,
            df.withColumn("v", pandas_udf(foo, "double", PandasUDFType.SCALAR)("id")).collect,
        )

    def test_stopiteration_in_grouped_map(self):
        def foo(x):
            raise StopIteration()

        def foofoo(x, y):
            raise StopIteration()

        exc_message = "StopIteration"
        df = self.spark.range(0, 100)

        # pandas grouped map
        self.assertRaisesRegex(
            PythonException,
            exc_message,
            df.groupBy("id").apply(pandas_udf(foo, df.schema, PandasUDFType.GROUPED_MAP)).collect,
        )

        self.assertRaisesRegex(
            PythonException,
            exc_message,
            df.groupBy("id")
            .apply(pandas_udf(foofoo, df.schema, PandasUDFType.GROUPED_MAP))
            .collect,
        )

    def test_stopiteration_in_grouped_agg(self):
        def foo(x):
            raise StopIteration()

        exc_message = "StopIteration"
        df = self.spark.range(0, 100)

        # pandas grouped agg
        self.assertRaisesRegex(
            PythonException,
            exc_message,
            df.groupBy("id")
            .agg(pandas_udf(foo, "double", PandasUDFType.GROUPED_AGG)("id"))
            .collect,
        )

    def test_pandas_udf_detect_unsafe_type_conversion(self):
        import pandas as pd
        import numpy as np

        values = [1.0] * 3
        pdf = pd.DataFrame({"A": values})
        df = self.spark.createDataFrame(pdf).repartition(1)

        @pandas_udf(returnType="int")
        def udf(column):
            return pd.Series(np.linspace(0, 1, len(column)))

        # Since 0.11.0, PyArrow supports the feature to raise an error for unsafe cast.
        with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": True}):
            with self.assertRaisesRegex(
                Exception, "Exception thrown when converting pandas.Series"
            ):
                df.select(["A"]).withColumn("udf", udf("A")).collect()

        # Disabling Arrow safe type check.
        with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": False}):
            df.select(["A"]).withColumn("udf", udf("A")).collect()

    def test_pandas_udf_arrow_overflow(self):
        import pandas as pd

        df = self.spark.range(0, 1)

        @pandas_udf(returnType="byte")
        def udf(column):
            return pd.Series([128] * len(column))

        # When enabling safe type check, Arrow 0.11.0+ disallows overflow cast.
        with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": True}):
            with self.assertRaisesRegex(
                Exception, "Exception thrown when converting pandas.Series"
            ):
                df.withColumn("udf", udf("id")).collect()

        # Disabling safe type check, let Arrow do the cast anyway.
        with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": False}):
            df.withColumn("udf", udf("id")).collect()

    def test_pandas_udf_int_to_decimal_coercion(self):
        import pandas as pd
        from decimal import Decimal

        df = self.spark.range(0, 3)

        @pandas_udf(returnType="decimal(10,2)")
        def int_to_decimal_udf(column):
            values = [123, 456, 789]
            return pd.Series([values[int(val) % len(values)] for val in column])

        with self.sql_conf(
            {"spark.sql.execution.pythonUDF.pandas.intToDecimalCoercionEnabled": True}
        ):
            result = df.withColumn("decimal_val", int_to_decimal_udf("id")).collect()
            self.assertEqual(result[0]["decimal_val"], 123.00)
            self.assertEqual(result[1]["decimal_val"], 456.00)
            self.assertEqual(result[2]["decimal_val"], 789.00)

        with self.sql_conf(
            {"spark.sql.execution.pythonUDF.pandas.intToDecimalCoercionEnabled": False}
        ):
            self.assertRaisesRegex(
                PythonException,
                "Exception thrown when converting pandas.Series",
                df.withColumn("decimal_val", int_to_decimal_udf("id")).collect,
            )

        @pandas_udf(returnType="decimal(25,1)")
        def high_precision_udf(column):
            values = [1, 2, 3]
            return pd.Series([values[int(val) % len(values)] for val in column])

        for intToDecimalCoercionEnabled in [True, False]:
            # arrow_cast is enabled by default for SQL_SCALAR_PANDAS_UDF and
            # and SQL_SCALAR_PANDAS_ITER_UDF, arrow can do this cast safely.
            # intToDecimalCoercionEnabled is not required for this case
            with self.sql_conf(
                {
                    "spark.sql.execution.pythonUDF.pandas.intToDecimalCoercionEnabled": intToDecimalCoercionEnabled  # noqa: E501
                }
            ):
                result = df.withColumn("decimal_val", high_precision_udf("id")).collect()
                self.assertEqual(len(result), 3)
                self.assertEqual(result[0]["decimal_val"], Decimal("1.0"))
                self.assertEqual(result[1]["decimal_val"], Decimal("2.0"))
                self.assertEqual(result[2]["decimal_val"], Decimal("3.0"))

    def test_pandas_udf_timestamp_ntz(self):
        # SPARK-36626: Test TimestampNTZ in pandas UDF
        @pandas_udf(returnType="timestamp_ntz")
        def noop(s):
            assert s.iloc[0] == datetime.datetime(1970, 1, 1, 0, 0)
            return s

        with self.sql_conf({"spark.sql.session.timeZone": "Asia/Hong_Kong"}):
            df = self.spark.createDataFrame(
                [(datetime.datetime(1970, 1, 1, 0, 0),)], schema="dt timestamp_ntz"
            ).select(noop("dt").alias("dt"))

            df.selectExpr("assert_true('1970-01-01 00:00:00' == CAST(dt AS STRING))").collect()
            self.assertEqual(df.schema[0].dataType.typeName(), "timestamp_ntz")
            self.assertEqual(df.first()[0], datetime.datetime(1970, 1, 1, 0, 0))

    def test_pandas_udf_day_time_interval_type(self):
        # SPARK-37277: Test DayTimeIntervalType in pandas UDF
        import pandas as pd

        @pandas_udf(DayTimeIntervalType(DayTimeIntervalType.DAY, DayTimeIntervalType.SECOND))
        def noop(s: pd.Series) -> pd.Series:
            assert s.iloc[0] == datetime.timedelta(microseconds=123)
            return s

        df = self.spark.createDataFrame(
            [(datetime.timedelta(microseconds=123),)], schema="td interval day to second"
        ).select(noop("td").alias("td"))

        df.select(
            assert_true(lit("INTERVAL '0 00:00:00.000123' DAY TO SECOND") == df.td.cast("string"))
        ).collect()
        self.assertEqual(df.schema[0].dataType.simpleString(), "interval day to second")
        self.assertEqual(df.first()[0], datetime.timedelta(microseconds=123))

    def test_pandas_udf_return_type_error(self):
        import pandas as pd

        @pandas_udf("s string")
        def upper(s: pd.Series) -> pd.Series:
            return s.str.upper()

        df = self.spark.createDataFrame([("a",)], schema="s string")

        self.assertRaisesRegex(
            PythonException, "Invalid return type", df.select(upper("s")).collect
        )

    def test_pandas_udf_empty_frame(self):
        import pandas as pd

        empty_df = self.spark.createDataFrame([], "id long")

        @pandas_udf("long")
        def add1(x: pd.Series) -> pd.Series:
            return x + 1

        result = empty_df.select(add1("id"))
        self.assertEqual(result.collect(), [])


class PandasUDFTests(PandasUDFTestsMixin, ReusedSQLTestCase):
    pass


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