API Reference

def run_pipeline(
    votes: list[dict],
    reducer: ReducerType = "pca",
    reducer_kwargs: dict = {},
    clusterer: ClustererType = "kmeans",
    clusterer_kwargs: dict = {},
    mod_out_statement_ids: list[int] = [],
    meta_statement_ids: list[int] = [],
    min_user_vote_threshold: int = 7,
    keep_participant_ids: list[int] = [],
    init_centers: Optional[list[list[float]]] = None,
    max_group_count: int = 5,
    force_group_count: Optional[int] = None,
    random_state: Optional[int] = None,
    pick_max: int = 5,
    confidence: float = 0.9,
    consensus_mode: Literal["standard", "legacy"] = "standard",
    candidate_group_counts: CandidateGroupCounts = None,
) -> TypedPolisPipelineResult:
    """
    An essentially feature-complete implementation of the Polis clustering algorithm.

    Still missing:
        - base-cluster calculations (so can't match output of conversations larger than 100 participants),
        - k-smoothing, which holds back k-value (group count) until re-calculated 3 consecutive times,
        - some advanced participant filtering that involves past state (you can use keep_participant_ids to mimic manually).

    Args:
        votes (list[dict]): Raw list of vote dicts, with keys for "participant_id", "statement_id", "vote" and "modified"
        reducer (ReducerType): Selects the type of reducer model to use.
        reducer_kwargs (dict): Extra params to pass to reducer model during initialization.
        clusterer (ClustererType): Selects the type of clusterer model to use.
        clusterer_kwargs (dict): Extra params to pass to clusterer model during initialization.
        mod_out_statement_ids (list[int]): List of statement IDs to moderate/zero out
        meta_statement_ids (list[int]): List of meta statement IDs
        min_user_vote_threshold (int): Minimum number of votes a participant must make to be included in clustering
        keep_participant_ids (list[int]): List of participant IDs to keep in clustering algorithm, regardless of normal filters.
        max_group_count (): Max number of group (k-values) to test using k-means and silhouette scores
        init_centers (list[list[float]]): Initial guesses of [x,y] coordinates for k-means (Length of list must match max_group_count)
        force_group_count (int): Instead of using silhouette scores, force a specific number of groups (k value)
        random_state (int): If set, will force determinism during k-means clustering
        confidence (float): Percent confidence interval (in decimal), within which selected statements are deemed significant
        pick_max (int): Max number of statements selected for consensus (per direction) and representative (per group).
        candidate_group_counts ("all" | list[int]): Return one forced-k result per candidate group count after reusing the PCA projection.


    Returns:
        AnalysisSuccess or AnalysisInsufficientData. On success, `result` contains either
        PolisClusteringResult or KMeansCandidatesResult when candidate_group_counts is set.
    """
    if force_group_count is not None and candidate_group_counts is not None:
        raise ValueError("force_group_count and candidate_group_counts cannot both be set")

    reason = get_insufficient_data_reason(
        votes=votes,
        mod_out_statement_ids=mod_out_statement_ids,
        min_user_vote_threshold=min_user_vote_threshold,
        keep_participant_ids=keep_participant_ids,
        force_group_count=None if candidate_group_counts is not None else force_group_count,
    )
    if reason is not None:
        return AnalysisInsufficientData(
            outcome=AnalysisOutcome.INSUFFICIENT_DATA,
            reason=reason,
        )

    if candidate_group_counts is not None:
        if reducer != "pca":
            raise ValueError("candidate_group_counts only supports the pca reducer")
        if clusterer != "kmeans":
            raise ValueError("candidate_group_counts only supports the kmeans clusterer")
        candidate_group_counts = _get_kmeans_candidate_group_counts(
            max_group_count=max_group_count,
            candidate_group_counts=candidate_group_counts,
        )
        projection = prepare_pca_projection(
            votes=votes,
            reducer_kwargs=reducer_kwargs,
            mod_out_statement_ids=mod_out_statement_ids,
            meta_statement_ids=meta_statement_ids,
            min_user_vote_threshold=min_user_vote_threshold,
            keep_participant_ids=keep_participant_ids,
            random_state=random_state,
        )
        return AnalysisSuccess(
            outcome=AnalysisOutcome.SUCCESS,
            result=run_kmeans_candidates_on_pca_projection(
                projection=projection,
                max_group_count=max_group_count,
                candidate_group_counts=candidate_group_counts,
                init_centers=init_centers,
                random_state=random_state,
                mod_out_statement_ids=mod_out_statement_ids,
                pick_max=pick_max,
                confidence=confidence,
                consensus_mode=consensus_mode,
            ),
        )

    return _run_single_pipeline(
        votes=votes,
        reducer=reducer,
        reducer_kwargs=reducer_kwargs,
        clusterer=clusterer,
        clusterer_kwargs=clusterer_kwargs,
        mod_out_statement_ids=mod_out_statement_ids,
        meta_statement_ids=meta_statement_ids,
        min_user_vote_threshold=min_user_vote_threshold,
        keep_participant_ids=keep_participant_ids,
        init_centers=init_centers,
        max_group_count=max_group_count,
        force_group_count=force_group_count,
        random_state=random_state,
        pick_max=pick_max,
        confidence=confidence,
        consensus_mode=consensus_mode,
    )

