"""General utilities for preprocessing and statistical helpers."""

from __future__ import annotations

import warnings
import json
from typing import Any, Optional, Literal

from pathlib import Path
import numpy as np
from numpy.typing import ArrayLike
import scipy.sparse
from tqdm import tqdm
import pandas as pd
import torch
from matplotlib.image import imread
from anndata import AnnData
from smoother import SpatialWeightMatrix, SpatialLoss
from splisosm.likelihood import liu_sf

__all__ = [
    "get_cov_sp",
    "counts_to_ratios",
    "false_discovery_control",
    "load_visium_sp_meta",
    "prepare_inputs_from_anndata",
    "extract_counts_n_ratios",
    "extract_gene_level_statistics",
    "run_sparkx",
    "run_hsic_gc",
]


def get_cov_sp(
    coords: np.ndarray | torch.Tensor, k: int = 4, rho: float = 0.99
) -> torch.Tensor:
    """Wrapper function to get the spatial covariance matrix from spatial coordinates.

    It will first construct a mutual-k-nearest neighbor graph from the euclidean spatial coordinates,
    then convert the adjacency matrix to a standardized spatial covariance matrix using the
    intrinsic conditional autoregressive (ICAR) model with spatial autocorrelation coefficient rho.
    See :cite:`su2023smoother` for details.

    Parameters
    ----------
    coords
        Shape (n_spots, n_dims). Euclidean spatial coordinates of spots.
    k
        Number of nearest neighbors.
    rho
        Spatial autocorrelation coefficient.

    Returns
    -------
    cov_sp : torch.Tensor
        Shape (n_spots, n_spots). Spatial covariance matrix with standardized variance (== 1).
    """
    # first calculate the spatial weights matrix (swm)
    # here swm is the binary adjacency matrix of the knn graph
    weights = SpatialWeightMatrix()
    weights.calc_weights_knn(coords, k=k, verbose=False)

    # # convert the swm to spatial covariance matrix with standardized variance (== 1)
    # spatial_loss = SpatialLoss("icar", weights, rho=rho, standardize_cov=True)
    # cov_sp = torch.cholesky_inverse(torch.linalg.cholesky(spatial_loss.inv_cov[0].to_dense())) # n_spots x n_spots
    spatial_loss = SpatialLoss("icar", weights, rho=rho, standardize_cov=False)
    cov_sp = torch.cholesky_inverse(
        torch.linalg.cholesky(spatial_loss.inv_cov[0].to_dense())
    )  # n_spots x n_spots
    inv_sds = torch.diagflat(torch.diagonal(cov_sp) ** (-0.5))
    cov_sp = inv_sds @ cov_sp @ inv_sds

    return cov_sp


def counts_to_ratios(
    counts: np.ndarray | torch.Tensor,
    transformation: Literal["none", "clr", "ilr", "alr", "radial"] = "none",
    nan_filling: Literal["mean", "none"] = "mean",
) -> torch.Tensor:
    """Convert isoform counts to proportions.

    By default, isoform ratios at zero-coverage spots are filled with the mean ratio per isoform across all spots.
    After conversion, the isoform ratios can be further transformed using log-ratio-based transformations
    (clr, ilr, alr) or radial transformation :cite:`park2022kernel`.

    Parameters
    ----------
    counts
        Shape (n_spots, n_isos). Isoform counts.
    transformation
        Transformation applied to the proportions. Can be one of the following:
        ``'none'``: no transformation, return isoform ratios.
        ``'clr'``: centered log-ratio transformation.
        ``'ilr'``: isometric log-ratio transformation.
        ``'alr'``: additive log-ratio transformation.
        ``'radial'``: radial transformation :cite:`park2022kernel`.
    nan_filling
        Method to fill all-zero rows.
        ``'mean'``: fill all-zero rows with the mean of the mean per column **before transformation**.
        ``'none'``: do not fill rows and return NaNs at all-zero rows.

    Returns
    -------
    ratios : torch.Tensor
        Shape (n_spots, n_isos) or (n_spots, n_isos - 1) if ilr or alr transformation is used.

    Notes
    -----
    Log-ratio-based transformations (clr, ilr, alr) are implemented via ``scikit-bio``, with
    a pseudocount of 1% of the global mean per isoform to avoid zeros in the ratio.
    """
    assert transformation in ["none", "clr", "ilr", "alr", "radial"]
    if transformation in ["clr", "ilr", "alr"]:
        try:
            from skbio.stats.composition import (
                clr,
                ilr,
                alr,
            )  # for ratio transformation
        except ImportError:
            warnings.warn(
                f"Please install scikit-bio to use ratio transformation='{transformation}'. Switching to 'none'."
            )
            transformation = "none"

    assert nan_filling in ["mean", "none"]

    if isinstance(counts, np.ndarray):
        counts = torch.from_numpy(counts).float()

    # identify zero rows to fill
    is_nan = counts.sum(1) == 0  # (n_spots,)

    # calculate isoform ratios
    if transformation in ["clr", "ilr", "alr"]:
        # add pseudocounts equal to 1% of the global mean per isoform to avoid zeros in the ratio
        y = (1 - 1e-2) * counts + 1e-2 * counts.mean(0, keepdim=True)
        y = y / y.sum(1, keepdim=True)  # isoform ratio without nans and zeros
    else:
        y = counts / counts.sum(1, keepdim=True)  # isoform ratio with nans
        # fill nan values with the mean ratio per column (isoform)
        if nan_filling == "mean":
            y[is_nan] = y[~is_nan].mean(0, keepdim=True)

    # apply transformation
    if transformation == "clr":
        y = torch.from_numpy(clr(y)).float()  # (n_spots, n_isos)
    elif transformation == "ilr":
        y = torch.from_numpy(ilr(y)).float()  # (n_spots, n_isos - 1)
    elif transformation == "alr":
        y = torch.from_numpy(alr(y)).float()  # (n_spots, n_isos - 1)
    elif transformation == "radial":
        y = y / y.norm(dim=1, keepdim=True)  # radial transformation with nans

    # fill back nan to rows with zero counts if needed
    if nan_filling == "none":
        y[is_nan] = torch.nan

    return y


