"""Pandas grouping specification visitor for therismos grouping.

This module provides a visitor that converts therismos grouping specifications into
pandas-compatible groupby parameters via the PandasGroupSpec dataclass.

Requires Pandas to be installed. Install with: pip install therismos[pandas]

Example
-------
>>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
>>> from therismos.grouping.visitors.pandas import PandasGroupSpecVisitor
>>>
>>> spec = GroupSpec(
...     group_by=["category", "region"],
...     aggregations=[
...         Aggregation("total", AggregationFunction.COUNT),
...         Aggregation("avg_price", AggregationFunction.AVERAGE, "price"),
...     ],
... )
>>>
>>> visitor = PandasGroupSpecVisitor()
>>> grp = spec.accept(visitor)
>>>
>>> # Use with DataFrame
>>> # df.groupby(list(grp.group_by)).agg(**grp.agg)

Notes
-----
- COUNT uses the first group_by field as the column and "count" as aggfunc
- COUNT raises ValueError if group_by is empty
- Percentile functions use named helper functions (not lambdas) for compatibility
- Named aggregations use pd.NamedAgg for explicit column/aggfunc mapping
"""

from __future__ import annotations

from dataclasses import dataclass

try:
    import pandas as pd
except ImportError as exc:
    msg = "pandas is required. Install with: pip install therismos[pandas]"
    raise ImportError(msg) from exc

from therismos.grouping._grouping import Aggregation, AggregationFunction, GroupSpec


def _quantile_01(s: pd.Series[float]) -> float:
    """Compute 1st percentile."""
    return float(s.quantile(0.01))


def _quantile_05(s: pd.Series[float]) -> float:
    """Compute 5th percentile."""
    return float(s.quantile(0.05))


def _quantile_10(s: pd.Series[float]) -> float:
    """Compute 10th percentile."""
    return float(s.quantile(0.10))


def _quantile_25(s: pd.Series[float]) -> float:
    """Compute 25th percentile (Q1)."""
    return float(s.quantile(0.25))


def _quantile_75(s: pd.Series[float]) -> float:
    """Compute 75th percentile (Q3)."""
    return float(s.quantile(0.75))


def _quantile_90(s: pd.Series[float]) -> float:
    """Compute 90th percentile."""
    return float(s.quantile(0.90))


def _quantile_95(s: pd.Series[float]) -> float:
    """Compute 95th percentile."""
    return float(s.quantile(0.95))


def _quantile_99(s: pd.Series[float]) -> float:
    """Compute 99th percentile."""
    return float(s.quantile(0.99))


@dataclass(frozen=True)
class PandasGroupSpec:
    """Dataclass holding pandas groupby parameters.

    :ivar group_by: Tuple of field names to group by.
    :vartype group_by: tuple[str, ...]
    :ivar agg: Dictionary mapping output column names to pd.NamedAgg specifications.
    :vartype agg: dict[str, pd.NamedAgg]

    Example
    -------
    >>> grp = PandasGroupSpec(
    ...     group_by=("category",),
    ...     agg={"total": pd.NamedAgg(column="category", aggfunc="count")},
    ... )
    >>> df.groupby(list(grp.group_by)).agg(**grp.agg)
    """

    group_by: tuple[str, ...]
    agg: dict[str, pd.NamedAgg]


