Aggregate

class BaseAggregator(abc.ABC):
    """
    Abstract base class for all aggregators.

    All subclasses expose a single ``aggregate(agg_df)`` method that accepts
    an input DataFrame and returns a per-(batch, label) summary DataFrame.
    """

    @abc.abstractmethod
    def aggregate(self, agg_df: pl.DataFrame) -> pl.DataFrame:
        """
        Compute per-(batch, label) statistics.

        Parameters
        ----------
        agg_df : pl.DataFrame
            Input DataFrame containing ``_meta_batch``, ``_meta_label``, and
            feature columns.

        Returns
        -------
        pl.DataFrame
            One row per (batch, label) group with computed statistics.
        """
        raise NotImplementedError

class NativeAggregator(BaseAggregator):
    """
    Base class for aggregators implemented as native Polars group-by expressions.

    Subclasses override :meth:`polars_exprs` to return a list of Polars
    expressions that are passed directly to ``group_by().agg()``.
    """

    @abc.abstractmethod
    def polars_exprs(self, feature_cols: list[str]) -> list[pl.Expr]:
        """
        Return Polars aggregation expressions for the given feature columns.

        Parameters
        ----------
        feature_cols : list[str]
            Names of the feature columns to aggregate.

        Returns
        -------
        list[pl.Expr]
            Polars expressions to pass to ``group_by().agg()``.
        """
        raise NotImplementedError

    def aggregate(self, agg_df: pl.DataFrame) -> pl.DataFrame:
        feature_cols = get_feature_cols(agg_df, as_string=True)
        return (
            agg_df.filter(pl.col("_meta_label") != "WT")
            .group_by("_meta_label", "_meta_batch")
            .agg(self.polars_exprs(feature_cols))
        )

class MeanAggregator(NativeAggregator):
    """Computes per-group mean for each feature column."""

    def polars_exprs(self, feature_cols: list[str]) -> list[pl.Expr]:
        return [pl.col(f).mean().alias(f"{f}_mean") for f in feature_cols]

class MedianAggregator(NativeAggregator):
    """Computes per-group median for each feature column."""

    def polars_exprs(self, feature_cols: list[str]) -> list[pl.Expr]:
        return [pl.col(f).median().alias(f"{f}_median") for f in feature_cols]

class MADAggregator(NativeAggregator):
    """Computes per-group median absolute deviation (MAD) for each feature column."""

    def polars_exprs(self, feature_cols: list[str]) -> list[pl.Expr]:
        return [
            (pl.col(f) - pl.col(f).median()).abs().median().alias(f"{f}_MAD")
            for f in feature_cols
        ]

class StdAggregator(NativeAggregator):
    """Computes per-group standard deviation for each feature column."""

    def polars_exprs(self, feature_cols: list[str]) -> list[pl.Expr]:
        return [pl.col(f).std().alias(f"{f}_std") for f in feature_cols]

class ReferenceBaseAggregator(BaseAggregator):
    """
    Base class for aggregators that compare each group against a reference
    distribution, typically built from control rows.

    Subclasses override :meth:`compute_statistic` to implement a specific
    scalar summary (e.g. EMD, KS statistic) for each feature.

    Parameters
    ----------
    reference_df : pl.DataFrame
        Reference DataFrame (typically control rows) used as the comparison
        distribution in :meth:`compute_statistic`.
    """

    def __init__(self, reference_df: pl.DataFrame) -> None:
        self.feature_cols = get_feature_cols(reference_df, as_string=True)
        self.reference_df = reference_df

    @abc.abstractmethod
    def compute_statistic(
        self, label: str, batch: str, group_df: pl.DataFrame, ref_df: pl.DataFrame
    ) -> dict[str, Any]:
        """
        Compute per-feature statistics for a single (batch, label) group.

        Parameters
        ----------
        label : str
            The variant label for this group.
        batch : str
            The batch identifier for this group.
        group_df : pl.DataFrame
            The subset of the input DataFrame corresponding to this
            (batch, label) group.
        ref_df : pl.DataFrame
            The subset of the reference DataFrame corresponding to this batch.

        Returns
        -------
        dict[str, Any]
            A dict with at minimum ``_meta_label`` and ``_meta_batch`` keys,
            plus one entry per feature column containing the computed statistic.
        """
        raise NotImplementedError

    def aggregate(
        self,
        agg_df: pl.DataFrame,
        n_jobs: int = -1,
        backend: str = "threading",
        verbose: int = 0,
    ) -> pl.DataFrame:
        """
        Compute per-(batch, label) statistics for all features using joblib.

        Filters out wildtype rows, groups by ``_meta_label`` and
        ``_meta_batch``, and dispatches one :meth:`compute_statistic` call
        per group in parallel.

        Parameters
        ----------
        agg_df : pl.DataFrame
            Input DataFrame containing ``_meta_batch``, ``_meta_label``, and
            feature columns.
        n_jobs : int, optional
            Number of parallel jobs passed to joblib.
        backend : str, optional
            Joblib backend.
        verbose : int, optional
            Joblib verbosity level. Defaults to ``0``.

        Returns
        -------
        pl.DataFrame
            A DataFrame with one row per (batch, label) group and one column
            per feature containing the computed statistic.
        """
        groups = agg_df.filter(pl.col("_meta_label") != "WT").group_by(
            "_meta_label", "_meta_batch"
        )
        tasks = [((label, batch), group_df) for (label, batch), group_df in groups]

        dicts = joblib.Parallel(n_jobs=n_jobs, backend=backend, verbose=verbose)(
            joblib.delayed(self.compute_statistic)(
                label=label,
                batch=batch,
                group_df=group_df,
                ref_df=self.reference_df.filter(pl.col("_meta_batch") == batch),
            )
            for (label, batch), group_df in tasks
        )

        return pl.DataFrame(dicts)

