"""
Aggregation functions for relational operations.

This module provides functions for aggregating data including:
- Statistical aggregates
- String aggregation
- Ranking and sorting
- Ordering modifiers
- Grouped aggregations
"""
from __future__ import annotations
from ..frontend.base import Field, Library, Aggregate, Group, TupleVariable, Value, Distinct
from ..frontend.core import Any, Boolean, Number, String, Integer, TypeVar, ScaledNumber, Numeric, Float
from relationalai.util.docutils import include_in_docs

__include_in_docs__ = True

library = Library("aggregates")

AggValue = Value | Distinct

#------------------------------------------------------
# Aggregates
#------------------------------------------------------

_sum = library.Relation("sum", fields=[Field.input("value", Numeric), Field("result", Numeric)],
                        overloads=[[Number, Number], [Float, Float]])

@include_in_docs
def sum(*args: AggValue) -> Aggregate:
    """
    Compute the sum of values.

    Parameters
    ----------
    *args: AggValue
        Values to sum. Can include Distinct wrapper for distinct aggregation.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the sum. Returns `Number` if the input is `Number`, or `Float` if the input is `Float`.

    Examples
    --------
    Sum order amounts:

    >>> select(aggregates.sum(Order.amount))

    Sum employee salaries per department:

    >>> select(Department, aggregates.sum(Employee.salary).per(Department))
    """
    return Aggregate(_sum, *args)

_count = library.Relation("count", fields=[Field("result", Integer)])

@include_in_docs
def count(*args: AggValue) -> Aggregate:
    """
    Count the number of values.

    Parameters
    ----------
    *args: AggValue
        Values to count. Can include Distinct wrapper for distinct count.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the count. Returns `Integer`.

    Examples
    --------
    Count all employees:

    >>> select(aggregates.count(Employee))

    Count employees per department:

    >>> select(Department, aggregates.count(Employee).per(Department))
    """
    return Aggregate(_count, *args)

_min = library.Relation("min", fields=[Field.input("over", TypeVar), Field("result", TypeVar)])

@include_in_docs
def min(*args: AggValue) -> Aggregate:
    """
    Find the minimum value.

    Parameters
    ----------
    *args: AggValue
        Values to find minimum from.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the minimum value. Returns the same type as the input.

    Examples
    --------
    Find minimum salary:

    >>> select(aggregates.min(Employee.salary))

    Find minimum price per category:

    >>> select(Category, aggregates.min(Product.price).per(Category).where(Product.category == Category))
    """
    return Aggregate(_min, *args)

_max = library.Relation("max", fields=[Field.input("over", TypeVar), Field("result", TypeVar)])

@include_in_docs
def max(*args: AggValue) -> Aggregate:
    """
    Find the maximum value.

    Parameters
    ----------
    *args: AggValue
        Values to find maximum from.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the maximum value. Returns the same type as the input.

    Examples
    --------
    Find maximum salary:

    >>> select(aggregates.max(Employee.salary))

    Find maximum order amount per customer:

    >>> select(Customer, aggregates.max(Order.amount).per(Customer).where(Order.customer == Customer))
    """
    return Aggregate(_max, *args)

_avg = library.Relation("avg", fields=[Field.input("over", Numeric), Field("result", Numeric)],
                        overloads=[[Number, ScaledNumber], [Float, Float]])

@include_in_docs
def avg(*args: AggValue) -> Aggregate:
    """
    Compute the average of values.

    Parameters
    ----------
    *args: AggValue
        Values to average.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the average. Returns `ScaledNumber` if the input is `Number`, or `Float` if the input is `Float`.

    Examples
    --------
    Compute average salary:

    >>> select(aggregates.avg(Employee.salary))

    Compute average order amount per customer:

    >>> select(Customer, aggregates.avg(Order.amount).per(Customer).where(Order.customer == Customer))
    """
    return Aggregate(_avg, *args)

_string_join = library.Relation("string_join",
        fields=[Field.input("index", Number), Field.input("sep", String), Field.input("over", String), Field("result", String)])

@include_in_docs
def string_join(*args: AggValue, sep="", index=1) -> Aggregate:
    """
    Join string values with a separator.

    Parameters
    ----------
    *args: AggValue
        String values to join.
    sep: str
        Separator string to use between values. Default: empty string.
    index: int
        Index parameter for ordering. Default: 1.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the joined string. Returns `String`.

    Examples
    --------
    Join employee names with comma separator:

    >>> select(aggregates.string_join(Employee.name, sep=", "))

    Join product tags per category:

    >>> select(Category, aggregates.string_join(Product.tag, sep="; ").per(Category).where(Product.category == Category))
    """
    return Aggregate(_string_join, index, sep, *args)

