#
# 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 inspect
from typing import Union, Iterator, Tuple

from pyspark.sql.functions import mean, lit
from pyspark.testing.sqlutils import ReusedSQLTestCase, \
    have_pandas, have_pyarrow, pandas_requirement_message, \
    pyarrow_requirement_message
from pyspark.sql.pandas.typehints import infer_eval_type
from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType
from pyspark.sql import Row

if have_pandas:
    import pandas as pd
    import numpy as np
    from pandas.testing import assert_frame_equal


@unittest.skipIf(
    not have_pandas or not have_pyarrow,
    pandas_requirement_message or pyarrow_requirement_message)  # type: ignore[arg-type]
class PandasUDFTypeHintsTests(ReusedSQLTestCase):
    def test_type_annotation_scalar(self):
        def func(col: pd.Series) -> pd.Series:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR)

        def func(col: pd.DataFrame, col1: pd.Series) -> pd.DataFrame:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR)

        def func(col: pd.DataFrame, *args: pd.Series) -> pd.Series:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR)

        def func(col: pd.Series, *args: pd.Series, **kwargs: pd.DataFrame) -> pd.Series:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR)

        def func(col: pd.Series, *, col2: pd.DataFrame) -> pd.DataFrame:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR)

        def func(col: Union[pd.Series, pd.DataFrame], *, col2: pd.DataFrame) -> pd.Series:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR)

    def test_type_annotation_scalar_iter(self):
        def func(iter: Iterator[pd.Series]) -> Iterator[pd.Series]:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR_ITER)

        def func(iter: Iterator[Tuple[pd.DataFrame, pd.Series]]) -> Iterator[pd.DataFrame]:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR_ITER)

        def func(iter: Iterator[Tuple[pd.DataFrame, ...]]) -> Iterator[pd.Series]:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR_ITER)

        def func(
            iter: Iterator[Tuple[Union[pd.DataFrame, pd.Series], ...]]
        ) -> Iterator[pd.Series]:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.SCALAR_ITER)

    def test_type_annotation_group_agg(self):

        def func(col: pd.Series) -> str:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.GROUPED_AGG)

        def func(col: pd.DataFrame, col1: pd.Series) -> int:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.GROUPED_AGG)

        def func(col: pd.DataFrame, *args: pd.Series) -> Row:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.GROUPED_AGG)

        def func(col: pd.Series, *args: pd.Series, **kwargs: pd.DataFrame) -> str:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.GROUPED_AGG)

        def func(col: pd.Series, *, col2: pd.DataFrame) -> float:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.GROUPED_AGG)

        def func(col: Union[pd.Series, pd.DataFrame], *, col2: pd.DataFrame) -> float:
            pass
        self.assertEqual(
            infer_eval_type(inspect.signature(func)), PandasUDFType.GROUPED_AGG)

    def test_type_annotation_negative(self):

        def func(col: str) -> pd.Series:
            pass
        self.assertRaisesRegex(
            NotImplementedError,
            "Unsupported signature.*str",
            infer_eval_type, inspect.signature(func))

        def func(col: pd.DataFrame, col1: int) -> pd.DataFrame:
            pass
        self.assertRaisesRegex(
            NotImplementedError,
            "Unsupported signature.*int",
            infer_eval_type, inspect.signature(func))

        def func(col: Union[pd.DataFrame, str], col1: int) -> pd.DataFrame:
            pass
        self.assertRaisesRegex(
            NotImplementedError,
            "Unsupported signature.*str",
            infer_eval_type, inspect.signature(func))

        def func(col: pd.Series) -> Tuple[pd.DataFrame]:
            pass
        self.assertRaisesRegex(
            NotImplementedError,
            "Unsupported signature.*Tuple",
            infer_eval_type, inspect.signature(func))

        def func(col, *args: pd.Series) -> pd.Series:
            pass
        self.assertRaisesRegex(
            ValueError,
            "should be specified.*Series",
            infer_eval_type, inspect.signature(func))

        def func(col: pd.Series, *args: pd.Series, **kwargs: pd.DataFrame):
            pass
        self.assertRaisesRegex(
            ValueError,
            "should be specified.*Series",
            infer_eval_type, inspect.signature(func))

        def func(col: pd.Series, *, col2) -> pd.DataFrame:
            pass
        self.assertRaisesRegex(
            ValueError,
            "should be specified.*Series",
            infer_eval_type, inspect.signature(func))

    def test_scalar_udf_type_hint(self):
        df = self.spark.range(10).selectExpr("id", "id as v")

        def plus_one(v: Union[pd.Series, pd.DataFrame]) -> pd.Series:
            return v + 1

        plus_one = pandas_udf("long")(plus_one)
        actual = df.select(plus_one(df.v).alias("plus_one"))
        expected = df.selectExpr("(v + 1) as plus_one")
        assert_frame_equal(expected.toPandas(), actual.toPandas())

    def test_scalar_iter_udf_type_hint(self):
        df = self.spark.range(10).selectExpr("id", "id as v")

        def plus_one(itr: Iterator[pd.Series]) -> Iterator[pd.Series]:
            for s in itr:
                yield s + 1

        plus_one = pandas_udf("long")(plus_one)

        actual = df.select(plus_one(df.v).alias("plus_one"))
        expected = df.selectExpr("(v + 1) as plus_one")
        assert_frame_equal(expected.toPandas(), actual.toPandas())

    def test_group_agg_udf_type_hint(self):
        df = self.spark.range(10).selectExpr("id", "id as v")

        def weighted_mean(v: pd.Series, w: pd.Series) -> float:
            return np.average(v, weights=w)

        weighted_mean = pandas_udf("double")(weighted_mean)

        actual = df.groupby('id').agg(weighted_mean(df.v, lit(1.0))).sort('id')
        expected = df.groupby('id').agg(mean(df.v).alias('weighted_mean(v, 1.0)')).sort('id')
        assert_frame_equal(expected.toPandas(), actual.toPandas())

    def test_ignore_type_hint_in_group_apply_in_pandas(self):
        df = self.spark.range(10)

        def pandas_plus_one(v: pd.DataFrame) -> pd.DataFrame:
            return v + 1

        actual = df.groupby('id').applyInPandas(pandas_plus_one, schema=df.schema).sort('id')
        expected = df.selectExpr("id + 1 as id")
        assert_frame_equal(expected.toPandas(), actual.toPandas())

    def test_ignore_type_hint_in_cogroup_apply_in_pandas(self):
        df = self.spark.range(10)

        def pandas_plus_one(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame:
            return left + 1

        actual = df.groupby('id').cogroup(
            self.spark.range(10).groupby("id")
        ).applyInPandas(pandas_plus_one, schema=df.schema).sort('id')
        expected = df.selectExpr("id + 1 as id")
        assert_frame_equal(expected.toPandas(), actual.toPandas())

    def test_ignore_type_hint_in_map_in_pandas(self):
        df = self.spark.range(10)

        def pandas_plus_one(iter: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
            return map(lambda v: v + 1, iter)

        actual = df.mapInPandas(pandas_plus_one, schema=df.schema)
        expected = df.selectExpr("id + 1 as id")
        assert_frame_equal(expected.toPandas(), actual.toPandas())


if __name__ == "__main__":
    from pyspark.sql.tests.test_pandas_udf_typehints import *  # noqa: #401

    try:
        import xmlrunner  # type: ignore[import]
        testRunner = xmlrunner.XMLTestRunner(output='target/test-reports', verbosity=2)
    except ImportError:
        testRunner = None
    unittest.main(testRunner=testRunner, verbosity=2)
