#
# 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.
#
"""
Incremental user-defined aggregators for PySpark, the Python analog of Scala's
``org.apache.spark.sql.expressions.Aggregator``.
"""

from abc import ABC, abstractmethod
from typing import Any, Tuple

from pyspark.errors import PySparkNotImplementedError
from pyspark.sql.types import DataType, StructType

__all__ = ["Aggregator"]


class Aggregator(ABC):
    """
    Base class for a user-defined *incremental* aggregator, the Python analog of Scala's
    :class:`org.apache.spark.sql.expressions.Aggregator`.

    Unlike a grouped-aggregate ``pandas_udf`` (which materializes the whole group and is invoked
    once), an :class:`Aggregator` is executed as a genuine two-stage aggregation with map-side
    combine: :meth:`reduce` folds input rows into a per-group *buffer* on the map side, the buffers
    are shuffled by the grouping key, :meth:`merge` combines the partial buffers of each group, and
    :meth:`finish` produces the final output value.

    The buffer is represented as a Python :class:`tuple` whose elements correspond, in order, to the
    fields of :attr:`bufferSchema`. An input row is likewise a tuple of the argument values passed
    to the aggregator call. :meth:`merge` must be associative and commutative (the framework may
    combine partial buffers in any order), and :meth:`zero` must be its identity element -- see
    :meth:`zero` for the identity law that makes the result independent of the partition count.

    .. versionadded:: 4.4.0

    Examples
    --------
    A mean aggregator::

        from pyspark.sql.aggregator import Aggregator
        from pyspark.sql.functions import udaf
        from pyspark.sql.types import StructType, StructField, DoubleType, LongType

        class Mean(Aggregator):
            @property
            def bufferSchema(self):
                return StructType([
                    StructField("sum", DoubleType()),
                    StructField("count", LongType()),
                ])

            @property
            def outputType(self):
                return DoubleType()

            def zero(self):
                return (0.0, 0)

            def reduce(self, buffer, value):
                (v,) = value
                if v is None:  # ignore null inputs, like SQL aggregates do
                    return buffer
                return (buffer[0] + v, buffer[1] + 1)

            def merge(self, b1, b2):
                return (b1[0] + b2[0], b1[1] + b2[1])

            def finish(self, buffer):
                return buffer[0] / buffer[1] if buffer[1] else None

        mean = udaf(Mean())
        df.groupBy("k").agg(mean(df.v)).show()
    """

    @property
    @abstractmethod
    def bufferSchema(self) -> StructType:
        """The schema of the intermediate buffer that crosses the shuffle."""
        ...

    @property
    @abstractmethod
    def outputType(self) -> DataType:
        """The data type of the aggregator's output value."""
        ...

    @abstractmethod
    def zero(self) -> Tuple[Any, ...]:
        """The initial (identity) buffer value, as a tuple matching :attr:`bufferSchema`.

        This must be the identity element for :meth:`merge`::

            merge(buffer, zero()) == buffer
            merge(zero(), buffer) == buffer

        A fresh ``zero()`` seeds every partition -- and every early-flushed chunk of the map-side
        combine -- so associativity and commutativity of :meth:`merge` alone do not guarantee a
        partition-independent result; the identity law above is what makes the aggregate value
        independent of how the input is split across partitions and batches.
        """
        ...

    @abstractmethod
    def reduce(self, buffer: Tuple[Any, ...], value: Tuple[Any, ...]) -> Tuple[Any, ...]:
        """Fold a single input row ``value`` into ``buffer`` and return the updated buffer."""
        ...

    @abstractmethod
    def merge(self, buffer1: Tuple[Any, ...], buffer2: Tuple[Any, ...]) -> Tuple[Any, ...]:
        """Merge two partial buffers into one. Must be associative and commutative."""
        ...

    @abstractmethod
    def finish(self, buffer: Tuple[Any, ...]) -> Any:
        """Produce the output value from the final merged buffer."""
        ...

    # The aggregator instance is shipped to the worker as the UDF "function"; making it callable
    # lets it satisfy ``UserDefinedFunction``'s ``callable`` check. It is never actually invoked as
    # a function -- the worker calls :meth:`zero`/:meth:`reduce`/:meth:`merge`/:meth:`finish`.
    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        raise PySparkNotImplementedError(
            errorClass="NOT_IMPLEMENTED",
            messageParameters={"feature": "calling an Aggregator directly; wrap it with udaf(...)"},
        )