# From scipy v1.13.1
# https://github.com/scipy/scipy/blob/v1.13.1/scipy/stats/_morestats.py#L4737
def false_discovery_control(
    ps: ArrayLike, *, axis: Optional[int] = 0, method: Literal["bh", "by"] = "bh"
) -> np.ndarray:
    """Adjust p-values to control the false discovery rate.

    The false discovery rate (FDR) is the expected proportion of rejected null
    hypotheses that are actually true.
    If the null hypothesis is rejected when the *adjusted* p-value falls below
    a specified level, the false discovery rate is controlled at that level.

    Parameters
    ----------
    ps
        The p-values to adjust. Elements must be real numbers between 0 and 1.
    axis
        The axis along which to perform the adjustment. The adjustment is
        performed independently along each axis-slice. If `axis` is None, `ps`
        is raveled before performing the adjustment.
    method
        The false discovery rate control procedure to apply: ``'bh'`` is for
        Benjamini-Hochberg :cite:`benjamini1995controlling` (Eq. 1), ``'by'`` is for Benjaminini-Yekutieli
        :cite:`benjamini2001control` (Theorem 1.3). The latter is more conservative, but it is
        guaranteed to control the FDR even when the p-values are not from
        independent tests.

    Returns
    -------
    ps_adjusted : numpy.ndarray
        The adjusted p-values. If the null hypothesis is rejected where these
        fall below a specified level, the false discovery rate is controlled
        at that level.

    Notes
    -----
    From `scipy.stats.false_discovery_control` in SciPy v1.13.1.
    See https://github.com/scipy/scipy/blob/v1.13.1/scipy/stats/_morestats.py#L4737.
    """
    # Input Validation and Special Cases
    ps = np.asarray(ps)

    # Handle NaNs
    if np.isnan(ps).any():
        warnings.warn("NaNs encountered in p-values. These will be ignored.")
        # ignore NaNs in the p-values
        ps_in_range = np.issubdtype(ps.dtype, np.number) and np.all(
            ps[~np.isnan(ps)] == np.clip(ps[~np.isnan(ps)], 0, 1)
        )
    else:
        ps_in_range = np.issubdtype(ps.dtype, np.number) and np.all(
            ps == np.clip(ps, 0, 1)
        )

    if not ps_in_range:
        raise ValueError("`ps` must include only numbers between 0 and 1.")

    methods = {"bh", "by"}
    if method.lower() not in methods:
        raise ValueError(
            f"Unrecognized `method` '{method}'." f"Method must be one of {methods}."
        )
    method = method.lower()

    if axis is None:
        axis = 0
        ps = ps.ravel()

    axis = np.asarray(axis)[()]
    if not np.issubdtype(axis.dtype, np.integer) or axis.size != 1:
        raise ValueError("`axis` must be an integer or `None`")

    if ps.size <= 1 or ps.shape[axis] <= 1:
        return ps[()]

    ps = np.moveaxis(ps, axis, -1)
    m = ps.shape[-1]

    # Main Algorithm
    # Equivalent to the ideas of [1] and [2], except that this adjusts the
    # p-values as described in [3]. The results are similar to those produced
    # by R's p.adjust.

    # "Let [ps] be the ordered observed p-values..."
    order = np.argsort(ps, axis=-1)
    ps = np.take_along_axis(ps, order, axis=-1)  # this copies ps

    # Equation 1 of [1] rearranged to reject when p is less than specified q
    i = np.arange(1, m + 1)
    ps *= m / i

    # Theorem 1.3 of [2]
    if method == "by":
        ps *= np.sum(1 / i)

    # accounts for rejecting all null hypotheses i for i < k, where k is
    # defined in Eq. 1 of either [1] or [2]. See [3]. Starting with the index j
    # of the second to last element, we replace element j with element j+1 if
    # the latter is smaller.
    np.fmin.accumulate(ps[..., ::-1], out=ps[..., ::-1], axis=-1)

    # Restore original order of axes and data
    np.put_along_axis(ps, order, values=ps.copy(), axis=-1)
    ps = np.moveaxis(ps, -1, axis)

    return np.clip(ps, 0, 1)


