#
# 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.
#

"""
A collections of builtin functions
"""
import sys
import functools
import warnings

from pyspark import since, SparkContext
from pyspark.rdd import PythonEvalType
from pyspark.sql.column import Column, _to_java_column, _to_seq, _create_column_from_literal
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.types import StringType, DataType
# Keep UserDefinedFunction import for backwards compatible import; moved in SPARK-22409
from pyspark.sql.udf import UserDefinedFunction, _create_udf  # noqa: F401
from pyspark.sql.udf import _create_udf
# Keep pandas_udf and PandasUDFType import for backwards compatible import; moved in SPARK-28264
from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType  # noqa: F401
from pyspark.sql.utils import to_str

# Note to developers: all of PySpark functions here take string as column names whenever possible.
# Namely, if columns are referred as arguments, they can be always both Column or string,
# even though there might be few exceptions for legacy or inevitable reasons.
# If you are fixing other language APIs together, also please note that Scala side is not the case
# since it requires to make every single overridden definition.


def _get_get_jvm_function(name, sc):
    """
    Retrieves JVM function identified by name from
    Java gateway associated with sc.
    """
    return getattr(sc._jvm.functions, name)


def _invoke_function(name, *args):
    """
    Invokes JVM function identified by name with args
    and wraps the result with :class:`~pyspark.sql.Column`.
    """
    jf = _get_get_jvm_function(name, SparkContext._active_spark_context)
    return Column(jf(*args))


def _invoke_function_over_column(name, col):
    """
    Invokes unary JVM function identified by name
    and wraps the result with :class:`~pyspark.sql.Column`.
    """
    return _invoke_function(name, _to_java_column(col))


def _invoke_binary_math_function(name, col1, col2):
    """
    Invokes binary JVM math function identified by name
    and wraps the result with :class:`~pyspark.sql.Column`.
    """
    return _invoke_function(
        name,
        # For legacy reasons, the arguments here can be implicitly converted into floats,
        # if they are not columns or strings.
        _to_java_column(col1) if isinstance(col1, (str, Column)) else float(col1),
        _to_java_column(col2) if isinstance(col2, (str, Column)) else float(col2)
    )


def _options_to_str(options=None):
    if options:
        return {key: to_str(value) for (key, value) in options.items()}
    return {}


def lit(col):
    """
    Creates a :class:`~pyspark.sql.Column` of literal value.

    .. versionadded:: 1.3.0

    Examples
    --------
    >>> df.select(lit(5).alias('height')).withColumn('spa
rk_user', lit(True)).take(1)
    [Row(height=5, spark_user=True)]
    """
    return col if isinstance(col, Column) else _invoke_function("lit", col)


@since(1.3)
def col(col):
    """
    Returns a :class:`~pyspark.sql.Column` based on the given column name.'
    """
    return _invoke_function("col", col)


@since(1.3)
def column(col):
    """
    Returns a :class:`~pyspa
rk.sql.Column` based on the given column name.'
    """
    return col(col)


@since(1.3)
def asc(col):
    """
    Returns a sort expression based on the ascending order of the given column name.
    """
    return (
        col.asc() if isinstance(col, Column)
        else _invoke_function("asc", col)
    )


@since(1.3)
def desc(col):
    """
    Returns a sort expression based on the descending order of the given column name.
    """
    return (
        col.desc() if isinstance(col, Column)
        else _invoke_function("desc", col)
    )


@since(1.3)
def sqrt(col):
    """
    Computes the square root of the specified float value.
    """
    return _invoke_function_over_column("sqrt", col)


@since(1.3)
def abs(col):
    """
    Computes the absolute value.
    """
    return _invoke_function_over_column("abs", col)


@since(1.3)
def max(col):
    """
    Aggregate function: returns the maximum value of the expression in a group.
    """
    return _invoke_function_over_column("max", col)


@since(1.3)
def min(col):
    """
    Aggregate function: returns the minimum value of the expression in a group.
    """
    return _invoke_function_over_column("min", col)


@since(1.3)
def count(col):
    """
    Aggregate function: returns the number of items in a group.
    """
    return _invoke_function_over_column("count", col)


@since(1.3)
def sum(col):
    """
    Aggregate function: returns the sum of all values in the expression.
    """
    return _invoke_function_over_column("sum", col)


@since(1.3)
def avg(col):
    """
    Aggregate function: returns the average of the values in a group.
    """
    return _invoke_function_over_column("avg", col)


@since(1.3)
def mean(col):
    """
    Aggregate function: returns the average of the values in a group.
    """
    return _invoke_function_over_column("mean", col)


@since(1.3)
def sumDistinct(col):
    """
    Aggregate function: returns the sum of distinct values in the expression.

    .. deprecated:: 3.2.0
        Use :func:`sum_distinct` instead.
    """
    warnings.warn("Deprecated in 3.2, use sum_distinct instead.", FutureWarning)
    return sum_distinct(col)


@since(3.2)
def sum_distinct(col):
    """
    Aggregate function: returns the sum of distinct values in the expression.
    """
    return _invoke_function_over_column("sum_distinct", col)


def acos(col):
    """
    .. versionadded:: 1.4.0

    Returns
    -------
    :class:`~pyspark.sql.Column`
        inverse cosine of `col`, as if computed by `java.lang.Math.acos()`
    """
    return _invoke_function_over_column("acos", col)


def acosh(col):
    """
    Computes inverse hyperbolic cosine of the input column.

    .. versionadded:: 3.1.0

    Returns
    -------
    :class:`~pyspark.sql.Column`
    """
    return _invoke_function_over_column("acosh", col)


def asin(col):
    """
    .. versionadded:: 1.3.0


    Returns
    -------
    :class:`~pyspark.sql.Column`
        inverse sine of `col`, as if computed by `java.lang.Math.asin()`
    """
    return _invoke_function_over_column("asin", col)


def asinh(col):
    """
    Computes inverse hyperbolic sine of the input column.

    .. versionadded:: 3.1.0

    Returns
    -------
    :class:`~pyspark.sql.Column`
    """
    return _invoke_function_over_column("asinh", col)


def atan(col):
    """
    .. versionadded:: 1.4.0

    Returns
    -------
    :class:`~pyspark.sql.Column`
        inverse tangent of `col`, as if computed by `java.lang.Math.atan()`
    """
    return _invoke_function_over_column("atan", col)


def atanh(col):
    """
    Computes inverse hyperbolic tangent of the input column.

    .. versionadded:: 3.1.0

    Returns
    -------
    :class:`~pyspark.sql.Column`
    """
    return _invoke_function_over_column("atanh", col)


@since(1.4)
def cbrt(col):
    """
    Computes the cube-root of the given value.
    """
    return _invoke_function_over_column("cbrt", col)


@since(1.4)
def ceil(col):
    """
    Computes the ceiling of the given value.
    """
    return _invoke_function_over_column("ceil", col)


def cos(col):
    """
    .. versionadded:: 1.4.0

    Parameters
    ----------
    col : :class:`~pyspark.sql.Column` or str
        angle in radians

    Returns