_sort = library.Relation("sort", fields=[Field.input("limit", Integer), Field.input("args", Any, is_list=True), Field.input("is_asc", Boolean, is_list=True)])
_rank = library.Relation("rank", fields=[Field.input("args", Any, is_list=True), Field.input("is_asc", Boolean, is_list=True), Field("rank", Integer)])
_limit = library.Relation("limit", fields=[Field.input("limit", Integer), Field.input("args", Any, is_list=True), Field.input("is_asc", Boolean, is_list=True)])


class Ordering:
    """Helper class for specifying sort order in aggregations."""
    def __init__(self, *values:AggValue, is_asc=True) -> None:
        self._values = values
        self._is_asc = is_asc

    @staticmethod
    def handle_arg(arg: AggValue|Ordering, ordering, ordering_args, is_asc=True) -> bool:
        has_distinct = False
        if isinstance(arg, Ordering):
            for v in arg._values:
                has_distinct = Ordering.handle_arg(v, ordering, ordering_args, arg._is_asc) or has_distinct
        elif isinstance(arg, Distinct):
            has_distinct = True
            for v in arg._items:
                Ordering.handle_arg(v, ordering, ordering_args, is_asc)
        else:
            ordering.append(is_asc)
            ordering_args.append(arg)
        return has_distinct

    @staticmethod
    def get_ordering_args(args: tuple[AggValue|Ordering, ...]) -> tuple[tuple[AggValue,...], tuple[bool,...], bool]:
        ordering = []
        ordering_args = []
        has_distinct = False
        for arg in args:
            has_distinct = Ordering.handle_arg(arg, ordering, ordering_args) or has_distinct
        return tuple(ordering_args), tuple(ordering), has_distinct

@include_in_docs
def asc(*args: AggValue) -> Ordering:
    """
    Specify ascending order for values in sorting and ranking operations.

    Parameters
    ----------
    *args: AggValue
        Values to sort in ascending order.

    Returns
    -------
    Ordering
        An `Ordering` object representing ascending order.

    Examples
    --------
    Rank employees by salary in ascending order:

    >>> select(Employee, aggregates.rank(aggregates.asc(Employee.salary)))
    """
    return Ordering(*args, is_asc=True)

@include_in_docs
def desc(*args: AggValue) -> Ordering:
    """
    Specify descending order for values in sorting and ranking operations.

    Parameters
    ----------
    *args: AggValue
        Values to sort in descending order.

    Returns
    -------
    Ordering
        An `Ordering` object representing descending order.

    Examples
    --------
    Rank employees by salary in descending order:

    >>> select(Employee, aggregates.rank(aggregates.desc(Employee.salary)))
    """
    return Ordering(*args, is_asc=False)

def _sort_agg(limit: int, *args: AggValue|Ordering) -> Aggregate:
    ordering_args, is_asc, has_distinct = Ordering.get_ordering_args(args)
    return Aggregate(_sort, limit, TupleVariable(ordering_args), TupleVariable(is_asc), distinct=has_distinct)

@include_in_docs
def rank(*args: AggValue|Ordering) -> Aggregate:
    """
    Compute rank based on the ordering of values.

    Parameters
    ----------
    *args: AggValue | Ordering
        Values or Ordering specifications to rank by.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the rank. Returns `Integer`.

    Examples
    --------
    Rank employees by salary:

    >>> select(Employee, aggregates.rank(aggregates.desc(Employee.salary)))

    Rank products by price per category:

    >>> select(Product, aggregates.rank(aggregates.asc(Product.price)).per(Category).where(Product.category == Category))
    """
    ordering_args, is_asc, has_distinct = Ordering.get_ordering_args(args)
    return Aggregate(_rank, TupleVariable(ordering_args), TupleVariable(is_asc), distinct=has_distinct)

@include_in_docs
def limit(limit: int, *args: AggValue|Ordering) -> Aggregate:
    """
    Limit results based on ordering.

    Parameters
    ----------
    limit: int
        Maximum number of results to return.
    *args: AggValue | Ordering
        Values or Ordering specifications to order by.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the limited results.

    Examples
    --------
    Get top 5 employees by salary:

    >>> select(Employee).where(aggregates.limit(5, aggregates.desc(Employee.salary)))

    Get top 3 products by sales per category:

    >>> select(Product).where(aggregates.limit(3, aggregates.desc(Product.sales)).per(Category).where(Product.category == Category))
    """
    ordering_args, is_asc, has_distinct = Ordering.get_ordering_args(args)
    return Aggregate(_limit, limit, TupleVariable(ordering_args), TupleVariable(is_asc), distinct=has_distinct)