# Similar to scanpy.read_visium
# https://github.com/scverse/scanpy/blob/main/scanpy/readwrite.py#L356-L512
def load_visium_sp_meta(
    adata: AnnData, path_to_spatial: str | Path, library_id: Optional[str] = None
) -> AnnData:
    """Helper function to load Visium spatial metadata.

    Parameters
    ----------
    adata
        Annotated data matrix to store the spatial metadata.
    path_to_spatial
        Path to the `spatial` folder generated by Space Ranger.
    library_id
        Library ID of the spatial data.

    Returns
    -------
    anndata : anndata.AnnData
        AnnData with spatial metadata.
    """
    if library_id is None:  # default library_id
        library_id = "library_id"

    adata.uns["spatial"] = dict()
    adata.uns["spatial"][library_id] = dict()

    path = Path(path_to_spatial)
    tissue_positions_file = (
        path / "tissue_positions.csv"
        if (path / "tissue_positions.csv").exists()
        else path / "tissue_positions_list.csv"
    )

    files = dict(
        tissue_positions_file=tissue_positions_file,
        scalefactors_json_file=path / "scalefactors_json.json",
        hires_image=path / "tissue_hires_image.png",
        lowres_image=path / "tissue_lowres_image.png",
    )

    # load images
    adata.uns["spatial"][library_id]["images"] = dict()
    for res in ["hires", "lowres"]:
        try:
            adata.uns["spatial"][library_id]["images"][res] = imread(
                str(files[f"{res}_image"])
            )
        except Exception:
            warnings.warn(
                f"Missing '{res}' image in {path_to_spatial}. Will be ignored."
            )
            adata.uns["spatial"][library_id]["images"][res] = None

    # read json scalefactors
    with open(files["scalefactors_json_file"]) as f:
        adata.uns["spatial"][library_id]["scalefactors"] = json.load(f)

    # read coordinates
    positions = pd.read_csv(
        files["tissue_positions_file"],
        header=0,
        index_col=0,
    )
    positions.columns = [
        "in_tissue",
        "array_row",
        "array_col",
        "pxl_col_in_fullres",
        "pxl_row_in_fullres",
    ]

    # add coordinates to spot metadata
    adata.obs = adata.obs.join(positions, how="left")
    adata.obsm["spatial"] = adata.obs[
        ["pxl_row_in_fullres", "pxl_col_in_fullres"]
    ].to_numpy()
    adata.obs.drop(
        columns=["pxl_row_in_fullres", "pxl_col_in_fullres"],
        inplace=True,
    )

    return adata


def _process_design_mtx(
    adata: AnnData,
    design_mtx: Optional[Any],
    covariate_names: Optional[list[str]],
) -> tuple[Optional[np.ndarray], Optional[list[str]]]:
    """Process and resolve design matrix from AnnData.obs columns.

    Handles extraction, categorical encoding, sparse array conversion, and
    dimension validation for design matrices extracted from AnnData object.
    Automatically infers covariate names when not explicitly provided.

    Supports three input modes:

    1. Single column name (str): extracts one ``adata.obs`` column
    2. Multiple column names (list[str]): extracts and concatenates columns
    3. Pre-computed matrix (array/dataframe/tensor): used as-is

    For columns containing categorical data (dtype 'category' or object),
    one-hot encoding is applied. Numerical columns are converted as-is.

    Parameters
    ----------
    adata
        Annotated data matrix.
    design_mtx
        Design matrix specification or pre-computed matrix.

        - str: single column name in ``adata.obs``
        - list[str]: list of column names in ``adata.obs``
        - ndarray/tensor/sparse/DataFrame: pre-computed design matrix

    covariate_names
        Explicit covariate names. If provided, dimensions must match
        the resulting design matrix. If None, names are inferred as follows:

        - **From obs columns (str/list[str])**: column names are used, with
          categorical columns expanded to one-hot encoded names.
        - **From DataFrame**: column names are extracted.
        - **Otherwise**: default names ``factor_1``, ``factor_2``, etc.

    Returns
    -------
    resolved_design : np.ndarray | None
        Design matrix of shape ``(n_spots, n_factors)``, dtype float32.
        Returns ``None`` if ``design_mtx`` is ``None``.
    resolved_covariates : list[str] | None
        Inferred or explicit covariate names matching design matrix columns.
        Resolves as:

        1. User-provided ``covariate_names`` (if valid dimensions)
        2. Extracted from obs columns or DataFrame
        3. Auto-generated ``factor_1``, ``factor_2``, etc.

        Returns ``None`` if ``design_mtx`` is ``None``.

    Raises
    ------
    ValueError
        If columns are missing, covariate_names dimensions don't match,
        or other validation failures occur.
    """
    if design_mtx is None:
        return None, None

    # Extract columns from adata.obs
    if isinstance(design_mtx, str):
        # Single column
        col_names = [design_mtx]
        columns_to_extract = col_names
    elif (
        isinstance(design_mtx, list)
        and len(design_mtx) > 0
        and all(isinstance(x, str) for x in design_mtx)
    ):
        # Multiple columns
        col_names = list(design_mtx)
        columns_to_extract = col_names
    else:
        # Pre-computed matrix: validate and return
        columns_to_extract = None
        col_names = None

    if columns_to_extract is not None:
        # Validate columns exist
        missing = [col for col in columns_to_extract if col not in adata.obs.columns]
        if missing:
            raise ValueError(
                f"Columns {missing} were not found in `adata.obs` for design_mtx."
            )

        # Extract and process columns
        design_list = []
        resolved_cov_names = []

        for col in columns_to_extract:
            col_data = adata.obs[col]
            col_dtype = col_data.dtype

            # Check if categorical (either pandas Categorical or object/string type)
            is_categorical = (
                isinstance(col_dtype, pd.CategoricalDtype)
                or col_dtype == "object"
                or col_dtype.kind == "O"  # Object array
            )

            if is_categorical:
                # One-hot encode categorical column
                # Use pd.get_dummies to create binary indicators
                dummies = pd.get_dummies(col_data, prefix=col, drop_first=False)
                design_list.append(dummies.values.astype(np.float32))
                resolved_cov_names.extend(dummies.columns.tolist())
            else:
                # Numerical column: convert to float
                col_array = col_data.values.astype(np.float32).reshape(-1, 1)
                design_list.append(col_array)
                resolved_cov_names.append(col)

        # Concatenate all columns
        resolved_design = np.concatenate(design_list, axis=1)

    else:
        # Pre-computed matrix: convert to appropriate format
        if isinstance(design_mtx, pd.DataFrame):
            resolved_design = design_mtx.values.astype(np.float32)
            resolved_cov_names = list(design_mtx.columns)
        elif isinstance(design_mtx, torch.Tensor):
            resolved_design = design_mtx.cpu().numpy().astype(np.float32)
            resolved_cov_names = None
        elif scipy.sparse.issparse(design_mtx):
            # Convert sparse matrix to dense
            resolved_design = design_mtx.toarray().astype(np.float32)
            resolved_cov_names = None
        elif isinstance(design_mtx, np.ndarray):
            resolved_design = design_mtx.astype(np.float32)
            resolved_cov_names = None
        else:
            raise TypeError(
                f"Unsupported design_mtx type: {type(design_mtx)}. "
                "Expected array-like, DataFrame, tensor, or sparse matrix."
            )

    # Validate shapes
    n_spots = adata.n_obs
    if resolved_design.shape[0] != n_spots:
        raise ValueError(
            f"Design matrix row count ({resolved_design.shape[0]}) "
            f"must match number of spots ({n_spots})."
        )

    # Validate and set covariate names
    n_factors = resolved_design.shape[1]
    if covariate_names is not None:
        if len(covariate_names) != n_factors:
            raise ValueError(
                f"covariate_names length ({len(covariate_names)}) must match "
                f"design matrix columns ({n_factors})."
            )
        resolved_cov_names = covariate_names
    elif resolved_cov_names is None:
        # Generate default names
        resolved_cov_names = [f"factor_{i+1}" for i in range(n_factors)]

    return resolved_design, resolved_cov_names