def run_pipeline(
    fdr_rate: float = 0.10,
    **kwargs,
) -> TypedAgoraPipelineResult:
    """
    Agora clustering pipeline. Runs the base pipeline and adds ranked
    representative/consensus statements with Benjamini-Hochberg selection.

    Accepts all the same arguments as base.run_pipeline(), plus:

    Args:
        fdr_rate (float): False discovery rate for Benjamini-Hochberg selection.
        **kwargs: All arguments forwarded to base.run_pipeline().

    Returns:
        AnalysisSuccess or AnalysisInsufficientData. On success, `result` contains either
        AgoraClusteringResult or KMeansCandidatesResult when candidate_group_counts is set.
    """
    base_pipeline_result = base.run_pipeline(**kwargs)
    if base_pipeline_result.outcome == base.AnalysisOutcome.INSUFFICIENT_DATA:
        return base_pipeline_result

    mod_out_statement_ids = kwargs.get("mod_out_statement_ids", [])
    if isinstance(base_pipeline_result.result, base.KMeansCandidatesResult):
        candidates: list[
            base.KMeansCandidateSuccess[AgoraClusteringResult]
            | base.KMeansCandidateInsufficientData
        ] = []
        for candidate in base_pipeline_result.result.candidates:
            if candidate.outcome == base.AnalysisOutcome.INSUFFICIENT_DATA:
                candidates.append(candidate)
                continue
            candidates.append(
                base.KMeansCandidateSuccess(
                    group_count=candidate.group_count,
                    outcome=base.AnalysisOutcome.SUCCESS,
                    silhouette_score=candidate.silhouette_score,
                    result=_build_agora_result(
                        base_result=candidate.result,
                        mod_out_statement_ids=mod_out_statement_ids,
                        fdr_rate=fdr_rate,
                    ),
                )
            )
        return base.AnalysisSuccess(
            outcome=base.AnalysisOutcome.SUCCESS,
            result=base.KMeansCandidatesResult(
                projection=base_pipeline_result.result.projection,
                candidates=candidates,
            ),
        )

    return base.AnalysisSuccess(
        outcome=base.AnalysisOutcome.SUCCESS,
        result=_build_agora_result(
            base_result=base_pipeline_result.result,
            mod_out_statement_ids=mod_out_statement_ids,
            fdr_rate=fdr_rate,
        ),
    )

def compute_effective_agreement_gac(
    grouped_stats_df: pd.DataFrame,
    statement_ids,
    n_groups: int,
) -> dict:
    """Compute GAC using effective agreement: prod(pa * (1-pd))^(1/n_groups).

    Unlike the Polis formula (raw pa product), this penalizes groups that are
    genuinely divided (high agree AND high disagree) by discounting each group's
    agreement by its disagreement.

    Args:
        grouped_stats_df: MultiIndex DataFrame with (group_id, statement_id) index, containing 'pa' and 'pd' columns.
        statement_ids: Statement IDs to compute GAC for.
        n_groups: Number of groups.

    Returns:
        Dict with 'agree' and 'disagree' keys, each mapping statement_id to GAC score.
    """
    agree = {}
    disagree = {}
    for sid in statement_ids:
        ea_prod = 1.0
        ed_prod = 1.0
        for gid in range(n_groups):
            pa = grouped_stats_df.loc[(gid, sid), "pa"]
            pd_val = grouped_stats_df.loc[(gid, sid), "pd"]
            ea_prod *= pa * (1 - pd_val)
            ed_prod *= pd_val * (1 - pa)
        agree[sid] = ea_prod ** (1.0 / n_groups)
        disagree[sid] = ed_prod ** (1.0 / n_groups)
    return {"agree": agree, "disagree": disagree}