@include_in_docs
def rank_asc(*args: AggValue) -> Aggregate:
    """
    Compute rank in ascending order.

    Parameters
    ----------
    *args: AggValue
        Values to rank in ascending order.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the ascending rank. Returns `Integer`.

    Examples
    --------
    Rank students by test score (lowest to highest):

    >>> select(Student, aggregates.rank_asc(Student.test_score))
    """
    return Aggregate(_rank, TupleVariable(args), TupleVariable([True for _ in args]))

@include_in_docs
def rank_desc(*args: AggValue) -> Aggregate:
    """
    Compute rank in descending order.

    Parameters
    ----------
    *args: AggValue
        Values to rank in descending order.

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the descending rank. Returns `Integer`.

    Examples
    --------
    Rank students by test score (highest to lowest):

    >>> select(Student, aggregates.rank_desc(Student.test_score))
    """
    return Aggregate(_rank, TupleVariable(args), TupleVariable([False for _ in args]))

@include_in_docs
def top(limit: int, *args: AggValue) -> Aggregate:
    """
    Get the top N results in descending order.

    Parameters
    ----------
    limit: int
        Number of top results to return.
    *args: AggValue
        Values to order by (descending).

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the top N results.

    Examples
    --------
    Get top 10 highest-paid employees:

    >>> select(Employee).where(aggregates.top(10, Employee.salary))

    Get top 3 products by revenue per store:

    >>> select(Product).where(aggregates.top(3, Product.revenue).per(Store).where(Product.store == Store))
    """
    return Aggregate(_limit, limit, TupleVariable(args), TupleVariable([False for _ in args]))

@include_in_docs
def bottom(limit: int, *args: AggValue) -> Aggregate:
    """
    Get the bottom N results in ascending order.

    Parameters
    ----------
    limit: int
        Number of bottom results to return.
    *args: AggValue
        Values to order by (ascending).

    Returns
    -------
    Aggregate
        An `Aggregate` representing the computation of the bottom N results.

    Examples
    --------
    Get 10 lowest-priced products:

    >>> select(Product).where(aggregates.bottom(10, Product.price))

    Get 5 least selling products per category:

    >>> select(Product).where(aggregates.bottom(5, Product.units_sold).per(Category).where(Product.category == Category))
    """
    return Aggregate(_limit, limit, TupleVariable(args), TupleVariable([True for _ in args]))

#------------------------------------------------------
# Per
#------------------------------------------------------