def prepare_inputs_from_anndata(
    adata: AnnData,
    layer: str,
    group_iso_by: str,
    spatial_key: str,
    min_counts: int,
    min_bin_pct: float,
    filter_single_iso_genes: bool,
    gene_names: Optional[str],
    design_mtx: Optional[Any],
    covariate_names: Optional[list[str]],
) -> tuple[
    list[torch.Tensor],
    torch.Tensor,
    list[str],
    Optional[Any],
    Optional[list[str]],
]:
    """Extract and filter isoform count tensors from an AnnData object.

    Shared helper used by both :class:`splisosm.hyptest_np.SplisosmNP` and
    :class:`splisosm.hyptest_glmm.SplisosmGLMM` to prepare legacy-compatible
    tensors from an AnnData input.  Feature filtering, sparse/dense handling,
    coordinate extraction, and design-matrix resolution are all performed here.

    Parameters
    ----------
    adata
        Annotated data matrix.
    layer
        Key in ``adata.layers`` containing raw isoform counts.
    group_iso_by
        Column in ``adata.var`` used to group isoforms by gene.
    spatial_key
        Key in ``adata.obsm`` for spatial coordinates.
    min_counts
        Minimum total isoform count across spots required to retain an isoform.
    min_bin_pct
        Minimum fraction/percentage of spots with non-zero expression for an
        isoform.  Values in ``[0, 1]`` are treated as fractions; values in
        ``(1, 100]`` are treated as percentages.
    filter_single_iso_genes
        Whether to discard genes with fewer than two retained isoforms.
    gene_names
        Column name in ``adata.var`` used as display names for grouped genes.
        If ``None``, the grouped gene IDs are used.
    design_mtx
        Design matrix for differential-usage tests.  Accepts a
        tensor/array/dataframe of shape ``(n_spots, n_factors)``, a single
        obs-column name (str), or a list of obs-column names.
    covariate_names
        Explicit covariate names.  When ``design_mtx`` is given as column
        name(s) and this is ``None``, the column names are used automatically.

    Returns
    -------
    counts_list : list[torch.Tensor]
        Per-gene isoform count tensors, each of shape ``(n_spots, n_isos)``.
        Sparse ``adata.layers[layer]`` input yields sparse COO tensors.
    coordinates : torch.Tensor
        Shape ``(n_spots, 2)`` spatial coordinates, dtype float32.
    resolved_gene_names : list[str]
        Display names for each gene in ``counts_list``.
    resolved_design : np.ndarray or tensor or None
        Resolved design matrix, or the original object if it was already
        array-like; ``None`` when ``design_mtx`` is ``None``.
    resolved_covariates : list[str] or None
        Resolved covariate names, or ``None`` when ``design_mtx`` is ``None``.

    Raises
    ------
    ValueError
        If required fields are missing from ``adata``, no isoforms survive
        filtering, or argument values are out of range.
    """
    if not isinstance(adata, AnnData):
        raise ValueError("`adata` must be an AnnData object.")
    if layer not in adata.layers:
        raise ValueError(f"Layer `{layer}` was not found in `adata.layers`.")
    if group_iso_by not in adata.var.columns:
        raise ValueError(f"`{group_iso_by}` was not found in `adata.var`.")
    if spatial_key not in adata.obsm:
        raise ValueError(f"`{spatial_key}` was not found in `adata.obsm`.")
    if min_counts < 0:
        raise ValueError("`min_counts` must be >= 0.")
    if min_bin_pct < 0 or min_bin_pct > 100:
        raise ValueError("`min_bin_pct` must be between 0 and 1, or between 0 and 100.")
    if isinstance(gene_names, list):
        raise ValueError(
            "In AnnData mode, `gene_names` must be a var-column name (str) or None."
        )

    iso_counts = adata.layers[layer]
    min_bin_frac = float(min_bin_pct)
    if min_bin_frac > 1.0:
        min_bin_frac /= 100.0

    total_counts = np.asarray(iso_counts.sum(axis=0)).ravel()
    if hasattr(iso_counts, "toarray"):
        bin_pct = np.asarray((iso_counts > 0).sum(axis=0)).ravel() / adata.n_obs
    else:
        bin_pct = np.count_nonzero(np.asarray(iso_counts) > 0, axis=0) / adata.n_obs

    var_df = adata.var.copy()
    var_df["__iso_name__"] = adata.var_names.astype(str)
    var_df["__total_counts__"] = total_counts
    var_df["__bin_pct__"] = bin_pct

    var_df = var_df[
        (var_df["__total_counts__"] >= min_counts)
        & (var_df["__bin_pct__"] >= min_bin_frac)
    ]
    if var_df.shape[0] == 0:
        raise ValueError(
            "No isoforms remained after applying `min_counts`/`min_bin_pct` filtering."
        )

    if filter_single_iso_genes:
        n_iso_per_gene = var_df.groupby(group_iso_by, observed=True).size()
        keep_genes = n_iso_per_gene[n_iso_per_gene >= 2].index
        var_df = var_df[var_df[group_iso_by].isin(keep_genes)]
        if var_df.shape[0] == 0:
            raise ValueError("No genes with >=2 isoforms remained after filtering.")

    is_sparse_input = scipy.sparse.issparse(iso_counts)
    if (
        is_sparse_input
        and not scipy.sparse.isspmatrix_csr(iso_counts)
        and not scipy.sparse.isspmatrix_csc(iso_counts)
    ):
        iso_counts = iso_counts.tocsr()

    iso_name_to_col = {name: i for i, name in enumerate(adata.var_names)}

    counts_list: list[torch.Tensor] = []
    grouped_genes: list[str] = []
    for gene_id, group in var_df.groupby(group_iso_by, observed=True, sort=False):
        iso_indices = [iso_name_to_col[name] for name in group["__iso_name__"]]
        if is_sparse_input:
            _counts_coo = iso_counts[:, iso_indices].tocoo()
            _i = torch.LongTensor(np.vstack((_counts_coo.row, _counts_coo.col)))
            _v = torch.FloatTensor(_counts_coo.data)
            _counts = torch.sparse_coo_tensor(_i, _v, torch.Size(_counts_coo.shape))
        else:
            _counts_np = np.asarray(iso_counts[:, iso_indices], dtype=np.float32)
            _counts = torch.from_numpy(_counts_np)
        counts_list.append(_counts)
        grouped_genes.append(gene_id)

    if len(counts_list) == 0:
        raise ValueError("No genes remained after extracting isoform counts.")

    if gene_names is not None:
        if gene_names not in adata.var.columns:
            raise ValueError(
                f"`gene_names` column `{gene_names}` was not found in `adata.var`."
            )
        gene_name_map = var_df.groupby(group_iso_by, observed=True)[gene_names].first()
        resolved_gene_names = [str(gene_name_map.loc[_g]) for _g in grouped_genes]
    else:
        resolved_gene_names = [str(_g) for _g in grouped_genes]

    coordinates = adata.obsm[spatial_key]
    coordinates = torch.as_tensor(np.asarray(coordinates), dtype=torch.float32)
    if coordinates.dim() != 2:
        raise ValueError("Coordinates in `adata.obsm[spatial_key]` must be a 2D array.")

    # Process design matrix (handles column extraction, categorical encoding, etc.)
    resolved_design, resolved_covariates = _process_design_mtx(
        adata, design_mtx, covariate_names
    )

    return (
        counts_list,
        coordinates,
        resolved_gene_names,
        resolved_design,
        resolved_covariates,
    )