def run_pipeline(
    votes: Sequence[Mapping[str, object]],
    reducer: base.ReducerType = "pca",
    reducer_kwargs: dict[str, object] | None = None,
    clusterer: base.ClustererType = "kmeans",
    clusterer_kwargs: dict[str, object] | None = None,
    mod_out_statement_ids: list[int] | None = None,
    meta_statement_ids: list[int] | None = None,
    min_user_vote_threshold: int = 7,
    keep_participant_ids: list[int] | None = None,
    init_centers: list[list[float]] | None = None,
    max_group_count: int = 5,
    force_group_count: int | None = None,
    random_state: int | None = None,
    pick_max: int = 5,
    confidence: float = 0.9,
    consensus_mode: base.Literal["standard", "legacy"] = "legacy",
    candidate_group_counts: base.CandidateGroupCounts = None,
) -> base.TypedPolisPipelineResult:
    return base.run_pipeline(
        votes=list(votes),
        reducer=reducer,
        reducer_kwargs=reducer_kwargs or {},
        clusterer=clusterer,
        clusterer_kwargs=clusterer_kwargs or {},
        mod_out_statement_ids=mod_out_statement_ids or [],
        meta_statement_ids=meta_statement_ids or [],
        min_user_vote_threshold=min_user_vote_threshold,
        keep_participant_ids=keep_participant_ids or [],
        init_centers=init_centers,
        max_group_count=max_group_count,
        force_group_count=force_group_count,
        random_state=random_state,
        pick_max=pick_max,
        confidence=confidence,
        consensus_mode=consensus_mode,
        candidate_group_counts=candidate_group_counts,
    )

def generate_raw_matrix(
        votes: List[Dict],
        cutoff: Optional[int] = None,
) -> VoteMatrix:
    """
    Generates a raw vote matrix from a list of vote records.

    See `filter_votes` method for details of `cutoff` arg.

    Args:
        votes (List[Dict]): An unsorted list of vote records, where each record is a dictionary containing:

            - "participant_id": The ID of the voter.
            - "statement_id": The ID of the statement being voted on.
            - "vote": The recorded vote value.
            - "modified": A unix timestamp object representing when the vote was made.

        cutoff (int): A cutoff unix timestamp (ms) or index position in date-sorted votes list.

    Returns:
        raw_matrix (pd.DataFrame): A full raw vote matrix DataFrame with NaN values where:

            1. rows are voters,
            2. columns are statements, and
            3. values are votes.

            This includes even voters that have no votes, and statements on which no votes were placed.
    """
    if cutoff:
        votes = filter_votes(votes=votes, cutoff=cutoff)

    raw_matrix = pd.DataFrame.from_dict(votes)
    raw_matrix = raw_matrix.pivot(
        values="vote",
        index="participant_id",
        columns="statement_id",
    )

    # Ensure consistent column ordering regardless of statement_id type
    # Sort numerically if all columns can be converted to integers
    # If not all values can be converted to int, fall back to natural sorting
    try:
        sorted_columns = sorted(raw_matrix.columns, key=lambda x: int(x))
    except (ValueError, TypeError):
        sorted_columns = sorted(raw_matrix.columns)
    raw_matrix = raw_matrix.reindex(columns=sorted_columns)

    # Ensure consistent index
    try:
        sorted_index = sorted(raw_matrix.index, key=lambda x: int(x))
    except (ValueError, TypeError):
        sorted_index = sorted(raw_matrix.index)
    raw_matrix = raw_matrix.reindex(index=sorted_index)

    return raw_matrix

def simple_filter_matrix(
        vote_matrix: VoteMatrix,
        mod_out_statement_ids: list[int] = [],
) -> VoteMatrix:
    """
    The simple filter on the vote_matrix that is used by Polis prior to running PCA.

    Args:
        vote_matrix (VoteMatrix): A raw vote_matrix (with missing values)
        mod_out_statement_ids (list): A list of moderated-out participant IDs to zero out.

    Returns:
        VoteMatrix: Copy of vote_matrix with statements zero'd out
    """
    vote_matrix = vote_matrix.copy()
    for tid in mod_out_statement_ids:
        # Zero out column only if already exists (ie. has votes)
        if tid in vote_matrix.columns:
            # TODO: Add a flag to try np.nan instead of zero.
            vote_matrix.loc[:, tid] = 0

    return vote_matrix

