_adata_row_operations.py

# _adata_row_operations.py
# module at adata_science_tools/_preprocessing/_adata_row_operations.py
'''Functions for performing row-wise operations on AnnData objects, such as computing paired means and differences between groups of observations.'''
# updated 2026-03-13 added CFG_filter_adata_by_obs() function to filter adata by obs column values specified in config dict keys. Returns filtered adata object. Compatible with yaml file generated config dict.
#



# module imports
import logging
from typing import Optional, Sequence

import pandas as pd
import numpy as np
import anndata as ad


def CFG_filter_adata_by_obs(
    adata: ad.AnnData,
    dataset_cfg: Optional[dict] = None,
    filter_obs_boolean_column: Optional[str] = None,
    filter_obs_column_key: Optional[str] = None,
    filter_obs_column_values_list: Optional[Sequence] = None,
    copy: bool = True,
    logger: Optional[logging.Logger] = None,
    **kwargs,
) -> ad.AnnData:
    """
    Filter anndata object by obs column values specified in config dict keys
    or by explicit filter args.

    Accepts either `dataset_cfg` (a dict taken from the YAML dataset block)
    or explicit args: filter_obs_boolean_column, filter_obs_column_key,
    filter_obs_column_values_list. If dataset_cfg is provided, explicit args
    are used as fallbacks.

    Returns filtered AnnData (copy by default).
    """
    logger = logger or logging.getLogger(__name__)

    # prefer dataset_cfg keys if provided
    if dataset_cfg:
        filter_obs_boolean_column = dataset_cfg.get("filter_obs_boolean_column", filter_obs_boolean_column)
        filter_obs_column_key = dataset_cfg.get("filter_obs_column_key", filter_obs_column_key)
        filter_obs_column_values_list = dataset_cfg.get("filter_obs_column_values_list", filter_obs_column_values_list)

    # nothing to do
    if filter_obs_boolean_column is None and (filter_obs_column_key is None or filter_obs_column_values_list is None):
        logger.info("CFG_filter_adata_by_obs: no filtering keys provided; returning original adata (copy=%s)", copy)
        return adata.copy() if copy else adata

    # start from the original adata; apply filters in sequence with AND semantics
    _adata = adata

    if filter_obs_boolean_column is not None:
        if filter_obs_boolean_column not in _adata.obs.columns:
            raise KeyError(f"filter_obs_boolean_column '{filter_obs_boolean_column}' not found in adata.obs")
        # coerce to boolean
        mask_bool = _adata.obs[filter_obs_boolean_column].astype(bool)
        _adata = _adata[mask_bool, :].copy() if copy else _adata[mask_bool, :]
        logger.info("Filtered adata by boolean column '%s': %d observations remain", filter_obs_boolean_column, _adata.n_obs)

    if filter_obs_column_key is not None and filter_obs_column_values_list is not None:
        if filter_obs_column_key not in _adata.obs.columns:
            raise KeyError(f"filter_obs_column_key '{filter_obs_column_key}' not found in adata.obs")
        vals = list(filter_obs_column_values_list)

        # Try numeric-aware matching first; fall back to string matching
        try:
            vals_num = [float(v) for v in vals]
            col_num = pd.to_numeric(_adata.obs[filter_obs_column_key], errors="coerce")
            if not col_num.isna().all():
                mask_vals = col_num.isin(vals_num)
                if mask_vals.any():
                    _adata = _adata[mask_vals, :].copy() if copy else _adata[mask_vals, :]
                    logger.info("Applied numeric value filter on column '%s' (kept %d obs)", filter_obs_column_key, _adata.n_obs)
                else:
                    # no numeric matches -> try string
                    raise ValueError("No numeric matches found; falling back to string matching")
            else:
                raise ValueError("Column numeric conversion yields all NaN; falling back to string matching")
        except Exception:
            mask_vals = _adata.obs[filter_obs_column_key].astype(str).isin([str(v) for v in vals])
            _adata = _adata[mask_vals, :].copy() if copy else _adata[mask_vals, :]
            logger.info("Applied string value filter on column '%s' (kept %d obs)", filter_obs_column_key, _adata.n_obs)

    return _adata