def extract_counts_n_ratios(
    adata: AnnData,
    layer: str = "counts",
    group_iso_by: str = "gene_symbol",
    return_sparse: bool = False,
    filter_single_iso_genes: bool = True,
) -> tuple[list[torch.Tensor], list[torch.Tensor], list[str], Optional[np.ndarray]]:
    """Extract per-gene lists of isoform counts and ratios from anndata.

    Parameters
    ----------
    adata
        Annotated data matrix.
    layer
        Layer to extract isoform counts (adata.layers[layer]).
    group_iso_by
        Gene index in adata.var to group isoforms by.
    return_sparse
        Whether to return sparse torch tensors for counts_list.
        If True, `ratios_list` will be empty and `ratio_obs_merged` will be None.
    filter_single_iso_genes
        Whether to filter out genes with only one isoform.
        By default True for compatibility with splisosm models.

    Returns
    -------
    counts_list : list[torch.Tensor]
        Isoform counts per gene, each of shape (n_spots, n_isos).
    ratios_list : list[torch.Tensor]
        Isoform ratios per gene, each of shape (n_spots, n_isos).
    gene_name_list : list[str]
        Gene names.
    ratio_obs_merged : np.ndarray | None
        Observed isoform ratios, shape (n_spots, n_isos_total), or None if `return_sparse` is True.
    """
    # extract isoform counts
    iso_counts = adata.layers[layer]  # (n_spots, n_isos_total)

    # Check if input is sparse
    is_sparse_input = scipy.sparse.issparse(iso_counts)
    if (
        is_sparse_input
        and not scipy.sparse.isspmatrix_csc(iso_counts)
        and not scipy.sparse.isspmatrix_csr(iso_counts)
    ):
        # convert to csr for efficient slicing if it is other sparse format
        iso_counts = iso_counts.tocsr()

    counts_list = []  # isoform counts per gene, each of (n_spots, n_isos)
    ratios_list = []  # isoform ratios per gene, each of (n_spots, n_isos)
    gene_name_list = []  # of length n_genes
    iso_ind_list = []  # of length n_genes

    for _gene, _group in tqdm(
        adata.var.reset_index().groupby(group_iso_by, observed=True)
    ):
        # filter single-isoform genes if needed
        if filter_single_iso_genes and _group.shape[0] < 2:
            continue

        # extract isoform name and index per gene
        gene_name_list.append(_gene)
        iso_indices = _group.index.tolist()
        iso_ind_list.append(iso_indices)

        # extract isoform counts and relative ratio
        if is_sparse_input and return_sparse:
            # since ratios are usually dense, we do not compute ratios for sparse input
            # ratios_list will be empty and ratio_obs_merged will be None

            # slice sparse matrix
            _counts_scipy = iso_counts[:, iso_indices]

            # convert to torch sparse coo tensor
            _counts_coo = _counts_scipy.tocoo()
            values = _counts_coo.data
            indices = np.vstack((_counts_coo.row, _counts_coo.col))

            i = torch.LongTensor(indices)
            v = torch.FloatTensor(values)
            shape = _counts_coo.shape
            _counts = torch.sparse_coo_tensor(i, v, torch.Size(shape))
            counts_list.append(_counts)
        else:
            if is_sparse_input:
                # slice sparse matrix and convert to dense
                _counts_np = iso_counts[:, iso_indices].toarray()
            else:
                _counts_np = iso_counts[:, iso_indices]

            # compute counts and ratios as dense tensors
            _counts = torch.from_numpy(_counts_np).float()  # (n_spots, n_isos)
            _ratios = counts_to_ratios(
                _counts, transformation="none"
            )  # (n_spots, n_isos)

            counts_list.append(_counts)
            ratios_list.append(_ratios)

    if is_sparse_input and return_sparse:
        ratio_obs_merged = None
    else:
        # reshape and store the observed ratio in anndata
        ratio_obs_merged = torch.concat(
            ratios_list, axis=1
        ).numpy()  # (n_spots, n_isos_total)
        ratio_obs_merged = ratio_obs_merged[
            :, np.argsort(np.concatenate(iso_ind_list))
        ]  # (n_spots, n_isos_total)

    return counts_list, ratios_list, gene_name_list, ratio_obs_merged