def get_clusterable_participant_ids(vote_matrix: VoteMatrix, vote_threshold: int) -> list:
    """
    Find participant IDs that meet a vote threshold in a vote_matrix.

    Args:
        vote_matrix (VoteMatrix): A raw vote_matrix (with missing values)
        vote_threshold (int): Vote threshold that each participant must meet

    Returns:
        participation_ids (list): A list of participant IDs that meet the threshold
    """
    # TODO: Make this available outside this function? To match polismath output.
    user_vote_counts = vote_matrix.count(axis="columns")
    participant_ids = list(vote_matrix[user_vote_counts >= vote_threshold].index)
    return participant_ids

def run_reducer(
    vote_matrix: NDArray,
    reducer: ReducerType = "pca",
    n_components: int = 2,
    **reducer_kwargs,
) -> Tuple[NDArray, Optional[NDArray], ReducerModel]:
    """
    Process a prepared vote matrix to be imputed and return participant and (optionally) statement data,
    projected into reduced n-dimensional space.

    The vote matrix should not yet be imputed, as this will happen within the method.

    Args:
        vote_matrix (NDArray): A vote matrix of data. Non-imputed values are expected.
        n_components (int): Number n of principal components to decompose the `vote_matrix` into.
        reducer (Literal["pca", "pacmap", "localmap"]): Dimensionality reduction method to use.

    Returns:
        X_participants (NDArray): A numpy array with n-d coordinates for each projected row/participant.
        X_statements (Optional[NDArray]): A numpy array with n-d coordinates for each projected col/statement.
        reducer_model (ReducerModel): The fitted dimensional reduction sci-kit learn estimator.
    """
    match reducer:
        case "pca":
            pipeline = PatchedPipeline(
                [
                    ("capture", SparsityAwareCapturer()),
                    ("impute", SimpleImputer(missing_values=np.nan, strategy="mean")),
                    ("reduce", get_reducer(reducer, n_components=n_components, **reducer_kwargs)),
                    ("scale", SparsityAwareScaler(capture_step="capture")),
                ]
            )
        case "pacmap" | "localmap":
            pipeline = PatchedPipeline(
                [
                    ("impute", SimpleImputer(missing_values=np.nan, strategy="mean")),
                    ("reduce", get_reducer(reducer, n_components=n_components, **reducer_kwargs)),
                ]
            )

    # Generate projections of participants.
    X_participants = pipeline.fit_transform(vote_matrix)

    if reducer == "pca":
        # Generate projections of statements via virtual vote matrix.
        # This projects unit vectors for each feature/statement into PCA space to
        # understand their placement.
        num_cols = vote_matrix.shape[1]
        n_statements = num_cols
        virtual_vote_matrix = generate_virtual_vote_matrix(n_statements)
        X_statements = pipeline.transform(virtual_vote_matrix)
    else:
        X_statements = None

    reducer_model: ReducerModel = pipeline.named_steps["reduce"]

    return X_participants, X_statements, reducer_model

def get_reducer(
    reducer: ReducerType = "pca",
    n_components: int = 2,
    random_state: Optional[int] = None,
    **reducer_kwargs,
) -> ReducerModel:
    # Setting n_neighbors to None defaults to 10 below 10,000 samples, and
    # slowly increases it according to a formula beyond that.
    # See: https://github.com/YingfanWang/PaCMAP?tab=readme-ov-file#parameters
    DEFAULT_N_NEIGHBORS = None
    match reducer:
        case "pacmap" | "localmap":
            pacmap = try_import("pacmap", extra="alt-algos")

            # Override with default if not set
            n_neighbors = reducer_kwargs.pop("n_neighbors", DEFAULT_N_NEIGHBORS)

            ReducerCls = pacmap.PaCMAP if reducer == "pacmap" else pacmap.LocalMAP
            return ReducerCls(
                n_components=n_components,
                random_state=random_state,
                n_neighbors=n_neighbors,  # type:ignore
                **reducer_kwargs,
            )
        case "pca" | _:
            from sklearn.decomposition import PCA

            return PCA(
                n_components=n_components,
                random_state=random_state,
                **reducer_kwargs,
            )