def ref_vs_target_adata(
    adata: ad.AnnData,
    groupby_key: str = "Pre_or_Post_obs_col",
    groupby_key_target_value: str = "Post",
    groupby_key_ref_value: str = "Pre",
    opperation_flavor: str | Sequence[str] = "subtraction",
    obs_dfs: str = "merge",
    ref_obs_suffix: str = ".src_pre",
    target_obs_suffix: str = ".src_post",
    keep_var_df: bool = True,
    **params,
) -> ad.AnnData | tuple[ad.AnnData, pd.DataFrame]:
    """
    Compute paired target-vs-reference transforms for matched AnnData observations.

    Pairing is controlled by `pair_by_key` in `params`. The returned AnnData has
    one observation per overlapping pair ID, indexed by the stringified pair ID.
    `opperation_flavor` may be a string or sequence of operation strings.
    """
    original_params = dict(params)
    original_opperation_flavor = opperation_flavor
    params = dict(params)
    operation_flavor = params.pop("operation_flavor", None)
    if operation_flavor is not None:
        opperation_flavor = operation_flavor

    pair_by_key = params.pop("pair_by_key", None)
    layer = params.pop("layer", None)
    layers_to_compute = params.pop("layers_to_compute", None)
    base_layer_provided = "base_layer" in params
    base_layer = params.pop("base_layer", None)
    epsilon = params.pop("epsilon", 1e-9)
    target_min_value = params.pop("target_min_value", None)
    target_max_value = params.pop("target_max_value", None)
    ref_min_value = params.pop("ref_min_value", None)
    ref_max_value = params.pop("ref_max_value", None)
    bounds_fill_missing = params.pop("bounds_fill_missing", False)
    bounds_fill_missing_paired_only = params.pop("bounds_fill_missing_paired_only", False)
    merge_shared_obs_cols = params.pop("merge_shared_obs_cols", False)
    return_df = params.pop("return_df", False)
    allow_unused_params = params.pop("allow_unused_params", False)
    logger = params.pop("logger", None)
    log_level = params.pop("log_level", "INFO")
    save_source_values_obsm = params.pop("save_source_values_obsm", False)
    target_values_obsm_key = params.pop("target_values_obsm_key", "post_values")
    ref_values_obsm_key = params.pop("ref_values_obsm_key", "pre_values")
    select_max_ref_value_var_groupby_key = params.pop("select_max_ref_value_var_groupby_key", None)
    select_max_ref_value_filter_vars_by_isin_lists = params.pop(
        "select_max_ref_value_filter_vars_by_isin_lists",
        None,
    )
    select_max_ref_value_source_obsm_key = params.pop(
        "select_max_ref_value_source_obsm_key",
        "selected_source_variable",
    )

    log = logger or logging.getLogger(__name__)
    if log_level is not None:
        log.setLevel(log_level)

    if log.isEnabledFor(logging.INFO):
        adata_summary = {
            "shape": adata.shape,
            "X_type": type(adata.X).__name__,
            "layers": list(adata.layers.keys()),
            "obs_columns": list(adata.obs.columns),
            "var_columns": list(adata.var.columns),
        }
        args_items = [
            ("adata", adata_summary),
            ("groupby_key", groupby_key),
            ("groupby_key_target_value", groupby_key_target_value),
            ("groupby_key_ref_value", groupby_key_ref_value),
            ("opperation_flavor", original_opperation_flavor),
            ("operation_flavor", operation_flavor),
            ("obs_dfs", obs_dfs),
            ("ref_obs_suffix", ref_obs_suffix),
            ("target_obs_suffix", target_obs_suffix),
            ("keep_var_df", keep_var_df),
            ("pair_by_key", pair_by_key),
            ("layer", layer),
            ("layers_to_compute", layers_to_compute),
            ("base_layer", base_layer),
            ("epsilon", epsilon),
            ("target_min_value", target_min_value),
            ("target_max_value", target_max_value),
            ("ref_min_value", ref_min_value),
            ("ref_max_value", ref_max_value),
            ("bounds_fill_missing", bounds_fill_missing),
            ("bounds_fill_missing_paired_only", bounds_fill_missing_paired_only),
            ("merge_shared_obs_cols", merge_shared_obs_cols),
            ("return_df", return_df),
            ("allow_unused_params", allow_unused_params),
            ("logger", logger),
            ("log_level", log_level),
            ("save_source_values_obsm", save_source_values_obsm),
            ("target_values_obsm_key", target_values_obsm_key),
            ("ref_values_obsm_key", ref_values_obsm_key),
            ("select_max_ref_value_var_groupby_key", select_max_ref_value_var_groupby_key),
            ("select_max_ref_value_filter_vars_by_isin_lists", select_max_ref_value_filter_vars_by_isin_lists),
            ("select_max_ref_value_source_obsm_key", select_max_ref_value_source_obsm_key),
            ("unused_params", params),
            ("raw_params", original_params),
        ]
        args_lines = "\n".join([f"  {key}: {value}" for key, value in args_items])
        log.info("ref_vs_target_adata args:\n%s", args_lines)

    if params and not allow_unused_params:
        raise ValueError(f"Unused params: {sorted(params)}")

    if pair_by_key is None:
        raise ValueError("pair_by_key is required.")
    if groupby_key not in adata.obs.columns:
        raise KeyError(f"groupby_key '{groupby_key}' not found in adata.obs.")
    if pair_by_key not in adata.obs.columns:
        raise KeyError(f"pair_by_key '{pair_by_key}' not found in adata.obs.")

    operation_aliases = {
        "subtraction": "subtraction",
        "subtract": "subtraction",
        "difference": "subtraction",
        "diff": "subtraction",
        "target_minus_ref": "subtraction",
        "relative_change_pct": "relative_change_pct",
        "relative_change_percent": "relative_change_pct",
        "percent_change": "relative_change_pct",
        "pct": "relative_change_pct",
        "relative_change_fc": "relative_change_fc",
        "fold_change": "relative_change_fc",
        "foldchange": "relative_change_fc",
        "fc": "relative_change_fc",
        "relative_change_l2fc": "relative_change_l2fc",
        "log2_fold_change": "relative_change_l2fc",
        "log2fc": "relative_change_l2fc",
        "l2fc": "relative_change_l2fc",
    }
    if isinstance(opperation_flavor, (str, bytes)):
        requested_operations = [str(opperation_flavor)]
        multi_operation_mode = False
    else:
        try:
            requested_operations = [str(operation_name) for operation_name in list(opperation_flavor)]
            multi_operation_mode = True
        except TypeError:
            requested_operations = [str(opperation_flavor)]
            multi_operation_mode = False
    if not requested_operations:
        raise ValueError("opperation_flavor must be a non-empty string or sequence.")

    operations = []
    for requested_operation_name in requested_operations:
        operation_name = operation_aliases.get(requested_operation_name.lower())
        if operation_name is None:
            raise ValueError(f"Unsupported opperation_flavor '{requested_operation_name}'.")
        operations.append(operation_name)
    duplicate_operations = sorted({operation_name for operation_name in operations if operations.count(operation_name) > 1})
    if duplicate_operations:
        raise ValueError(
            "Duplicate opperation_flavor entries after alias normalization: "
            f"{duplicate_operations}"
        )
    requested_operation = requested_operations[0]
    operation = operations[0]
    log.info(
        "ref_vs_target_adata normalized operations: requested=%s, canonical=%s, multi_operation_mode=%s",
        requested_operations,
        operations,
        multi_operation_mode,
    )

    if layers_to_compute is None:
        sources_to_compute = [layer]
    elif isinstance(layers_to_compute, (str, bytes)):
        sources_to_compute = [layers_to_compute]
    else:
        sources_to_compute = list(layers_to_compute)
    if not sources_to_compute:
        raise ValueError("layers_to_compute must be a non-empty list or None.")
    if not base_layer_provided:
        base_layer = sources_to_compute[0]
    if base_layer not in sources_to_compute:
        raise ValueError("base_layer must be included in layers_to_compute or None.")
    for source in sources_to_compute:
        if source is not None and source not in adata.layers:
            raise KeyError(f"layer '{source}' not found in adata.layers.")

    def _source_label(source):
        return ".X" if source is None else str(source)

    def _source_layer_label(source):
        return "X" if source is None else str(source)

    source_labels = [_source_label(source) for source in sources_to_compute]
    log.info(
        "ref_vs_target_adata selected sources: layers_to_compute=%s, base_layer=%s",
        source_labels,
        _source_label(base_layer),
    )

    select_max_ref_value_enabled = select_max_ref_value_var_groupby_key is not None
    select_max_ref_value_filters = None
    if select_max_ref_value_filter_vars_by_isin_lists is not None:
        if not select_max_ref_value_enabled:
            raise ValueError(
                "select_max_ref_value_var_groupby_key is required when "
                "select_max_ref_value_filter_vars_by_isin_lists is provided."
            )
        select_max_ref_value_filters = {}
        for filter_key, filter_values in select_max_ref_value_filter_vars_by_isin_lists.items():
            if filter_key not in adata.var.columns:
                raise KeyError(
                    "Column "
                    f"'{filter_key}' from select_max_ref_value_filter_vars_by_isin_lists "
                    "not found in adata.var."
                )
            if isinstance(filter_values, (str, bytes)):
                select_max_ref_value_filters[filter_key] = [filter_values]
            else:
                select_max_ref_value_filters[filter_key] = list(filter_values)
    if select_max_ref_value_enabled:
        if select_max_ref_value_var_groupby_key not in adata.var.columns:
            raise KeyError(
                "Column "
                f"'{select_max_ref_value_var_groupby_key}' from "
                "select_max_ref_value_var_groupby_key not found in adata.var."
            )
        if (
            save_source_values_obsm
            and select_max_ref_value_source_obsm_key in {target_values_obsm_key, ref_values_obsm_key}
        ):
            raise ValueError(
                "select_max_ref_value_source_obsm_key must not match target_values_obsm_key "
                "or ref_values_obsm_key."
            )

    target_mask = adata.obs[groupby_key] == groupby_key_target_value
    ref_mask = adata.obs[groupby_key] == groupby_key_ref_value
    target_obs = adata.obs.loc[target_mask].copy()
    ref_obs = adata.obs.loc[ref_mask].copy()
    log.info(
        "ref_vs_target_adata selected observations: target_count=%d, ref_count=%d",
        int(target_mask.sum()),
        int(ref_mask.sum()),
    )

    target_missing_pair_ids = target_obs.index[target_obs[pair_by_key].isna()].tolist()
    ref_missing_pair_ids = ref_obs.index[ref_obs[pair_by_key].isna()].tolist()
    if target_missing_pair_ids or ref_missing_pair_ids:
        raise ValueError(
            "Missing pair IDs found in selected observations: "
            f"target={target_missing_pair_ids}, ref={ref_missing_pair_ids}"
        )

    target_pair_ids = target_obs[pair_by_key].astype(str)
    ref_pair_ids = ref_obs[pair_by_key].astype(str)
    target_duplicate_ids = sorted(target_pair_ids[target_pair_ids.duplicated(keep=False)].unique())
    ref_duplicate_ids = sorted(ref_pair_ids[ref_pair_ids.duplicated(keep=False)].unique())
    if target_duplicate_ids or ref_duplicate_ids:
        raise ValueError(
            "Duplicate pair IDs found in selected observations: "
            f"target={target_duplicate_ids}, ref={ref_duplicate_ids}"
        )

    target_pair_id_set = set(target_pair_ids.tolist())
    ref_pair_id_set = set(ref_pair_ids.tolist())
    matched_pair_ids = sorted(target_pair_id_set.intersection(ref_pair_id_set))
    if not matched_pair_ids:
        raise ValueError("No overlapping pair IDs found between target and reference observations.")

    dropped_target_only_pair_ids = sorted(target_pair_id_set.difference(ref_pair_id_set))
    dropped_ref_only_pair_ids = sorted(ref_pair_id_set.difference(target_pair_id_set))
    log.info(
        "ref_vs_target_adata matched pairs: n_matched=%d, dropped_target_only=%d, dropped_ref_only=%d",
        len(matched_pair_ids),
        len(dropped_target_only_pair_ids),
        len(dropped_ref_only_pair_ids),
    )
    if dropped_target_only_pair_ids or dropped_ref_only_pair_ids:
        log.info(
            "Dropping unmatched pair IDs: target_only=%s, ref_only=%s",
            dropped_target_only_pair_ids,
            dropped_ref_only_pair_ids,
        )

    target_obs_name_by_pair = pd.Series(target_obs.index.to_numpy(), index=target_pair_ids.to_numpy())
    ref_obs_name_by_pair = pd.Series(ref_obs.index.to_numpy(), index=ref_pair_ids.to_numpy())
    target_obs_names = target_obs_name_by_pair.loc[matched_pair_ids].to_numpy()
    ref_obs_names = ref_obs_name_by_pair.loc[matched_pair_ids].to_numpy()
    target_positions = adata.obs_names.get_indexer(target_obs_names)
    ref_positions = adata.obs_names.get_indexer(ref_obs_names)

    try:
        import scipy.sparse as sp

        def _is_sparse(value):
            return sp.issparse(value)
    except Exception:
        def _is_sparse(value):
            return False

    def _as_dense_array(value) -> np.ndarray:
        if _is_sparse(value):
            return value.toarray()
        return np.asarray(value)

    bounds_requested = any(
        value is not None
        for value in (target_min_value, target_max_value, ref_min_value, ref_max_value)
    )
    bounds_missing_requested = bounds_requested and (bounds_fill_missing or bounds_fill_missing_paired_only)

    def _apply_bounds(value, min_value, max_value):
        if min_value is None and max_value is None:
            return value
        bounded_value = _as_dense_array(value).astype(float, copy=True)
        if min_value is not None:
            bounded_value = np.maximum(bounded_value, min_value)
        if max_value is not None:
            bounded_value = np.minimum(bounded_value, max_value)
        return bounded_value

    def _get_source_matrix(source):
        if source is None:
            return adata.X
        return adata.layers[source]

    def _compute_source_values(source):
        matrix = _get_source_matrix(source)
        target_values = matrix[target_positions, :].copy()
        ref_values = matrix[ref_positions, :].copy()
        if bounds_requested:
            if bounds_missing_requested:
                target_values = _as_dense_array(target_values).astype(float, copy=True)
                ref_values = _as_dense_array(ref_values).astype(float, copy=True)
                target_missing_fill_value = target_min_value if target_min_value is not None else target_max_value
                ref_missing_fill_value = ref_min_value if ref_min_value is not None else ref_max_value
                if bounds_fill_missing_paired_only:
                    target_missing_mask = np.isnan(target_values) & ~np.isnan(ref_values)
                    ref_missing_mask = np.isnan(ref_values) & ~np.isnan(target_values)
                    if target_missing_fill_value is not None:
                        target_values[target_missing_mask] = target_missing_fill_value
                    if ref_missing_fill_value is not None:
                        ref_values[ref_missing_mask] = ref_missing_fill_value
                else:
                    if target_missing_fill_value is not None:
                        target_values[np.isnan(target_values)] = target_missing_fill_value
                    if ref_missing_fill_value is not None:
                        ref_values[np.isnan(ref_values)] = ref_missing_fill_value
            target_values = _apply_bounds(target_values, target_min_value, target_max_value)
            ref_values = _apply_bounds(ref_values, ref_min_value, ref_max_value)
        return target_values, ref_values

    matched_index = pd.Index(matched_pair_ids, name=pair_by_key)
    source_values_by_source = {}
    for source in sources_to_compute:
        source_label = _source_label(source)
        log.info("ref_vs_target_adata computing source '%s'.", source_label)
        source_values_by_source[source] = _compute_source_values(source)

    selected_group_values = []
    selected_group_labels = []
    select_max_ref_value_group_source_counts = []
    selected_source_positions = None
    selected_source_variable_df = None
    select_max_ref_value_tied_selection_count = 0
    if select_max_ref_value_enabled:
        var_mask = pd.Series(True, index=adata.var_names)
        for filter_key, filter_values in (select_max_ref_value_filters or {}).items():
            var_mask &= adata.var[filter_key].isin(filter_values)
        grouped_var = adata.var.loc[
            var_mask & adata.var[select_max_ref_value_var_groupby_key].notna()
        ]
        if grouped_var.empty:
            raise ValueError("No variables remain after select_max_ref_value filtering.")

        selected_group_values = list(pd.unique(grouped_var[select_max_ref_value_var_groupby_key]))
        selected_group_labels = [str(group_value) for group_value in selected_group_values]
        if len(set(selected_group_labels)) != len(selected_group_labels):
            raise ValueError("select_max_ref_value group labels must be unique after string conversion.")

        base_ref_values_for_selection = _as_dense_array(source_values_by_source[base_layer][1]).astype(float, copy=False)
        selected_source_positions = np.full(
            (len(matched_pair_ids), len(selected_group_labels)),
            -1,
            dtype=int,
        )
        selected_source_variable_values = np.full(
            selected_source_positions.shape,
            "",
            dtype=object,
        )
        source_var_names = np.asarray(adata.var_names)
        for group_idx, group_value in enumerate(selected_group_values):
            group_var_names = grouped_var.index[
                grouped_var[select_max_ref_value_var_groupby_key] == group_value
            ]
            group_var_positions = adata.var_names.get_indexer(group_var_names)
            select_max_ref_value_group_source_counts.append(len(group_var_positions))
            group_ref_values = base_ref_values_for_selection[:, group_var_positions]
            valid_ref_values = ~np.isnan(group_ref_values)
            pairs_with_values = valid_ref_values.any(axis=1)
            if not pairs_with_values.any():
                continue

            masked_ref_values = np.where(valid_ref_values, group_ref_values, -np.inf)
            local_selected_positions = np.argmax(masked_ref_values, axis=1)
            pair_indices = np.flatnonzero(pairs_with_values)
            selected_var_positions = group_var_positions[local_selected_positions[pair_indices]]
            selected_source_positions[pair_indices, group_idx] = selected_var_positions
            selected_source_variable_values[pair_indices, group_idx] = source_var_names[selected_var_positions]

            selected_ref_values = masked_ref_values[pair_indices, local_selected_positions[pair_indices]]
            tied_counts = (
                (masked_ref_values[pair_indices] == selected_ref_values[:, None])
                & valid_ref_values[pair_indices]
            ).sum(axis=1)
            select_max_ref_value_tied_selection_count += int((tied_counts > 1).sum())

        selected_source_variable_df = pd.DataFrame(
            selected_source_variable_values,
            index=matched_index,
            columns=selected_group_labels,
        )
        log.info(
            "ref_vs_target_adata select_max_ref_value enabled: groups=%d, tied_selections=%d.",
            len(selected_group_labels),
            select_max_ref_value_tied_selection_count,
        )

    def _select_max_ref_value_columns(values):
        dense_values = _as_dense_array(values).astype(float, copy=False)
        collapsed_values = np.full(selected_source_positions.shape, np.nan, dtype=float)
        pair_indices = np.arange(selected_source_positions.shape[0])
        for group_idx in range(selected_source_positions.shape[1]):
            source_positions = selected_source_positions[:, group_idx]
            pairs_with_selection = source_positions >= 0
            collapsed_values[pairs_with_selection, group_idx] = dense_values[
                pair_indices[pairs_with_selection],
                source_positions[pairs_with_selection],
            ]
        return collapsed_values

    multi_source_mode = len(sources_to_compute) > 1
    layer_results = {}
    operation_layer_results = {}
    operation_layer_keys = []
    operation_layer_key_by_source_operation = {}
    generated_layer_keys = set()
    base_operation_layer = None
    base_result_values = None
    base_target_values = None
    base_ref_values = None
    for source in sources_to_compute:
        source_label = _source_label(source)
        target_values, ref_values = source_values_by_source[source]
        if select_max_ref_value_enabled:
            target_values = _select_max_ref_value_columns(target_values)
            ref_values = _select_max_ref_value_columns(ref_values)
        dense_target_values = None
        dense_ref_values = None
        source_operation_layer_keys = {}
        for operation_name in operations:
            if operation_name == "subtraction":
                result_values = target_values - ref_values
            else:
                if dense_target_values is None:
                    dense_target_values = _as_dense_array(target_values).astype(float, copy=False)
                    dense_ref_values = _as_dense_array(ref_values).astype(float, copy=False)
                with np.errstate(divide="ignore", invalid="ignore"):
                    if operation_name == "relative_change_pct":
                        result_values = ((dense_target_values - dense_ref_values) / (dense_ref_values + epsilon)) * 100
                    elif operation_name == "relative_change_fc":
                        result_values = (dense_target_values + epsilon) / (dense_ref_values + epsilon)
                    else:
                        result_values = np.log2((dense_target_values + epsilon) / (dense_ref_values + epsilon))

            if multi_operation_mode:
                layer_key = operation_name
                if multi_source_mode:
                    layer_key = f"{_source_layer_label(source)}__{operation_name}"
                if layer_key in generated_layer_keys:
                    raise ValueError(f"Generated layer key collision: '{layer_key}'.")
                generated_layer_keys.add(layer_key)
                operation_layer_results[layer_key] = result_values
                operation_layer_keys.append(layer_key)
                source_operation_layer_keys[operation_name] = layer_key
                if source == base_layer and operation_name == operation:
                    base_result_values = result_values
                    base_operation_layer = layer_key
            else:
                layer_results[source] = result_values

        log.info(
            "ref_vs_target_adata computed source '%s' for operations=%s.",
            source_label,
            operations,
        )
        if source == base_layer:
            if not multi_operation_mode:
                base_result_values = layer_results[source]
            base_target_values = target_values
            base_ref_values = ref_values
        if multi_operation_mode:
            operation_layer_key_by_source_operation[_source_label(source)] = source_operation_layer_keys

    if multi_operation_mode:
        log.info("ref_vs_target_adata generated operation layer keys: %s", operation_layer_keys)

    target_aligned_obs = adata.obs.loc[target_obs_names].copy()
    ref_aligned_obs = adata.obs.loc[ref_obs_names].copy()
    target_aligned_obs.index = matched_index
    ref_aligned_obs.index = matched_index

    if obs_dfs == "merge":
        result_obs_data = {}
        for column in ref_aligned_obs.columns:
            ref_values = ref_aligned_obs[column]
            target_values = target_aligned_obs[column]
            if merge_shared_obs_cols and ref_values.equals(target_values):
                result_obs_data[column] = ref_values.to_numpy()
            else:
                result_obs_data[f"{column}{ref_obs_suffix}"] = ref_values.to_numpy()
                result_obs_data[f"{column}{target_obs_suffix}"] = target_values.to_numpy()
        if pair_by_key not in result_obs_data:
            result_obs_data = {pair_by_key: matched_pair_ids, **result_obs_data}
        else:
            result_obs_data[pair_by_key] = matched_pair_ids
        result_obs = pd.DataFrame(result_obs_data, index=matched_index)
    elif obs_dfs == "keep_ref":
        result_obs = ref_aligned_obs.copy()
        result_obs[pair_by_key] = matched_pair_ids
    elif obs_dfs == "keep_target":
        result_obs = target_aligned_obs.copy()
        result_obs[pair_by_key] = matched_pair_ids
    else:
        raise ValueError("obs_dfs must be 'merge', 'keep_ref', or 'keep_target'.")

    result_obs["ref_obs_name"] = ref_obs_names
    result_obs["target_obs_name"] = target_obs_names
    result_obs["pair_order"] = np.arange(len(matched_pair_ids), dtype=int)
    result_obs["ref_groupby_value"] = groupby_key_ref_value
    result_obs["target_groupby_value"] = groupby_key_target_value
    result_obs["ref_vs_target_operation"] = operation

    if select_max_ref_value_enabled:
        result_var = pd.DataFrame(
            index=pd.Index(selected_group_labels, name=adata.var_names.name),
        )
        if keep_var_df:
            result_var[select_max_ref_value_var_groupby_key] = selected_group_values
            result_var["select_max_ref_value_n_source_variables"] = select_max_ref_value_group_source_counts
    else:
        result_var = adata.var.copy() if keep_var_df else pd.DataFrame(index=adata.var_names.copy())
    result_var["ref_vs_target_operation"] = operation
    result_var["ref_vs_target_groupby_key"] = groupby_key
    result_var["ref_vs_target_target_value"] = groupby_key_target_value
    result_var["ref_vs_target_ref_value"] = groupby_key_ref_value
    result_var["ref_vs_target_pair_by_key"] = pair_by_key

    result_adata = ad.AnnData(
        X=base_result_values,
        obs=result_obs,
        var=result_var,
    )
    if multi_operation_mode:
        for layer_key, result_values in operation_layer_results.items():
            result_adata.layers[layer_key] = result_values
    else:
        for source in sources_to_compute:
            if source is not None:
                result_adata.layers[source] = layer_results[source]

    log.info(
        "ref_vs_target_adata final adata.X selected: base_layer=%s, base_operation=%s, "
        "operation_layer_key=%s, shape=%s, dtype=%s",
        _source_label(base_layer),
        operation,
        base_operation_layer,
        result_adata.X.shape,
        getattr(result_adata.X, "dtype", None),
    )

    result_adata.uns["ref_vs_target_adata"] = {
        "pair_by_key": pair_by_key,
        "groupby_key": groupby_key,
        "groupby_key_target_value": groupby_key_target_value,
        "groupby_key_ref_value": groupby_key_ref_value,
        "requested_operation_flavor": requested_operation,
        "requested_operation_flavors": requested_operations,
        "operation_flavor": operation,
        "operation_flavors": operations,
        "base_operation": operation,
        "base_operation_layer": base_operation_layer,
        "operation_layer_keys": operation_layer_keys,
        "operation_layer_key_by_source_operation": operation_layer_key_by_source_operation,
        "obs_dfs": obs_dfs,
        "ref_obs_suffix": ref_obs_suffix,
        "target_obs_suffix": target_obs_suffix,
        "keep_var_df": keep_var_df,
        "layer": _source_label(layer),
        "layers_to_compute": source_labels,
        "base_layer": _source_label(base_layer),
        "epsilon": epsilon,
        "target_min_value": target_min_value,
        "target_max_value": target_max_value,
        "ref_min_value": ref_min_value,
        "ref_max_value": ref_max_value,
        "bounds_fill_missing": bounds_fill_missing,
        "bounds_fill_missing_paired_only": bounds_fill_missing_paired_only,
        "merge_shared_obs_cols": merge_shared_obs_cols,
        "n_target_obs": int(target_mask.sum()),
        "n_ref_obs": int(ref_mask.sum()),
        "n_matched_pairs": len(matched_pair_ids),
        "matched_pair_ids": matched_pair_ids,
        "dropped_target_only_pair_ids": dropped_target_only_pair_ids,
        "dropped_ref_only_pair_ids": dropped_ref_only_pair_ids,
        "save_source_values_obsm": save_source_values_obsm,
        "target_values_obsm_key": target_values_obsm_key if save_source_values_obsm else None,
        "ref_values_obsm_key": ref_values_obsm_key if save_source_values_obsm else None,
        "select_max_ref_value_enabled": select_max_ref_value_enabled,
        "select_max_ref_value_var_groupby_key": select_max_ref_value_var_groupby_key,
        "select_max_ref_value_filter_vars_by_isin_lists": select_max_ref_value_filters,
        "select_max_ref_value_source_obsm_key": (
            select_max_ref_value_source_obsm_key if select_max_ref_value_enabled else None
        ),
        "select_max_ref_value_n_output_groups": len(selected_group_labels),
        "select_max_ref_value_tied_selection_count": select_max_ref_value_tied_selection_count,
    }
    result_adata.uns["ref_vs_target_dropped_target_only_pair_ids"] = dropped_target_only_pair_ids
    result_adata.uns["ref_vs_target_dropped_ref_only_pair_ids"] = dropped_ref_only_pair_ids

    if save_source_values_obsm:
        log.info(
            "ref_vs_target_adata saving source values to obsm: target_key=%s, ref_key=%s",
            target_values_obsm_key,
            ref_values_obsm_key,
        )
        result_adata.obsm[target_values_obsm_key] = pd.DataFrame(
            _as_dense_array(base_target_values),
            index=result_adata.obs_names.copy(),
            columns=result_adata.var_names.copy(),
        )
        result_adata.obsm[ref_values_obsm_key] = pd.DataFrame(
            _as_dense_array(base_ref_values),
            index=result_adata.obs_names.copy(),
            columns=result_adata.var_names.copy(),
        )

    if select_max_ref_value_enabled:
        result_adata.obsm[select_max_ref_value_source_obsm_key] = selected_source_variable_df

    if return_df:
        log.info("ref_vs_target_adata returning result AnnData and DataFrame.")
        result_df = pd.DataFrame(
            _as_dense_array(base_result_values),
            index=result_adata.obs_names.copy(),
            columns=result_adata.var_names.copy(),
        )
        return result_adata, result_df
    return result_adata