def extract_gene_level_statistics(
    adata: AnnData, layer: str = "counts", group_iso_by: str = "gene_symbol"
) -> pd.DataFrame:
    """Extract gene-level metadata from isoform-level counts anndata.

    Parameters
    ----------
    adata
        Annotated data matrix.
    layer
        Layer to extract isoform counts (adata.layers[layer]).
    group_iso_by
        Gene index in adata.var to group isoforms by.

    Returns
    -------
    pandas.DataFrame
        Gene-level metadata with columns:

        - ``'n_iso'``: int. Number of isoforms per gene.
        - ``'pct_spot_on'``: float. Percentage of spots with non-zero counts.
        - ``'count_avg'``: float. Average counts per gene.
        - ``'count_std'``: float. Standard deviation of counts per gene.
        - ``'perplexity'``: float. Expression-based effective number of isoforms.
        - ``'major_ratio_avg'``: float. Average ratio of the major isoform.
    """
    # extract isoform counts
    iso_counts = adata.layers[layer]  # (n_spots, n_isos_total)

    # Check if input is sparse
    is_sparse_input = scipy.sparse.issparse(iso_counts)
    if (
        is_sparse_input
        and not scipy.sparse.isspmatrix_csc(iso_counts)
        and not scipy.sparse.isspmatrix_csr(iso_counts)
    ):
        iso_counts = iso_counts.tocsr()

    df_list = []
    # loop through genes
    for _gene, _group in tqdm(
        adata.var.reset_index().groupby(group_iso_by, observed=True)
    ):
        # extract isoform counts and relative ratio
        iso_indices = _group.index.tolist()
        _counts = iso_counts[:, iso_indices]

        if is_sparse_input:
            # Calculate statistics on sparse matrix without densifying
            _sum_per_iso = np.asarray(_counts.sum(0)).flatten()  # (n_isos,)
            _row_sums = np.asarray(_counts.sum(1)).flatten()  # (n_spots,)
            _total_sum = _sum_per_iso.sum()

            pct_spot_on = (_row_sums > 0).mean()
            count_avg = _row_sums.mean()
            count_std = _row_sums.std()
        else:
            if not isinstance(_counts, np.ndarray):
                _counts = np.asarray(_counts)

            _sum_per_iso = _counts.sum(0)
            _row_sums = _counts.sum(1)
            _total_sum = _counts.sum()

            pct_spot_on = (_row_sums > 0).mean()
            count_avg = _row_sums.mean()
            count_std = _row_sums.std()

        # Avoid division by zero
        if _total_sum == 0:
            _ratios_avg = np.zeros_like(_sum_per_iso)
        else:
            _ratios_avg = _sum_per_iso / _total_sum  # (n_isos,)

        # calculate and store gene-level statistics
        # handle zeros in log
        with np.errstate(divide="ignore", invalid="ignore"):
            entropy = -(np.log(_ratios_avg) * _ratios_avg)
            entropy = np.nan_to_num(entropy).sum()

        df_list.append(
            {
                "gene": _gene,
                "n_iso": _group.shape[0],
                "pct_spot_on": pct_spot_on,
                "count_avg": count_avg,
                "count_std": count_std,
                "perplexity": np.exp(entropy),
                "major_ratio_avg": _ratios_avg.max() if _ratios_avg.size > 0 else 0.0,
            }
        )

    df_gene_meta = pd.DataFrame(df_list).set_index("gene")

    return df_gene_meta