def run_clusterer(
    X_participants_clusterable: NDArray,
    clusterer: ClustererType = "kmeans",
    force_group_count=None,
    max_group_count=5,
    init_centers=None,
    random_state=None,
    **clusterer_kwargs,
) -> Optional[ClustererModel]:
    match clusterer:
        case "kmeans":
            if force_group_count:
                k_bounds = [force_group_count, force_group_count]
            else:
                k_bounds = [2, max_group_count]

            _, _, kmeans = find_best_kmeans(
                X_to_cluster=X_participants_clusterable,
                k_bounds=k_bounds,
                # Force polis strategy of initiating cluster centers. See: PolisKMeans.
                init="polis",
                init_centers=init_centers,
                random_state=random_state,
                # TODO: Support passing in arbitrary clusterer_kwargs.
            )

            return kmeans

        case "hdbscan":
            hdbscan = try_import("hdbscan", extra="alt-algos")

            hdb = hdbscan.HDBSCAN(**clusterer_kwargs)
            hdb.fit(X_participants_clusterable)

            return hdb
        case _:
            raise NotImplementedError("clusterer type unknown")

def find_best_kmeans(
        X_to_cluster: NDArray,
        k_bounds: RangeLike = [2, 5],
        init: InitStrategy = "k-means++",
        init_centers: Optional[List] = None,
        random_state: Optional[int] = None,
) -> tuple[int, float, PolisKMeans | None]:
    """
    Use silhouette scores to find the best number of clusters k to assume to fit the data.

    Args:
        X_to_cluster (NDArray): A n-D numpy array.
        k_bounds (RangeLike): An upper and low bound on n_clusters to test for. (Default: [2, 5])
        init_centers (List): A list of xy coordinates to use as initial center guesses.
        random_state (int): Determines random number generation for centroid initialization. Use an int to make the randomness deterministic.

    Returns:
        best_k (int): Ideal number of clusters.
        best_silhouette_score (float): Silhouette score for this K value.
        best_kmeans (PolisKMeans | None): The optimal fitted estimator returned from PolisKMeans.
    """
    param_grid = {
        "n_clusters": to_range(k_bounds),
    }

    def scoring_function(estimator, X):
        labels = estimator.fit_predict(X)
        score = calculate_kmeans_silhouette_score(X_to_cluster=X, labels=labels)
        return -1 if score is None else score

    search = GridSearchNonCV(
        param_grid=param_grid,
        scoring=scoring_function,
        estimator=PolisKMeans(
            init=init, # strategy
            init_centers=init_centers, # guesses
            random_state=random_state,
        ),
    )

    search.fit(X_to_cluster)

    best_k = search.best_params_['n_clusters']
    best_silhouette_score = search.best_score_
    best_kmeans = search.best_estimator_

    return best_k, best_silhouette_score, best_kmeans

