from typing import Any, Dict, Optional, Tuple
import pandas as pd

from aboba.base.base_processors import BaseDataProcessor



class MeanByBucketProcessor(BaseDataProcessor):
    """
    Compute mean metric value by bucket.

    Aggregates data directly at bucket level by taking the mean metric value
    for each bucket. If `group_column` is present, computes the mean inside
    each `(bucket, group)` pair.
    Args:
    - value_column: Name of the column with metric values to average.
    - bucket_column: Name of the column with bucket labels.
    - group_column: Optional name of the column with group labels. If provided,
      mean is computed separately for each group within each bucket.
    """

    def __init__(
        self,
        value_column: str = "purchase_amount",
        bucket_column: str = "bucket",
        group_column: Optional[str] = "group",
    ):
        super().__init__()
        self.value_column = value_column
        self.bucket_column = bucket_column
        self.group_column = group_column

    def fit(self, data: pd.DataFrame):
        """Fit method for pipeline compatibility."""
        return self

    def transform(self, data: pd.DataFrame) -> Tuple[pd.DataFrame, Dict[str, Any]]:
        """
        Aggregate metric directly by bucket.

        Returns:
            Tuple[pd.DataFrame, Dict[str, Any]]:
                - DataFrame with one row per bucket (and group if present)
                - Empty artifacts dict
        """
        bucket_group_keys = [self.bucket_column]
        if self.group_column and self.group_column in data.columns:
            bucket_group_keys.append(self.group_column)

        bucketed_data = (
            data.groupby(bucket_group_keys, as_index=False)[self.value_column].mean()
        )

        return bucketed_data, dict()