def compute_paired_mean_adata(
    adata,
    layer='RFU',
    pair_by_key='AnimalID_Tattoo',
    groupby_key='Treatment_unique',
    datapoint_1='drug78hr',
    datapoint_2='drug30hr',
    debug_mode=True,
    layers_to_compute=None,
    base_layer=None,
):
    """
    Compute the average of two selected groups in an AnnData object and return a new AnnData object.
    
    Parameters:
    - adata: AnnData object
    - layer: The layer to use (default: 'RFU'). If None, uses .X
    - layers_to_compute: Optional list of layers to compute the mean for.
      If provided, the `layer` argument is ignored.
    - base_layer: The layer to use as `.X` in the returned AnnData when
      `layers_to_compute` is provided. Defaults to the first entry in
      `layers_to_compute`. Use None to base on `.X`.
    - pair_by_key: The key in .obs to pair the data
    - groupby_key: The key in .obs used to select groups
    - datapoint_1: First group to compute the mean
    - datapoint_2: Second group to compute the mean
    - debug_mode: If True, prints debug information
    
    Returns:
    - adata_mean_datapoint_12: AnnData object with averaged values
    - mean_datapoint_12_df: DataFrame of averaged values
    """
    _adata = adata.copy()

    multi_layer_mode = layers_to_compute is not None
    if layers_to_compute is None:
        layers_to_compute = [layer]
    if not isinstance(layers_to_compute, list) or len(layers_to_compute) == 0:
        raise ValueError("layers_to_compute must be a non-empty list or None.")
    if base_layer is None:
        base_layer = layers_to_compute[0]
    if base_layer is not None and base_layer not in layers_to_compute:
        raise ValueError("base_layer must be included in layers_to_compute or None.")

    for layer_name in layers_to_compute:
        if layer_name is not None and layer_name not in _adata.layers:
            raise KeyError(f"layer '{layer_name}' not found in adata.layers.")
    
    # Boolean indexing to select cells in each group
    datapoint_1_idx = _adata.obs[groupby_key].isin([datapoint_1])
    datapoint_2_idx = _adata.obs[groupby_key].isin([datapoint_2])
    
    # Ensure `pair_by_key` is categorical for proper sorting
    if not isinstance(_adata.obs[pair_by_key].dtype, pd.CategoricalDtype):
        _adata.obs[pair_by_key] = _adata.obs[pair_by_key].astype('category')

    order_1 = np.argsort(_adata.obs.loc[datapoint_1_idx, pair_by_key].cat.codes)
    order_2 = np.argsort(_adata.obs.loc[datapoint_2_idx, pair_by_key].cat.codes)

    obs_names1 = _adata[datapoint_1_idx].obs_names.to_numpy()[order_1]
    obs_names2 = _adata[datapoint_2_idx].obs_names.to_numpy()[order_2]
    obs_names = 'avg_' + obs_names1 + '_' + obs_names2

    def _get_layer_matrix(layer_name):
        if layer_name is None:
            return _adata.X
        return _adata.layers[layer_name]

    layer_results = {}
    for layer_name in layers_to_compute:
        X = _get_layer_matrix(layer_name)
        # Select the data matrix for each group
        data_datapoint_1 = X[datapoint_1_idx, :].copy()
        data_datapoint_2 = X[datapoint_2_idx, :].copy()

        if debug_mode and layer_name == base_layer:
            display('data_datapoint_1:', data_datapoint_1.shape, data_datapoint_1[:11,:2])
            display('data_datapoint_2:', data_datapoint_2.shape, data_datapoint_2[:11,:2])

        # Sort the data by `pair_by_key`
        data_datapoint_1_rel = data_datapoint_1[order_1]
        data_datapoint_2_rel = data_datapoint_2[order_2]

        if debug_mode and layer_name == base_layer:
            display('data_datapoint_1_rel:', data_datapoint_1_rel.shape, data_datapoint_1_rel[:11,:2])
            display('data_datapoint_2_rel:', data_datapoint_2_rel.shape, data_datapoint_2_rel[:11,:2])

        # Compute the mean
        mean_datapoint_12 = (data_datapoint_1_rel + data_datapoint_2_rel) / 2
        layer_results[layer_name] = mean_datapoint_12

        if debug_mode and layer_name == base_layer:
            display('mean_datapoint_12:', mean_datapoint_12.shape, mean_datapoint_12[:11,:2])

    base_matrix = layer_results[base_layer]

    # Convert base layer to DataFrame
    mean_datapoint_12_df = pd.DataFrame(base_matrix, columns=_adata.var.index)
    mean_datapoint_12_df['obs_names1'] = obs_names1
    mean_datapoint_12_df['obs_names2'] = obs_names2
    mean_datapoint_12_df['obs_names'] = obs_names
    mean_datapoint_12_df.set_index('obs_names', inplace=True)
    mean_datapoint_12_df.drop(columns=['obs_names1', 'obs_names2'], inplace=True)

    if debug_mode:
        display(mean_datapoint_12_df.head(5))

    # Create new AnnData object
    adata_mean_datapoint_12 = ad.AnnData(X=mean_datapoint_12_df)
    adata_mean_datapoint_12.var_names = _adata.var_names
    adata_mean_datapoint_12.var = _adata.var.copy()

    if multi_layer_mode:
        for layer_name in layers_to_compute:
            if layer_name is None:
                continue
            adata_mean_datapoint_12.layers[layer_name] = layer_results[layer_name]

    if debug_mode:
        display(adata_mean_datapoint_12[:11,:2])
    
    return adata_mean_datapoint_12, mean_datapoint_12_df