def select_consensus_statements(
    vote_matrix: VoteMatrix,
    mod_out_statement_ids: list[int] = [],
    pick_max=5,
    prob_threshold: float = 0.5,
    confidence: float = 0.9,
) -> ConsensusResult:
    """
    Select consensus statements from a given vote matrix.

    Args:
        vote_matrix (VoteMatrix): The full raw vote matrix (not just clusterable participants)
        mod_out_statement_ids (Optional[list[int]]): Statements to ignore from consensus statement selection
        pick_max (int): Max number of statements selected per agree/disagree direction
        prob_threshold (float): The cutoff probability below which statements won't be considered for consensus
        confidence (float): Percent confidence interval (in decimal), within which selected statements are deemed significant

    Returns:
        A dict of agree/disagree formatted statement dicts.
    """
    N_g_c, N_v_g_c, P_v_g_c, _, P_v_g_c_test, *_ = calculate_comment_statistics(
        vote_matrix=vote_matrix,
        cluster_labels=None,
    )
    # When no labels provided above, mock group is used.
    MOCK_GID = 0
    df = pd.DataFrame(
        {
            "na": N_v_g_c[votes.A, MOCK_GID, :],
            "nd": N_v_g_c[votes.D, MOCK_GID, :],
            "ns": N_g_c[MOCK_GID, :],
            "pa": P_v_g_c[votes.A, MOCK_GID, :],
            "pd": P_v_g_c[votes.D, MOCK_GID, :],
            "pat": P_v_g_c_test[votes.A, MOCK_GID, :],
            "pdt": P_v_g_c_test[votes.D, MOCK_GID, :],
            # agree metric = pa * pat
            "am": P_v_g_c[votes.A, MOCK_GID, :] * P_v_g_c_test[votes.A, MOCK_GID, :],
            # disagree metric = pd * pdt
            "dm": P_v_g_c[votes.D, MOCK_GID, :] * P_v_g_c_test[votes.D, MOCK_GID, :],
        },
        index=vote_matrix.columns,
    )

    # Optional filtering for mod_out_statement_ids
    if mod_out_statement_ids:
        df = df[~df.index.isin(mod_out_statement_ids)]

    # Agree candidates: pa > threshold and significant
    agree_mask = (df["pa"] > prob_threshold) & df["pat"].apply(
        lambda x: is_significant(x, confidence)
    )
    agree_candidates = df[agree_mask].copy()
    agree_candidates["consensus_agree_rank"] = (
        (-agree_candidates["am"]).rank(method="dense").astype("Int64")
    )

    # Disagree candidates: pd > threshold and significant
    disagree_mask = (df["pd"] > prob_threshold) & df["pdt"].apply(
        lambda x: is_significant(x, confidence)
    )
    disagree_candidates = df[disagree_mask].copy()
    disagree_candidates["consensus_disagree_rank"] = (
        (-disagree_candidates["dm"]).rank(method="dense").astype("Int64")
    )

    # Merge ranks back into full df
    df["consensus_agree_rank"] = agree_candidates["consensus_agree_rank"]
    df["consensus_disagree_rank"] = disagree_candidates["consensus_disagree_rank"]

    # Select top N agree/disagree statements
    if agree_candidates.empty:
        top_agree = []
    else:
        top_agree = [
            # Drop the cons-for key from final output.
            {k: v for k, v in st.items() if k != "cons-for"}
            for st in agree_candidates.sort_values("consensus_agree_rank")
            .head(pick_max)
            .reset_index()
            .apply(format_comment_stats, axis=1)
        ]

    if disagree_candidates.empty:
        top_disagree = []
    else:
        top_disagree = [
            # Drop the cons-for key from final output.
            {k: v for k, v in st.items() if k != "cons-for"}
            for st in disagree_candidates.sort_values("consensus_disagree_rank")
            .head(pick_max)
            .reset_index()
            .apply(format_comment_stats, axis=1)
        ]

    return {
        "agree": top_agree,
        "disagree": top_disagree,
    }

def rank_consensus_statements(
    vote_matrix: VoteMatrix,
    mod_out_statement_ids: list[int] = [],
    fdr_rate: float = 0.10,
) -> RankedConsensusResult:
    """
    Ranks ALL statements by consensus strength and marks selections via Benjamini-Hochberg.

    Args:
        vote_matrix (VoteMatrix): The full raw vote matrix.
        mod_out_statement_ids (list[int]): Statements to ignore.
        fdr_rate (float): False discovery rate for BH procedure.

    Returns:
        RankedConsensusResult with agree and disagree lists, each containing all statements
        ranked by effect size (pa or pd) with BH selection flags.
    """
    N_g_c, N_v_g_c, P_v_g_c, _, P_v_g_c_test, *_ = calculate_comment_statistics(
        vote_matrix=vote_matrix,
        cluster_labels=None,
    )
    MOCK_GID = 0
    df = pd.DataFrame(
        {
            "na": N_v_g_c[votes.A, MOCK_GID, :],
            "nd": N_v_g_c[votes.D, MOCK_GID, :],
            "ns": N_g_c[MOCK_GID, :],
            "pa": P_v_g_c[votes.A, MOCK_GID, :],
            "pd": P_v_g_c[votes.D, MOCK_GID, :],
            "pat": P_v_g_c_test[votes.A, MOCK_GID, :],
            "pdt": P_v_g_c_test[votes.D, MOCK_GID, :],
        },
        index=vote_matrix.columns,
    )

    if mod_out_statement_ids:
        df = df[~df.index.isin(mod_out_statement_ids)]

    def _rank_direction(
        df: pd.DataFrame,
        z_col: str,
        effect_col: str,
        direction: str,
    ) -> list[RankedConsensusStatement]:
        p_values = z_to_pvalue(df[z_col].values)
        effect_sizes = df[effect_col].values

        has_votes = (df["na"].values > 0) | (df["nd"].values > 0)
        selected_mask, adjusted = apply_bh_with_vote_filter(
            np.asarray(p_values), has_votes, fdr_rate
        )

        # Rank by effect size descending.
        n = len(effect_sizes)
        rank_order = np.argsort(-effect_sizes)
        ranks = np.empty(n, dtype=int)
        ranks[rank_order] = np.arange(1, n + 1)

        statements: list[RankedConsensusStatement] = []
        for idx in rank_order:
            row = df.iloc[idx]
            statements.append(
                RankedConsensusStatement(
                    statement_id=int(df.index[idx]),
                    direction=direction,
                    na=int(row["na"]),
                    nd=int(row["nd"]),
                    ns=int(row["ns"]),
                    pa=float(row["pa"]),
                    pd=float(row["pd"]),
                    pat=float(row["pat"]),
                    pdt=float(row["pdt"]),
                    effect_size=float(effect_sizes[idx]),
                    p_value=float(p_values[idx]),
                    adjusted_p_value=float(adjusted[idx]),
                    selected=bool(selected_mask[idx]),
                    rank=int(ranks[idx]),
                )
            )
        return statements

    agree_ranked = _rank_direction(df, "pat", "pa", "agree")
    disagree_ranked = _rank_direction(df, "pdt", "pd", "disagree")

    return RankedConsensusResult(agree=agree_ranked, disagree=disagree_ranked)