def run_sparkx(
    counts_gene: np.ndarray | torch.Tensor,
    coordinates: np.ndarray | torch.Tensor,
) -> dict[str, Any]:
    """Wrapper for running the SPARK-X test for spatial gene expression variability.

    It runs the R-package SPARK :cite:`zhu2021spark` via rpy2.

    Parameters
    ----------
    counts_gene
        Shape (n_spots, n_genes), the observed gene counts.
    coordinates
        Shape (n_spots, 2), the spatial coordinates.

    Returns
    -------
    dict
        Results of the SPARK-X spatial variability test with keys:

        - ``'statistic'``: np.ndarray of shape (n_genes,). Mean SPARK-X statistics.
        - ``'pvalue'``: np.ndarray of shape (n_genes,). Combined p-values.
        - ``'pvalue_adj'``: np.ndarray of shape (n_genes,). Adjusted combined p-values.
        - ``'method'``: str. Method name "spark-x".
    """
    if scipy.sparse.issparse(counts_gene):
        raise ValueError("run_sparkx does not support sparse input for counts_gene.")
    if isinstance(counts_gene, torch.Tensor) and counts_gene.is_sparse:
        raise ValueError("run_sparkx does not support sparse input for counts_gene.")

    # load packages neccessary for running SPARK-X
    import rpy2.robjects as ro
    from rpy2.robjects import r
    from rpy2.robjects import numpy2ri
    from rpy2.robjects.packages import importr

    spark = importr("SPARK")

    # prepare robject inputs
    if isinstance(counts_gene, torch.Tensor):
        counts_gene = counts_gene.numpy()
    if isinstance(coordinates, torch.Tensor):
        coordinates = coordinates.numpy()

    with (ro.default_converter + numpy2ri.converter).context():
        coords_r = ro.conversion.get_conversion().py2rpy(coordinates)  # (n_spots, 2)
        counts_r = ro.conversion.get_conversion().py2rpy(
            counts_gene.T
        )  # (n_genes, n_spots)

    counts_r.colnames = ro.vectors.StrVector(r["rownames"](coords_r))

    # run SPARK-X and extract outputs
    sparkx_res = spark.sparkx(counts_r, coords_r)
    with (ro.default_converter + numpy2ri.converter).context():
        sv_sparkx = {
            "statistic": ro.conversion.get_conversion()
            .rpy2py(sparkx_res.rx2("stats"))
            .mean(1),
            "pvalue": ro.conversion.get_conversion().rpy2py(
                sparkx_res.rx2("res_mtest")
            )["combinedPval"],
            "pvalue_adj": ro.conversion.get_conversion().rpy2py(
                sparkx_res.rx2("res_mtest")
            )["adjustedPval"],
            "method": "spark-x",
        }

    return sv_sparkx