class PandasGroupSpecVisitor:
    """Visitor that converts grouping specifications to pandas groupby parameters.

    This visitor traverses grouping specifications and generates a PandasGroupSpec
    dataclass that can be passed to DataFrame.groupby().agg().

    Example
    -------
    >>> import pandas as pd
    >>> from therismos.grouping import GroupSpec, Aggregation, AggregationFunction
    >>> from therismos.grouping.visitors.pandas import PandasGroupSpecVisitor
    >>>
    >>> spec = GroupSpec(
    ...     group_by=["status"],
    ...     aggregations=[
    ...         Aggregation("count", AggregationFunction.COUNT),
    ...         Aggregation("avg_age", AggregationFunction.AVERAGE, "age"),
    ...     ],
    ... )
    >>>
    >>> visitor = PandasGroupSpecVisitor()
    >>> grp = spec.accept(visitor)
    >>> df = pd.DataFrame({"status": ["a", "b", "a"], "age": [20, 30, 25]})
    >>> df.groupby(list(grp.group_by)).agg(**grp.agg)
    """

    def _convert_aggregation(self, agg: Aggregation, first_group_field: str | None) -> pd.NamedAgg:
        """Convert a single aggregation to a pandas NamedAgg.

        :param agg: The aggregation to convert.
        :type agg: Aggregation
        :param first_group_field: The first group_by field, used for COUNT.
        :type first_group_field: str | None
        :returns: pandas NamedAgg specification.
        :rtype: pd.NamedAgg
        :raises ValueError: If COUNT is used with no group_by fields,
            or if the aggregation function is not supported.
        """
        func = agg.function

        match func:
            case AggregationFunction.COUNT:
                if first_group_field is None:
                    msg = "COUNT aggregation requires at least one group_by field"
                    raise ValueError(msg)
                return pd.NamedAgg(column=first_group_field, aggfunc="count")
            case AggregationFunction.SUM:
                return pd.NamedAgg(column=agg.field, aggfunc="sum")  # type: ignore[arg-type]
            case AggregationFunction.MIN:
                return pd.NamedAgg(column=agg.field, aggfunc="min")  # type: ignore[arg-type]
            case AggregationFunction.MAX:
                return pd.NamedAgg(column=agg.field, aggfunc="max")  # type: ignore[arg-type]
            case AggregationFunction.AVERAGE:
                return pd.NamedAgg(column=agg.field, aggfunc="mean")  # type: ignore[arg-type]
            case AggregationFunction.STDDEV:
                return pd.NamedAgg(column=agg.field, aggfunc="std")  # type: ignore[arg-type]
            case AggregationFunction.MEDIAN:
                return pd.NamedAgg(column=agg.field, aggfunc="median")  # type: ignore[arg-type]
            case AggregationFunction.Q1:
                return pd.NamedAgg(column=agg.field, aggfunc=_quantile_25)  # type: ignore[arg-type]
            case AggregationFunction.Q3:
                return pd.NamedAgg(column=agg.field, aggfunc=_quantile_75)  # type: ignore[arg-type]
            case AggregationFunction.P01:
                return pd.NamedAgg(column=agg.field, aggfunc=_quantile_01)  # type: ignore[arg-type]
            case AggregationFunction.P05:
                return pd.NamedAgg(column=agg.field, aggfunc=_quantile_05)  # type: ignore[arg-type]
            case AggregationFunction.P10:
                return pd.NamedAgg(column=agg.field, aggfunc=_quantile_10)  # type: ignore[arg-type]
            case AggregationFunction.P90:
                return pd.NamedAgg(column=agg.field, aggfunc=_quantile_90)  # type: ignore[arg-type]
            case AggregationFunction.P95:
                return pd.NamedAgg(column=agg.field, aggfunc=_quantile_95)  # type: ignore[arg-type]
            case AggregationFunction.P99:
                return pd.NamedAgg(column=agg.field, aggfunc=_quantile_99)  # type: ignore[arg-type]
            case _:
                msg = f"Unsupported aggregation function: {func}"
                raise ValueError(msg)

    def visit_group_spec(self, spec: GroupSpec) -> PandasGroupSpec:
        """Convert a grouping specification to pandas groupby parameters.

        :param spec: The grouping specification to visit.
        :type spec: GroupSpec
        :returns: PandasGroupSpec with group_by fields and NamedAgg specifications.
        :rtype: PandasGroupSpec
        :raises ValueError: If COUNT is used with no group_by fields.
        """
        first_group_field = spec.group_by[0] if spec.group_by else None

        agg_dict = {
            agg.id: self._convert_aggregation(agg, first_group_field)
            for agg in spec.aggregations.values()
        }

        return PandasGroupSpec(group_by=spec.group_by, agg=agg_dict)


__all__ = [
    "PandasGroupSpec",
    "PandasGroupSpecVisitor",
]