def select_representative_statements(
    grouped_stats_df: pd.DataFrame,
    mod_out_statement_ids: list[int] = [],
    pick_max: int = 5,
    confidence: float = 0.90,
) -> PolisRepness:
    """
    Selects statistically representative statements from each group cluster.

    This is expected to match the Polis outputs when all defaults are set.

    Args:
        grouped_stats_df (pd.DataFrame): MultiIndex Dataframe of statement statistics, indexed by group and statement.
        mod_out_statement_ids (list[int]): A list of statements to ignore from selection algorithm
        pick_max (int): Max number of statements selected per group
        confidence (float): Percent confidence interval (in decimal), within which selected statements are deemed significant

    Returns:
        PolisRepness: A dict object with lists of statements keyed to groups, matching Polis format.
    """
    repness = {}
    # TODO: Should this be done elsewhere? A column in MultiIndex dataframe?
    mod_out_mask = grouped_stats_df.index.get_level_values("statement_id").isin(
        mod_out_statement_ids
    )
    grouped_stats_df = grouped_stats_df[~mod_out_mask]  # type: ignore
    for gid, group_df in grouped_stats_df.groupby(level="group_id"):
        # Bring statement_id into regular column.
        group_df = group_df.reset_index()

        best_agree = None
        # Track the best-agree, to bring to top if exists.
        for _, row in group_df.iterrows():
            if beats_best_of_agrees(row, best_agree, confidence):
                best_agree = row

        sig_filter = lambda row: is_statement_significant(row, confidence)
        sufficient_statements_row_mask = group_df.apply(sig_filter, axis="columns")
        sufficient_statements = group_df[sufficient_statements_row_mask]

        # Track the best, even if doesn't meet sufficient minimum, to have at least one.
        best_overall = None
        if len(sufficient_statements) == 0:
            for _, row in group_df.iterrows():
                if beats_best_by_repness_test(row, best_overall):
                    best_overall = row
        else:
            # Finalize statements into output format.
            # TODO: Figure out how to finalize only at end in output. Change repness_metric?
            sufficient_statements = (
                pd.DataFrame(
                    [
                        format_comment_stats(row)
                        for _, row in sufficient_statements.iterrows()
                    ]
                )
                # Create a column to sort repnress, then remove.
                .assign(repness_metric=repness_metric)
                .sort_values(by="repness_metric", ascending=False)
                .drop(columns="repness_metric")
            )

        if best_agree is not None:
            best_agree = format_comment_stats(best_agree)
            best_agree.update({"n-agree": best_agree["n-success"], "best-agree": True})
            best_head = [best_agree]
        elif best_overall is not None:
            best_overall = format_comment_stats(best_overall)
            best_head = [best_overall]
        else:
            best_head = []

        selected = best_head
        selected = selected + [
            row.to_dict()
            for _, row in sufficient_statements.iterrows()
            if best_head
            # Skip any statements already in best_head
            and best_head[0]["tid"] != row["tid"]
        ]
        selected = selected[:pick_max]
        # Does the work of agrees-before-disagrees sort in polismath, since "a" before "d".
        selected = sorted(selected, key=lambda row: row["repful-for"])
        repness[gid] = selected

    return repness  # type:ignore