def compute_paired_difference_adata(
    adata,
    layer='RFU',
    pair_by_key='AnimalID_Tattoo',
    groupby_key='Treatment_unique',
    datapoint_1='drug78hr',
    datapoint_2='drug30hr',
    debug_mode=True,
    layers_to_compute=None,
    base_layer=None,
):
    """
    Compute the paired difference (datapoint_1 - datapoint_2) in an AnnData object and return a new AnnData object.
    
    Parameters:
    - adata: AnnData object
    - layer: The layer to use (default: 'RFU'). If None, uses .X
    - layers_to_compute: Optional list of layers to compute the difference for.
      If provided, the `layer` argument is ignored.
    - base_layer: The layer to use as `.X` in the returned AnnData when
      `layers_to_compute` is provided. Defaults to the first entry in
      `layers_to_compute`. Use None to base on `.X`.
    - pair_by_key: The key in .obs to pair the data
    - groupby_key: The key in .obs used to select groups
    - datapoint_1: Group treated as minuend
    - datapoint_2: Group treated as subtrahend
    - debug_mode: If True, prints debug information
    
    Returns:
    - adata_diff_datapoint_12: AnnData object with difference values
    - diff_datapoint_12_df: DataFrame of difference values
    """
    _adata = adata.copy()

    multi_layer_mode = layers_to_compute is not None
    if layers_to_compute is None:
        layers_to_compute = [layer]
    if not isinstance(layers_to_compute, list) or len(layers_to_compute) == 0:
        raise ValueError("layers_to_compute must be a non-empty list or None.")
    if base_layer is None:
        base_layer = layers_to_compute[0]
    if base_layer is not None and base_layer not in layers_to_compute:
        raise ValueError("base_layer must be included in layers_to_compute or None.")

    for layer_name in layers_to_compute:
        if layer_name is not None and layer_name not in _adata.layers:
            raise KeyError(f"layer '{layer_name}' not found in adata.layers.")
    
    datapoint_1_idx = _adata.obs[groupby_key].isin([datapoint_1])
    datapoint_2_idx = _adata.obs[groupby_key].isin([datapoint_2])
    
    if not isinstance(_adata.obs[pair_by_key].dtype, pd.CategoricalDtype):
        _adata.obs[pair_by_key] = _adata.obs[pair_by_key].astype('category')

    order_1 = np.argsort(_adata.obs.loc[datapoint_1_idx, pair_by_key].cat.codes)
    order_2 = np.argsort(_adata.obs.loc[datapoint_2_idx, pair_by_key].cat.codes)

    obs_names1 = _adata[datapoint_1_idx].obs_names.to_numpy()[order_1]
    obs_names2 = _adata[datapoint_2_idx].obs_names.to_numpy()[order_2]
    obs_names = 'diff_' + obs_names1 + '_' + obs_names2

    def _get_layer_matrix(layer_name):
        if layer_name is None:
            return _adata.X
        return _adata.layers[layer_name]

    layer_results = {}
    for layer_name in layers_to_compute:
        X = _get_layer_matrix(layer_name)
        data_datapoint_1 = X[datapoint_1_idx, :].copy()
        data_datapoint_2 = X[datapoint_2_idx, :].copy()

        if debug_mode and layer_name == base_layer:
            display('data_datapoint_1:', data_datapoint_1.shape, data_datapoint_1[:11,:2])
            display('data_datapoint_2:', data_datapoint_2.shape, data_datapoint_2[:11,:2])

        data_datapoint_1_rel = data_datapoint_1[order_1]
        data_datapoint_2_rel = data_datapoint_2[order_2]

        if debug_mode and layer_name == base_layer:
            display('data_datapoint_1_rel:', data_datapoint_1_rel.shape, data_datapoint_1_rel[:11,:2])
            display('data_datapoint_2_rel:', data_datapoint_2_rel.shape, data_datapoint_2_rel[:11,:2])

        diff_datapoint_12 = data_datapoint_1_rel - data_datapoint_2_rel
        layer_results[layer_name] = diff_datapoint_12

        if debug_mode and layer_name == base_layer:
            display('diff_datapoint_12:', diff_datapoint_12.shape, diff_datapoint_12[:11,:2])

    base_matrix = layer_results[base_layer]

    diff_datapoint_12_df = pd.DataFrame(base_matrix, columns=_adata.var.index)
    diff_datapoint_12_df['obs_names1'] = obs_names1
    diff_datapoint_12_df['obs_names2'] = obs_names2
    diff_datapoint_12_df['obs_names'] = obs_names
    diff_datapoint_12_df.set_index('obs_names', inplace=True)
    diff_datapoint_12_df.drop(columns=['obs_names1', 'obs_names2'], inplace=True)

    if debug_mode:
        display(diff_datapoint_12_df.head(5))

    adata_diff_datapoint_12 = ad.AnnData(X=diff_datapoint_12_df)
    adata_diff_datapoint_12.var_names = _adata.var_names
    adata_diff_datapoint_12.var = _adata.var.copy()

    if multi_layer_mode:
        for layer_name in layers_to_compute:
            if layer_name is None:
                continue
            adata_diff_datapoint_12.layers[layer_name] = layer_results[layer_name]

    if debug_mode:
        display(adata_diff_datapoint_12[:11,:2])
    
    return adata_diff_datapoint_12, diff_datapoint_12_df

'''_adata = adata.copy()
adata_mean_datapoint_12, mean_datapoint_12_df = compute_mean_adata(
    adata=_adata,
    layer='Volume_norm',
    pair_by_key='AnimalID_Tattoo',
    groupby_key='Treatment_unique',
    datapoint_1='drug78hr',
    datapoint_2='drug30hr',
    debug_mode=False
)
display(adata_mean_datapoint_12[:11,:2], adata_mean_datapoint_12.obs.head(5))
display(mean_datapoint_12_df.head(5))'''