class EMDAggregator(ReferenceBaseAggregator):
    """
    Computes per-group 1D Wasserstein distances (Earth Mover's Distance)
    against a reference distribution for each feature column.
    """

    def compute_statistic(
        self, label: str, batch: str, group_df: pl.DataFrame, ref_df: pl.DataFrame
    ) -> dict[str, Any]:
        row: dict[str, Any] = {"_meta_label": label, "_meta_batch": batch}
        for feat in self.feature_cols:
            variant = group_df.get_column(feat).to_numpy()
            reference = ref_df.get_column(feat).to_numpy()
            row[f"{feat}_EMD"] = scipy.stats.wasserstein_distance(variant, reference)
        return row

class KSAggregator(ReferenceBaseAggregator):
    """
    Computes per-group two-sample Kolmogorov-Smirnov statistics against
    a reference distribution for each feature column.
    """

    def compute_statistic(
        self, label: str, batch: str, group_df: pl.DataFrame, ref_df: pl.DataFrame
    ) -> dict[str, Any]:
        row: dict[str, Any] = {"_meta_label": label, "_meta_batch": batch}
        for feat in self.feature_cols:
            variant = group_df.get_column(feat).to_numpy()
            reference = ref_df.get_column(feat).to_numpy()
            row[f"{feat}_KS"] = scipy.stats.ks_2samp(variant, reference).statistic
        return row

class QQCorrelationAggregator(ReferenceBaseAggregator):
    """
    Computes per-group Q-Q correlation against a reference distribution for
    each feature column.

    Parameters
    ----------
    reference_df : pl.DataFrame
        Reference DataFrame (typically control rows).
    n_quantiles : int, optional
        Number of quantile points to evaluate. Defaults to ``100``.
    """

    def __init__(self, reference_df: pl.DataFrame, n_quantiles: int = 100) -> None:
        super().__init__(reference_df)
        self.quantile_points = np.linspace(0, 1, n_quantiles)

    def compute_statistic(
        self, label: str, batch: str, group_df: pl.DataFrame, ref_df: pl.DataFrame
    ) -> dict[str, Any]:
        row: dict[str, Any] = {"_meta_label": label, "_meta_batch": batch}
        for feat in self.feature_cols:
            variant = group_df.get_column(feat).to_numpy()
            reference = ref_df.get_column(feat).to_numpy()
            variant_quantiles = np.quantile(variant, self.quantile_points)
            reference_quantiles = np.quantile(reference, self.quantile_points)
            row[f"{feat}_QQ"] = scipy.stats.pearsonr(
                variant_quantiles, reference_quantiles
            ).statistic
        return row

class AUROCAggregator(ReferenceBaseAggregator):
    """
    Computes per-group AUROC against a reference distribution for each
    feature column.

    Variant samples are labelled ``1`` and reference samples are labelled
    ``0``. ``0.5`` indicates identical distributions and ``1.0`` indicates
    perfect separability.
    """

    def compute_statistic(
        self, label: str, batch: str, group_df: pl.DataFrame, ref_df: pl.DataFrame
    ) -> dict[str, Any]:
        row: dict[str, Any] = {"_meta_label": label, "_meta_batch": batch}
        for feat in self.feature_cols:
            variant = group_df.get_column(feat).to_numpy()
            reference = ref_df.get_column(feat).to_numpy()
            values = np.concatenate([reference, variant])
            labels = np.concatenate(
                [
                    np.zeros(len(reference)),
                    np.ones(len(variant)),
                ]
            )

            auroc = sklearn.metrics.roc_auc_score(labels, values)
            if auroc < 0.5:
                auroc = 1 - auroc

            row[f"{feat}_AUROC"] = auroc

        return row