def rank_representative_statements(
    grouped_stats_df: pd.DataFrame,
    mod_out_statement_ids: list[int] = [],
    fdr_rate: float = 0.10,
) -> dict[int, list[RankedRepnessStatement]]:
    """
    Ranks ALL statements per group by effect size and marks selections via Benjamini-Hochberg.

    Args:
        grouped_stats_df (pd.DataFrame): MultiIndex DataFrame of statement statistics, indexed by group and statement.
        mod_out_statement_ids (list[int]): A list of statements to ignore.
        fdr_rate (float): False discovery rate for BH procedure.

    Returns:
        A dict mapping group_id to a list of RankedRepnessStatement, sorted by effect_size descending.
    """
    mod_out_mask = grouped_stats_df.index.get_level_values("statement_id").isin(
        mod_out_statement_ids
    )
    grouped_stats_df = grouped_stats_df[~mod_out_mask]  # type: ignore

    result: dict[int, list[RankedRepnessStatement]] = {}

    for gid, group_df in grouped_stats_df.groupby(level="group_id"):
        group_df = group_df.reset_index()

        # Per-direction p-values: probability test and representativeness test.
        p_prob_a = z_to_pvalue(group_df["pat"].values)
        p_rep_a = z_to_pvalue(group_df["rat"].values)
        p_prob_d = z_to_pvalue(group_df["pdt"].values)
        p_rep_d = z_to_pvalue(group_df["rdt"].values)

        # Combine probability + representativeness p-values using Simes' method,
        # which is valid under positive dependence (both tests use the same votes).
        # For 2 tests: p_simes = min(2 * p_min, p_max).
        p_simes_a = np.minimum(
            2 * np.minimum(p_prob_a, p_rep_a), np.maximum(p_prob_a, p_rep_a)
        )
        p_simes_d = np.minimum(
            2 * np.minimum(p_prob_d, p_rep_d), np.maximum(p_prob_d, p_rep_d)
        )

        # Pick the better direction (lower combined p-value).
        agree_better = p_simes_a <= p_simes_d
        p_combined = np.where(agree_better, p_simes_a, p_simes_d)
        directions = np.where(agree_better, "agree", "disagree")

        # Effect size: ra*pa for agree, rd*pd for disagree.
        effect_agree = group_df["ra"].values * group_df["pa"].values
        effect_disagree = group_df["rd"].values * group_df["pd"].values
        effect_sizes = np.where(agree_better, effect_agree, effect_disagree)

        # Apply BH with zero-vote filter.
        has_votes = (group_df["na"].values > 0) | (group_df["nd"].values > 0)
        selected_mask, adjusted = apply_bh_with_vote_filter(
            p_combined, has_votes, fdr_rate
        )

        # Rank by effect size descending (1-based).
        n = len(p_combined)
        rank_order = np.argsort(-effect_sizes)
        ranks = np.empty(n, dtype=int)
        ranks[rank_order] = np.arange(1, n + 1)

        # Build RankedRepnessStatement list, sorted by rank.
        statements: list[RankedRepnessStatement] = []
        for idx in rank_order:
            row = group_df.iloc[idx]
            statements.append(
                RankedRepnessStatement(
                    statement_id=int(row["statement_id"]),
                    repful_for=directions[idx],
                    na=int(row["na"]),
                    nd=int(row["nd"]),
                    ns=int(row["ns"]),
                    pa=float(row["pa"]),
                    pd=float(row["pd"]),
                    pat=float(row["pat"]),
                    pdt=float(row["pdt"]),
                    ra=float(row["ra"]),
                    rd=float(row["rd"]),
                    rat=float(row["rat"]),
                    rdt=float(row["rdt"]),
                    effect_size=float(effect_sizes[idx]),
                    p_value=float(p_combined[idx]),
                    adjusted_p_value=float(adjusted[idx]),
                    selected=bool(selected_mask[idx]),
                    rank=int(ranks[idx]),
                )
            )

        result[int(gid)] = statements

    return result