def run_hsic_gc(
    counts_gene: np.ndarray | torch.Tensor,
    coordinates: np.ndarray | torch.Tensor,
    approx_rank: Optional[int] = None,
    **spatial_kernel_kwargs: Any,
) -> dict[str, Any]:
    """Function to compute HSIC-GC statistic for gene-level counts.

    This function is designed to be a plugin replacement for SPARK-X.

    Parameters
    ----------
    counts_gene
        Shape (n_spots, n_genes). Gene counts.
    coordinates
        Shape (n_spots, 2). Spatial coordinates of spots.
    approx_rank
        Approximate rank of the spatial kernel matrix.
    **spatial_kernel_kwargs
        Additional arguments for SpatialCovKernel.

    Returns
    -------
    dict
        Results of the HSIC-GC spatial variability test with keys:

        - ``'statistic'``: np.ndarray of shape (n_genes,). HSIC-GC statistics.
        - ``'pvalue'``: np.ndarray of shape (n_genes,). P-values.
        - ``'pvalue_adj'``: np.ndarray of shape (n_genes,). Adjusted p-values.
        - ``'method'``: str. Method name "hsic-gc".
    """

    from splisosm.kernel import SpatialCovKernel

    n_spots = counts_gene.shape[0]
    n_genes = counts_gene.shape[1]

    # determine the maximum rank for spatial kernel computation
    if n_spots > 5000:
        # 10x Visium has 4992 spots per slide. For larger datasets (i.e. Slideseq-V2),
        # it is recommended to use low-rank approximation
        max_rank = np.ceil(np.sqrt(n_spots) * 4).astype(int)
        approx_rank = (
            min(approx_rank, max_rank) if approx_rank is not None else max_rank
        )
    else:
        if approx_rank is not None:
            approx_rank = approx_rank if approx_rank < n_spots else None

    # set default spatial kernel kwargs
    default_spatial_kernel_kwargs = {
        "k_neighbors": 4,
        "model": "icar",
        "rho": 0.99,
        "standardize_cov": True,
        "centering": True,
        "approx_rank": approx_rank,
    }
    if spatial_kernel_kwargs is not None:
        if "centering" in spatial_kernel_kwargs:
            warnings.warn(
                "The 'centering' argument in spatial_kernel_kwargs will be ignored. It is always set to True for HSIC-GC."
            )
            spatial_kernel_kwargs.pop("centering")
        default_spatial_kernel_kwargs.update(spatial_kernel_kwargs)

    # compute the spatial kernel
    # Ensure coordinates is numpy for SpatialCovKernel/smoother
    if isinstance(coordinates, torch.Tensor):
        coordinates = coordinates.detach().cpu().numpy()

    K_sp = SpatialCovKernel(coordinates, **default_spatial_kernel_kwargs)

    # get the eigenvalues
    lambda_sp = K_sp.eigenvalues()  # (rank,)
    lambda_sp = lambda_sp[lambda_sp > 1e-5]  # filter small eigenvalues

    # compute the HSIC-GC statistic per-gene
    is_scipy_sparse = scipy.sparse.issparse(counts_gene)
    is_torch_sparse = isinstance(counts_gene, torch.Tensor) and counts_gene.is_sparse

    if is_scipy_sparse:
        n_spots, n_genes = counts_gene.shape
        # Ensure efficient column slicing
        if not scipy.sparse.isspmatrix_csc(
            counts_gene
        ) and not scipy.sparse.isspmatrix_csr(counts_gene):
            counts_gene = counts_gene.tocsc()
    elif is_torch_sparse:
        n_spots, n_genes = counts_gene.shape
        if counts_gene.dtype != torch.float32 and counts_gene.dtype != torch.float64:
            counts_gene = counts_gene.float()
        if not counts_gene.is_coalesced():
            counts_gene = counts_gene.coalesce()
        # Pre-allocate selection vector for column extraction
        selection_vec = torch.zeros(
            n_genes, 1, device=counts_gene.device, dtype=counts_gene.dtype
        )
    else:
        if not isinstance(counts_gene, torch.Tensor):
            counts_gene = torch.from_numpy(counts_gene).float()  # (n_spots, n_genes)
        n_spots, n_genes = counts_gene.shape
        y_dense = counts_gene - counts_gene.mean(
            0, keepdim=True
        )  # center the counts, (n_spots, n_genes)

    hsic_list, pvals_list = [], []
    for i in tqdm(range(n_genes)):
        if is_scipy_sparse:
            col = counts_gene[:, i].toarray()  # dense (n_spots, 1)
            counts = torch.from_numpy(col).float()
            counts = counts - counts.mean()
        elif is_torch_sparse:
            # Extract column i using Sparse x Dense matrix multiplication
            # (N, G) @ (G, 1) -> (N, 1)
            # This is efficient and avoids densifying the full matrix
            selection_vec.zero_()
            selection_vec[i, 0] = 1.0

            col = torch.sparse.mm(counts_gene, selection_vec)
            counts = col - col.mean()
        else:
            counts = y_dense[:, i : i + 1]  # (n_spots, 1)

        lambda_y = counts.T @ counts  # (1, 1)
        lambda_spy = lambda_sp * lambda_y  # (rank, 1)
        hsic_scaled = torch.trace(K_sp.xtKx(counts))  # scalar
        pval = liu_sf((hsic_scaled * n_spots).numpy(), lambda_spy.numpy())

        hsic_list.append(hsic_scaled / (n_spots - 1) ** 2)  # HSIC statistic
        pvals_list.append(pval)

    sv_test_results = {
        "statistic": torch.tensor(hsic_list).numpy(),
        "pvalue": torch.tensor(pvals_list).numpy(),
        "method": "hsic-gc",
    }

    # calculate adjusted p-values
    sv_test_results["pvalue_adj"] = false_discovery_control(sv_test_results["pvalue"])

    return sv_test_results