@include_in_docs
class Per(Group):
    """
    Group aggregation context for computing per-group aggregates.

    Use this class to perform aggregations within groups defined by specific values.
    """

    @include_in_docs
    def sum(self, *args: AggValue) -> Aggregate:
        """
        Compute sum per group.

        Parameters
        ----------
        *args: AggValue
            Values to sum within each group.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group sum. Returns `Number` if the input is `Number`, or `Float` if the input is `Float`.

        Examples
        --------
        Sum order amounts per customer:

        >>> select(Customer, aggregates.sum(Order.amount).per(Customer).where(Order.customer == Customer))
        """
        return sum(*args).per(*self._args)

    @include_in_docs
    def count(self, *args: AggValue) -> Aggregate:
        """
        Count values per group.

        Parameters
        ----------
        *args: AggValue
            Values to count within each group.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group count. Returns `Integer`.

        Examples
        --------
        Count employees per department:

        >>> select(Department, aggregates.count(Employee).per(Department).where(Employee.department == Department))
        """
        return count(*args).per(*self._args)

    @include_in_docs
    def min(self, *args: AggValue) -> Aggregate:
        """
        Find minimum per group.

        Parameters
        ----------
        *args: AggValue
            Values to find minimum from within each group.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group minimum. Returns the same type as the input.

        Examples
        --------
        Find minimum salary per department:

        >>> select(Department, aggregates.min(Employee.salary).per(Department).where(Employee.department == Department))
        """
        return min(*args).per(*self._args)

    @include_in_docs
    def max(self, *args: AggValue) -> Aggregate:
        """
        Find maximum per group.

        Parameters
        ----------
        *args: AggValue
            Values to find maximum from within each group.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group maximum. Returns the same type as the input.

        Examples
        --------
        Find maximum order amount per customer:

        >>> select(Customer, aggregates.max(Order.amount).per(Customer).where(Order.customer == Customer))
        """
        return max(*args).per(*self._args)

    @include_in_docs
    def avg(self, *args: AggValue) -> Aggregate:
        """
        Compute average per group.

        Parameters
        ----------
        *args: AggValue
            Values to average within each group.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group average. Returns `ScaledNumber` if the input is `Number`, or `Float` if the input is `Float`.

        Examples
        --------
        Compute average salary per department:

        >>> select(Department, aggregates.avg(Employee.salary).per(Department).where(Employee.department == Department))
        """
        return avg(*args).per(*self._args)

    @include_in_docs
    def string_join(self, *args: AggValue, sep="", index=1) -> Aggregate:
        """
        Join strings per group.

        Parameters
        ----------
        *args: AggValue
            String values to join within each group.
        sep: str
            Separator string. Default: empty string.
        index: int
            Index parameter. Default: 1.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group string join. Returns `String`.

        Examples
        --------
        Join employee names per department:

        >>> select(Department, aggregates.string_join(Employee.name, sep=", ").per(Department).where(Employee.department == Department))
        """
        return string_join(*args, sep=sep, index=index).per(*self._args)

    @include_in_docs
    def rank(self, *args: AggValue|Ordering) -> Aggregate:
        """
        Compute rank per group.

        Parameters
        ----------
        *args: AggValue | Ordering
            Values or Ordering specs to rank by within each group.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group rank. Returns `Integer`.

        Examples
        --------
        Rank employees by salary within each department:

        >>> select(Employee, aggregates.rank(aggregates.desc(Employee.salary)).per(Department).where(Employee.department == Department))
        """
        return rank(*args).per(*self._args)

    @include_in_docs
    def limit(self, limit_: int, *args: AggValue|Ordering) -> Aggregate:
        """
        Limit results per group.

        Parameters
        ----------
        limit_: int
            Maximum results per group.
        *args: AggValue | Ordering
            Values or Ordering specs to order by.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group limited results.

        Examples
        --------
        Get top 3 employees per department by salary:

        >>> select(Employee).where(aggregates.limit(3, aggregates.desc(Employee.salary)).per(Department).where(Employee.department == Department))
        """
        return limit(limit_, *args).per(*self._args)

    @include_in_docs
    def rank_asc(self, *args: AggValue) -> Aggregate:
        """
        Compute ascending rank per group.

        Parameters
        ----------
        *args: AggValue
            Values to rank in ascending order within each group.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group ascending rank. Returns `Integer`.

        Examples
        --------
        Rank products by price within each category (lowest first):

        >>> select(Product, aggregates.rank_asc(Product.price).per(Category))
        """
        return rank_asc(*args).per(*self._args)

    @include_in_docs
    def rank_desc(self, *args: AggValue) -> Aggregate:
        """
        Compute descending rank per group.

        Parameters
        ----------
        *args: AggValue
            Values to rank in descending order within each group.

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group descending rank. Returns `Integer`.

        Examples
        --------
        Rank products by sales within each category (highest first):

        >>> select(Product, aggregates.rank_desc(Product.sales).per(Category).where(Product.category == Category))
        """
        return rank_desc(*args).per(*self._args)

    @include_in_docs
    def top(self, limit: int, *args: AggValue) -> Aggregate:
        """
        Get top N per group.

        Parameters
        ----------
        limit: int
            Number of top results per group.
        *args: AggValue
            Values to order by (descending).

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group top N.

        Examples
        --------
        Get top 5 products by revenue per store:

        >>> select(Product).where(aggregates.top(5, Product.revenue).per(Store).where(Product.store == Store))
        """
        return top(limit, *args).per(*self._args)

    @include_in_docs
    def bottom(self, limit: int, *args: AggValue) -> Aggregate:
        """
        Get bottom N per group.

        Parameters
        ----------
        limit: int
            Number of bottom results per group.
        *args: AggValue
            Values to order by (ascending).

        Returns
        -------
        Aggregate
            An `Aggregate` expression for per-group bottom N.

        Examples
        --------
        Get bottom 3 products by sales per category:

        >>> select(Product).where(aggregates.bottom(3, Product.sales).per(Category).where(Product.category == Category))
        """
        return bottom(limit, *args).per(*self._args)

@include_in_docs
def per(*args: Value) -> Per:
    """
    Create a grouped aggregation context.

    Parameters
    ----------
    *args: Value
        Values defining the grouping.

    Returns
    -------
    Per
        A `Per` object for performing grouped aggregations.

    Examples
    --------
    Compute sum per category:

    >>> aggregates.per(Category).sum(amount)

    Count items per department:

    >>> aggregates.per(Department).count()
    """
    return Per(*args)
