"""
Connectivity matrix building for homogeneous spot identification
Implements the spatial → anchor → homogeneous neighbor finding algorithm
"""

import logging
from functools import partial

import anndata as ad
import jax
import jax.numpy as jnp
import numpy as np
import scanpy as sc
from jax import jit
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TaskProgressColumn,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
    track,
)
from scipy.spatial import cKDTree

from gsMap.config import DatasetType, LatentToGeneConfig

logger = logging.getLogger(__name__)

# Configure JAX
jax.config.update("jax_enable_x64", False)  # Use float32 for speed


def find_spatial_neighbors_with_slices(
    coords: np.ndarray,
    slice_ids: np.ndarray | None = None,
    query_cell_mask: np.ndarray | None = None,
    high_quality_cell_mask: np.ndarray = None,
    k_central: int = 101,
    k_adjacent: int = 50,
    n_adjacent_slices: int = 1
) -> tuple[np.ndarray, dict[int, np.ndarray]]:
    """
    Find spatial neighbors with slice-aware search for 3D data.

    For 2D data (slice_ids is None), performs standard KNN.
    For 3D data, implements slice-aware neighbor search:
    - Finds k_central neighbors on the same slice
    - Finds k_adjacent neighbors on each of n_adjacent_slices above and below

    KDTrees are built using only high quality cells, and neighbors are searched
    within high quality cells only.

    Args:
        coords: Spatial coordinates (n_cells, 2) - only x,y coordinates
        slice_ids: Slice/z-coordinate indices (n_cells,) - sequential integers
        query_cell_mask: Boolean mask for cells to find neighbors for
        high_quality_cell_mask: Boolean mask for high quality cells (defines the neighbor search pool)
        k_central: Number of neighbors to find on the central slice
        k_adjacent: Number of neighbors to find on each adjacent slice
        n_adjacent_slices: Number of slices to search above and below

    Returns:
        Tuple of:
        - spatial_neighbors: Array of neighbor indices (n_query_cells, k_central + 2*n_adjacent_slices*k_adjacent)
        - neighbor_pool_per_slice: Dict mapping slice_id to local indices of neighbor pool cells on that slice
    """
    n_cells = len(coords)
    if query_cell_mask is None:
        query_cell_mask = np.ones(n_cells, dtype=bool)

    if high_quality_cell_mask is None:
        high_quality_cell_mask = np.ones(n_cells, dtype=bool)

    logger.debug(f"Finding neighbors: {high_quality_cell_mask.sum()}/{n_cells} cells are high quality")

    query_cell_indices = np.where(query_cell_mask)[0]
    n_query_cells = len(query_cell_indices)

    # Neighbor pool: high quality cells that will be used to build KDTrees
    # (intersection of query_cell_mask and high_quality_cell_mask)
    neighbor_pool_mask = query_cell_mask & high_quality_cell_mask
    neighbor_pool_indices = np.where(neighbor_pool_mask)[0]

    # If no slice_ids provided, perform standard 2D KNN
    if slice_ids is None:
        logger.info(f"No slice IDs provided, performing standard 2D KNN with k={k_central}")
        # Build tree with neighbor pool cells only
        kdtree = cKDTree(coords[neighbor_pool_mask])
        # Query for all query cells (including non-HQ ones)
        _, neighbor_local_indices = kdtree.query(
            coords[query_cell_mask],
            k=min(k_central, len(neighbor_pool_indices)),
            workers=-1  # Use all available cores
        )
        # Convert local indices to global indices
        spatial_neighbors = neighbor_pool_indices[neighbor_local_indices]
        return spatial_neighbors, {}

    # Slice-aware neighbor search with fixed-size arrays
    logger.info(f"Performing slice-aware neighbor search: k_central={k_central}, "
                f"k_adjacent={k_adjacent}, n_adjacent_slices={n_adjacent_slices}")

    query_cell_slice_ids = slice_ids[query_cell_mask]
    query_cell_coords = coords[query_cell_mask]

    # Pre-allocate output with fixed size, initialized with -1 (invalid)
    total_neighbors_per_cell = k_central + 2 * n_adjacent_slices * k_adjacent
    # Always use int32 for spatial neighbors since they contain global cell indices
    # which can easily exceed int16 range in large spatial datasets
    spatial_neighbors = np.full((n_query_cells, total_neighbors_per_cell), -1, dtype=np.int32)

    # Get unique slices and create mapping for all query cells
    unique_slice_ids = np.unique(query_cell_slice_ids)
    slice_to_query_cell_indices = {s: np.where(query_cell_slice_ids == s)[0] for s in unique_slice_ids}

    # Create mapping for neighbor pool cells per slice (for building KDTrees)
    neighbor_pool_per_slice = {}

    # Get slice IDs and coordinates for neighbor pool cells
    neighbor_pool_slice_ids = slice_ids[neighbor_pool_mask]
    neighbor_pool_coords = coords[neighbor_pool_mask]

    for slice_id in unique_slice_ids:
        # Find neighbor pool cells on this slice (in the neighbor_pool array)
        slice_neighbor_pool_mask = neighbor_pool_slice_ids == slice_id
        slice_neighbor_pool_local_indices = np.where(slice_neighbor_pool_mask)[0]
        neighbor_pool_per_slice[slice_id] = slice_neighbor_pool_local_indices
        logger.debug(f"Slice {slice_id}: {len(slice_neighbor_pool_local_indices)} high quality cells")

    # Pre-compute KDTree for each slice using neighbor pool cells only
    # Note: KDTree can handle cases where k > number of points in tree
    # It will return valid neighbors first, then fill remaining slots with invalid indices
    slice_kdtrees = {}
    for slice_id in unique_slice_ids:
        # Use neighbor pool cells to build tree
        slice_neighbor_pool_local = neighbor_pool_per_slice.get(slice_id, np.array([]))
        if len(slice_neighbor_pool_local) > 0:  # Build tree as long as there's at least 1 hq cell
            # Get coordinates of neighbor pool cells on this slice
            slice_neighbor_pool_coords = neighbor_pool_coords[slice_neighbor_pool_local]
            kdtree = cKDTree(slice_neighbor_pool_coords)
            # Store tree with global indices of neighbor pool cells
            slice_kdtrees[slice_id] = (kdtree, neighbor_pool_indices[slice_neighbor_pool_local])

            if len(slice_neighbor_pool_local) < max(k_central, k_adjacent):
                logger.warning(f"Slice {slice_id} has only {len(slice_neighbor_pool_local)} high quality cells, "
                             f"which is less than required k={max(k_central, k_adjacent)}. "
                             # f"Some neighbors will be invalid (-1)."
                               )

    # Build a mapping of slice_id -> list of adjacent slice_ids to search
    # This ensures all slices search the same total number of adjacent slices
    slice_adjacent_mapping = {}
    min_slice_id = min(unique_slice_ids)
    max_slice_id = max(unique_slice_ids)

    for slice_id in unique_slice_ids:
        adjacent_slices = []

        # Collect slices in both directions up to n_adjacent_slices
        for offset in range(1, n_adjacent_slices + 1):
            neg_slice = slice_id - offset
            pos_slice = slice_id + offset
            if neg_slice in slice_kdtrees:
                adjacent_slices.append((neg_slice, offset, -1))  # (slice_id, offset, direction)
            if pos_slice in slice_kdtrees:
                adjacent_slices.append((pos_slice, offset, 1))

        # If we don't have enough adjacent slices (2*n_adjacent_slices total), compensate
        target_count = 2 * n_adjacent_slices
        if len(adjacent_slices) < target_count:
            # Search deeper in available directions
            extra_offset = n_adjacent_slices + 1
            while len(adjacent_slices) < target_count and extra_offset <= max(max_slice_id - min_slice_id, 10):
                neg_slice = slice_id - extra_offset
                pos_slice = slice_id + extra_offset

                if neg_slice >= min_slice_id and neg_slice in slice_kdtrees:
                    if not any(s[0] == neg_slice for s in adjacent_slices):
                        adjacent_slices.append((neg_slice, extra_offset, -1))
                        if len(adjacent_slices) >= target_count:
                            break

                if pos_slice <= max_slice_id and pos_slice in slice_kdtrees:
                    if not any(s[0] == pos_slice for s in adjacent_slices):
                        adjacent_slices.append((pos_slice, extra_offset, 1))
                        if len(adjacent_slices) >= target_count:
                            break

                extra_offset += 1

        # Sort adjacent slices: first by offset (closer slices first), then by direction
        # This maintains consistency: [offset=1,dir=-1], [offset=1,dir=1], [offset=2,dir=-1], [offset=2,dir=1], ...
        adjacent_slices.sort(key=lambda x: (x[1], -x[2]))  # Sort by offset, then direction (-1 before 1)

        slice_adjacent_mapping[slice_id] = adjacent_slices

        if len(adjacent_slices) < target_count:
            logger.warning(f"Slice {slice_id} only has {len(adjacent_slices)} adjacent slices available "
                         f"(target: {target_count}). Some neighbor slots will remain empty.")

    # Log the adjacent slice mapping for verification
    for slice_id in sorted(slice_adjacent_mapping.keys()):
        adj_slice_ids = [s[0] for s in slice_adjacent_mapping[slice_id]]
        logger.debug(f"Slice {slice_id} will search in adjacent slices: {adj_slice_ids}")

    # Batch process all query cells (including edge cases with few neighbor pool cells)
    for slice_id in unique_slice_ids:
        if slice_id in slice_kdtrees:
            # Get all query cells on this slice
            query_cells_on_slice = slice_to_query_cell_indices[slice_id]
            query_coords = query_cell_coords[query_cells_on_slice]

            # Central slice neighbors (fixed k_central)
            central_kdtree, central_neighbor_pool_global_indices = slice_kdtrees[slice_id]
            _, central_neighbor_local_indices = central_kdtree.query(query_coords, k=k_central, workers=-1)

            # Handle invalid indices returned by KDTree when k > number of points in tree
            # KDTree returns index >= tree_size for invalid slots
            n_central_pool_cells = len(central_neighbor_pool_global_indices)
            # Initialize with -1 (invalid)
            central_neighbor_global_indices = np.full_like(central_neighbor_local_indices, -1, dtype=np.int32)
            # Create mask for valid indices (< tree size)
            central_valid_mask = central_neighbor_local_indices < n_central_pool_cells
            # Map valid local indices to global indices
            central_neighbor_global_indices[central_valid_mask] = central_neighbor_pool_global_indices[
                central_neighbor_local_indices[central_valid_mask]
            ]
            spatial_neighbors[query_cells_on_slice, :k_central] = central_neighbor_global_indices

            # Adjacent slices - use the pre-computed mapping
            # Get the list of adjacent slices to search for this slice_id
            adjacent_slices_to_search = slice_adjacent_mapping.get(slice_id, [])

            neighbor_column_offset = k_central

            # We need to fill exactly 2*n_adjacent_slices slots in the output array
            # The array structure is: [offset=1,dir=-1], [offset=1,dir=+1], [offset=2,dir=-1], [offset=2,dir=+1], ...
            for slot_idx in range(2 * n_adjacent_slices):
                if slot_idx < len(adjacent_slices_to_search):
                    # We have a slice to search for this slot
                    adjacent_slice_id, _, _ = adjacent_slices_to_search[slot_idx]

                    adjacent_kdtree, adjacent_neighbor_pool_global_indices = slice_kdtrees[adjacent_slice_id]
                    _, adjacent_neighbor_local_indices = adjacent_kdtree.query(query_coords, k=k_adjacent, workers=-1)

                    # Handle invalid indices for adjacent slices
                    n_adjacent_pool_cells = len(adjacent_neighbor_pool_global_indices)
                    adjacent_neighbor_global_indices = np.full_like(adjacent_neighbor_local_indices, -1, dtype=np.int32)
                    adjacent_valid_mask = adjacent_neighbor_local_indices < n_adjacent_pool_cells
                    adjacent_neighbor_global_indices[adjacent_valid_mask] = adjacent_neighbor_pool_global_indices[
                        adjacent_neighbor_local_indices[adjacent_valid_mask]
                    ]
                    spatial_neighbors[query_cells_on_slice, neighbor_column_offset:neighbor_column_offset+k_adjacent] = adjacent_neighbor_global_indices
                # else: slot remains as -1 (already initialized) if no slice available

                neighbor_column_offset += k_adjacent

    # Note: Slices with 0 neighbor pool cells will not be in slice_kdtrees
    # Query cells on such slices will have all neighbors remain as -1 (already initialized)

    return spatial_neighbors, neighbor_pool_per_slice


@partial(jit, static_argnums=(5, 6))
def _find_homogeneous_spot_dual_embedding_batch_jit(
        emb_niche_batch_norm: jnp.ndarray,  # (batch_size, d1)
        emb_indv_batch_norm: jnp.ndarray,  # (batch_size, d2)
        spatial_neighbors: jnp.ndarray,  # (batch_size, k1)
        all_emb_niche_norm: jnp.ndarray,  # (n_all, d1)
        all_emb_indv_norm: jnp.ndarray,  # (n_all, d2)
        homogeneous_neighbors: int,
        cell_embedding_similarity_threshold: float = 0.0,
        spatial_domain_similarity_threshold: float = 0.5
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """
    Finds homogeneous neighbors using a soft priority (Bonus) system.

    Returns:
        homogeneous_neighbors_array: Indices of selected neighbors (batch, homogeneous_neighbors)
        selected_cell_sims: Cell similarity scores for selected neighbors (batch, homogeneous_neighbors)
        selected_niche_sims: Niche/Spatial similarity scores for selected neighbors (batch, homogeneous_neighbors)
    """
    batch_size = emb_indv_batch_norm.shape[0]

    # Step 1: Extract spatial neighbors' embeddings
    # Handle padding: map -1 to 0 temporarily for extraction
    safe_neighbors = jnp.where(spatial_neighbors >= 0, spatial_neighbors, 0)
    spatial_emb_indv_norm = all_emb_indv_norm[safe_neighbors]
    spatial_emb_niche_norm = all_emb_niche_norm[safe_neighbors]

    # Step 2: Compute Similarities
    # 2a. Cell Similarity
    raw_cell_sims = jnp.einsum('bd,bkd->bk', emb_indv_batch_norm, spatial_emb_indv_norm)
    # Apply cell threshold: values below threshold become 0.0
    cell_sims = jnp.where(raw_cell_sims >= cell_embedding_similarity_threshold, raw_cell_sims, 0.0)

    # 2b. Niche (Spatial) Similarity
    niche_sims = jnp.einsum('bd,bkd->bk', emb_niche_batch_norm, spatial_emb_niche_norm)

    # Step 3: Compute Ranking Score with Bonus
    # Logic:
    #   If niche_sim >= threshold: Score = cell_sim + 2.0
    #   Else:                      Score = cell_sim
    # This prioritizes niche matches first, then falls back to highest cell sim.
    BONUS = 2.0
    pass_spatial_mask = niche_sims >= spatial_domain_similarity_threshold
    ranking_score = cell_sims + jnp.where(pass_spatial_mask, BONUS, 0.0)

    # Step 4: Mask Invalid Neighbors in Ranking
    # Force padding indices (original -1s) to the bottom of the list
    valid_neighbor_mask = spatial_neighbors >= 0
    ranking_score = jnp.where(valid_neighbor_mask, ranking_score, -100.0)

    # Step 5: Sort and Select
    # Get indices of top scores
    top_homo_idx = jnp.argsort(-ranking_score, axis=1)[:, :homogeneous_neighbors]
    batch_idx = jnp.arange(batch_size)[:, None]

    # Gather the results
    homogeneous_neighbors_array = spatial_neighbors[batch_idx, top_homo_idx]
    selected_cell_sims = cell_sims[batch_idx, top_homo_idx]
    selected_niche_sims = niche_sims[batch_idx, top_homo_idx]

    # Step 6: Final Masking
    # If the selected neighbor is invalid (which happens if k1 < homogeneous_neighbors
    # or all valid neighbors were exhausted), ensure returned similarities are 0.0.
    final_valid_mask = homogeneous_neighbors_array >= 0

    selected_cell_sims = jnp.where(final_valid_mask, selected_cell_sims, 0.0)
    selected_niche_sims = jnp.where(final_valid_mask, selected_niche_sims, 0.0)

    return homogeneous_neighbors_array, selected_cell_sims, selected_niche_sims
def _find_homogeneous_3d_memory_efficient(
    emb_niche_masked_jax: jnp.ndarray,
    emb_indv_masked_jax: jnp.ndarray,
    spatial_neighbors: np.ndarray,
    all_emb_niche_norm_jax: jnp.ndarray,
    all_emb_indv_norm_jax: jnp.ndarray,
    num_homogeneous_per_slice: int,
    k_central: int,
    k_adjacent: int,
    n_adjacent_slices: int,
    cell_embedding_similarity_threshold: float,
    spatial_domain_similarity_threshold: float,
    find_homogeneous_batch_size: int
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Memory-efficient version of 3D homogeneous neighbor finding.
    Processes slices separately to avoid large memory allocations.

    Args:
        emb_niche_masked_jax: Niche embeddings for masked cells (JAX array, float16, always exists, may be dummy ones)
        emb_indv_masked_jax: Individual embeddings for masked cells (JAX array, float16)
        spatial_neighbors: Spatial neighbors array with structure [central | adj1 | adj2 | ...] (numpy array)
        all_emb_niche_norm_jax: All normalized niche embeddings (JAX array, float16, always exists, may be dummy ones)
        all_emb_indv_norm_jax: All normalized individual embeddings (JAX array, float16)
        num_homogeneous_per_slice: Number of neighbors to select per slice
        k_central: Number of neighbors in central slice
        k_adjacent: Number of neighbors per adjacent slice
        n_adjacent_slices: Number of adjacent slices above and below
        cell_embedding_similarity_threshold: Minimum similarity threshold for cell embedding
        spatial_domain_similarity_threshold: Minimum similarity threshold for spatial domain embedding
        find_homogeneous_batch_size: Batch size for processing

    Returns:
        homogeneous_neighbors: Selected neighbors
        homogeneous_cell_sims: Cell similarity scores for selected neighbors
        homogeneous_niche_sims: Niche/Spatial similarity scores for selected neighbors
    """
    n_masked = emb_indv_masked_jax.shape[0]
    1 + 2 * n_adjacent_slices

    homogeneous_neighbors_all_slices = []
    homogeneous_cell_sims_all_slices = []
    homogeneous_niche_sims_all_slices = []

    # Process all slices (central + adjacent) in a single loop
    total_slices = 1 + 2 * n_adjacent_slices

    # Create overall slice progress tracking
    with Progress(
        SpinnerColumn(),
        TextColumn("[bold blue]{task.description}"),
        BarColumn(),
        MofNCompleteColumn(),
        TaskProgressColumn(),
        TimeRemainingColumn(),
        TimeElapsedColumn(),
        refresh_per_second=1
    ) as slice_progress:
        # Overall slice progress task
        slice_task = slice_progress.add_task(
            "Finding homogeneous neighbors (3D cross-slice)...",
            total=total_slices
        )

        for slice_num in range(total_slices):
            # Determine slice name and parameters
            if slice_num == 0:
                slice_name = "central slice"
                slice_start = 0
                slice_end = k_central
            else:
                # Adjacent slices
                adj_idx = slice_num - 1
                if adj_idx < n_adjacent_slices:
                    slice_name = f"adjacent slice -{n_adjacent_slices - adj_idx}"
                else:
                    slice_name = f"adjacent slice +{adj_idx - n_adjacent_slices + 1}"

                slice_start = k_central + adj_idx * k_adjacent
                slice_end = slice_start + k_adjacent

            # Convert slice neighbors to JAX array once per slice
            spatial_neighbors_slice = jnp.asarray(spatial_neighbors[:, slice_start:slice_end])

            homogeneous_neighbors_slice_list = []
            homogeneous_cell_sims_slice_list = []
            homogeneous_niche_sims_slice_list = []

            # Process batches for this slice using simple transient track
            for batch_start in track(range(0, n_masked, find_homogeneous_batch_size),
                                    description=f"Finding homogeneous neighbors ({slice_name})",
                                    transient=True):
                batch_end = min(batch_start + find_homogeneous_batch_size, n_masked)
                batch_indices = slice(batch_start, batch_end)

                # Get batch data (emb_niche always exists now, may be dummy ones)
                emb_niche_batch_norm = emb_niche_masked_jax[batch_indices]
                emb_indv_batch_norm = emb_indv_masked_jax[batch_indices]

                # Extract batch of neighbors for this slice
                spatial_neighbors_slice_batch = spatial_neighbors_slice[batch_indices, :]

                # Process with 2D function
                homo_neighbors_batch, cell_sims_batch, niche_sims_batch = _find_homogeneous_spot_dual_embedding_batch_jit(
                    emb_niche_batch_norm=emb_niche_batch_norm,
                    emb_indv_batch_norm=emb_indv_batch_norm,
                    spatial_neighbors=spatial_neighbors_slice_batch,
                    all_emb_niche_norm=all_emb_niche_norm_jax,
                    all_emb_indv_norm=all_emb_indv_norm_jax,
                    homogeneous_neighbors=num_homogeneous_per_slice,
                    cell_embedding_similarity_threshold=cell_embedding_similarity_threshold,
                    spatial_domain_similarity_threshold=spatial_domain_similarity_threshold
                )

                homogeneous_neighbors_slice_list.append(np.array(homo_neighbors_batch))
                homogeneous_cell_sims_slice_list.append(np.array(cell_sims_batch))
                homogeneous_niche_sims_slice_list.append(np.array(niche_sims_batch))

            # Concatenate this slice's results
            homogeneous_neighbors_slice = np.vstack(homogeneous_neighbors_slice_list)
            homogeneous_cell_sims_slice = np.vstack(homogeneous_cell_sims_slice_list)
            homogeneous_niche_sims_slice = np.vstack(homogeneous_niche_sims_slice_list)
            homogeneous_neighbors_all_slices.append(homogeneous_neighbors_slice)
            homogeneous_cell_sims_all_slices.append(homogeneous_cell_sims_slice)
            homogeneous_niche_sims_all_slices.append(homogeneous_niche_sims_slice)

            # Update slice progress
            slice_progress.update(slice_task, advance=1)

    # Concatenate all slices along axis 1
    homogeneous_neighbors = np.concatenate(homogeneous_neighbors_all_slices, axis=1)
    homogeneous_cell_sims = np.concatenate(homogeneous_cell_sims_all_slices, axis=1)
    homogeneous_niche_sims = np.concatenate(homogeneous_niche_sims_all_slices, axis=1)

    return homogeneous_neighbors, homogeneous_cell_sims, homogeneous_niche_sims


def build_scrna_connectivity(
    emb_cell: np.ndarray,
    cell_mask: np.ndarray | None = None,
    n_neighbors: int = 21,
    metric: str = 'euclidean'
) -> tuple[np.ndarray, np.ndarray, None]:
    """
    Build connectivity for scRNA-seq data using KNN on cell embeddings.

    Args:
        emb_cell: Cell embeddings (n_cells, d)
        cell_mask: Boolean mask for cells to process
        n_neighbors: Number of nearest neighbors
        metric: Distance metric for KNN

    Returns:
        neighbor_indices: (n_masked, n_neighbors) array of neighbor indices
        neighbor_weights: (n_masked, n_neighbors) array of weights from KNN graph
        niche_sims: None (scRNA-seq has no niche embedding)
    """
    n_cells = len(emb_cell)
    if cell_mask is None:
        cell_mask = np.ones(n_cells, dtype=bool)

    cell_indices = np.where(cell_mask)[0]
    n_masked = len(cell_indices)

    logger.info(f"Building scRNA-seq connectivity using KNN with k={n_neighbors}")

    # Create temporary AnnData for using scanpy's neighbors function
    adata_temp = ad.AnnData(X=emb_cell[cell_mask])
    adata_temp.obsm['X_emb'] = emb_cell[cell_mask]

    # Compute neighbors using scanpy
    sc.pp.neighbors(
        adata_temp,
        n_neighbors=n_neighbors,
        use_rep='X_emb',
        metric=metric,
        method='umap'
    )

    # Extract connectivity matrix
    connectivities = adata_temp.obsp['connectivities'].tocsr()

    # Convert to dense format for consistency with spatial methods
    # Check the actual max cell index value to determine dtype
    max_cell_idx = cell_indices.max() if len(cell_indices) > 0 else 0
    idx_dtype = np.int16 if max_cell_idx < 32768 else np.int32
    neighbor_indices = np.zeros((n_masked, n_neighbors), dtype=idx_dtype)
    neighbor_weights = np.zeros((n_masked, n_neighbors), dtype=np.float16)

    for i in range(n_masked):
        row = connectivities.getrow(i)
        neighbors = row.indices
        weights = row.data

        # Sort by weight (descending)
        sorted_idx = np.argsort(-weights)[:n_neighbors]

        if len(sorted_idx) < n_neighbors:
            # Pad with self-index if needed
            n_found = len(sorted_idx)
            neighbor_indices[i, :n_found] = cell_indices[neighbors[sorted_idx]]
            neighbor_indices[i, n_found:] = cell_indices[i]
            neighbor_weights[i, :n_found] = weights[sorted_idx]
            neighbor_weights[i, n_found:] = 0.0
        else:
            neighbor_indices[i] = cell_indices[neighbors[sorted_idx]]
            neighbor_weights[i] = weights[sorted_idx]

    # Normalize weights to sum to 1
    weight_sums = neighbor_weights.sum(axis=1, keepdims=True)
    weight_sums = np.where(weight_sums > 0, weight_sums, 1.0)
    neighbor_weights = neighbor_weights / weight_sums

    logger.info(f"scRNA-seq connectivity built: {n_masked} cells × {n_neighbors} neighbors")

    return neighbor_indices, neighbor_weights, None


class ConnectivityMatrixBuilder:
    """Build connectivity matrix using JAX-accelerated computation with GPU memory optimization"""

    def __init__(self, config: LatentToGeneConfig):
        """
        Initialize with configuration

        Args:
            config: LatentToGeneConfig object
        """
        self.config = config
        # Use configured batch size for GPU processing
        self.find_homogeneous_batch_size = config.find_homogeneous_batch_size
        self.dataset_type = config.dataset_type

    def build_connectivity_matrix(
        self,
        coords: np.ndarray | None = None,
        emb_niche: np.ndarray = None,
        emb_indv: np.ndarray | None = None,
        cell_mask: np.ndarray | None = None,
        high_quality_mask: np.ndarray = None,
        slice_ids: np.ndarray | None = None,
        k_central: int | None = None,
        k_adjacent: int | None = None,
        n_adjacent_slices: int | None = None
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
        """
        Build connectivity matrix for a group of cells based on dataset type.

        For scRNA-seq: Uses KNN on cell embeddings (emb_indv)
        For spatial2D: Uses spatial anchors and homogeneous neighbors
        For spatial3D: Uses slice-aware spatial anchors and homogeneous neighbors

        Args:
            coords: Spatial coordinates (n_cells, 2) - required for spatial datasets
            emb_niche: Niche embeddings (n_cells, d1) - always provided (may be dummy ones for scRNA-seq)
            emb_indv: Cell identity embeddings (n_cells, d2) - required for all datasets
            cell_mask: Boolean mask for cells to process
            high_quality_mask: Boolean mask for high quality cells (used for neighbor search in spatial data)
            slice_ids: Optional slice/z-coordinate indices (n_cells,) for spatial3D
            k_central: Number of neighbors on central slice (defaults to config settings)
            k_adjacent: Number of neighbors on adjacent slices for spatial3D
            n_adjacent_slices: Number of slices to search above/below for spatial3D

        Returns:
            Tuple of (neighbor_indices, cell_sims, niche_sims) arrays.
            niche_sims is None for scRNA-seq datasets.
        """
        # Check dataset type and call appropriate method
        if self.dataset_type == DatasetType.SCRNA_SEQ:
            logger.info("Building connectivity for scRNA-seq dataset")
            if emb_indv is None:
                raise ValueError("emb_indv (cell embeddings) required for scRNA-seq dataset")

            return build_scrna_connectivity(
                emb_cell=emb_indv,
                cell_mask=cell_mask,
                n_neighbors=self.config.homogeneous_neighbors,
                metric='euclidean'
            )

        elif self.dataset_type in ['spatial2D', 'spatial3D']:
            logger.info(f"Building connectivity for {self.dataset_type} dataset")

            # Validate required inputs for spatial datasets
            if coords is None or emb_indv is None:
                raise ValueError("coords and emb_indv required for spatial datasets")

            # Use config defaults if not provided
            if k_central is None:
                k_central = self.config.spatial_neighbors
            if k_adjacent is None:
                k_adjacent = self.config.adjacent_slice_spatial_neighbors
            if n_adjacent_slices is None:
                if self.dataset_type == DatasetType.SPATIAL_2D:
                    n_adjacent_slices = 0
                else:  # spatial3D
                    n_adjacent_slices = self.config.n_adjacent_slices

            # For spatial2D, ensure no cross-slice search (but can still have slice_ids)
            if self.dataset_type == DatasetType.SPATIAL_2D:
                n_adjacent_slices = 0  # No cross-slice search, but keep slice_ids if provided

            return self._build_spatial_connectivity(
                coords=coords,
                emb_niche=emb_niche,
                emb_indv=emb_indv,
                cell_mask=cell_mask,
                high_quality_mask=high_quality_mask,
                slice_ids=slice_ids,
                k_central=k_central,
                k_adjacent=k_adjacent,
                n_adjacent_slices=n_adjacent_slices
            )

        else:
            raise ValueError(f"Unknown dataset type: {self.dataset_type}")

    def _build_spatial_connectivity(
        self,
        coords: np.ndarray,
        emb_indv: np.ndarray,
        emb_niche: np.ndarray = None,
        cell_mask: np.ndarray | None = None,
        high_quality_mask: np.ndarray = None,
        slice_ids: np.ndarray | None = None,
        k_central: int = 101,
        k_adjacent: int = 50,
        n_adjacent_slices: int = 1
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Internal method for building spatial connectivity matrix.

        Args:
            coords: Spatial coordinates (n_cells, 2) - only x,y coordinates
            emb_niche: Niche embeddings (n_cells, d1) - always provided (may be dummy ones)
            emb_indv: Cell identity embeddings (n_cells, d2)
            cell_mask: Boolean mask for cells to process
            high_quality_mask: Boolean mask for high quality cells (used for neighbor search)
            slice_ids: Optional slice/z-coordinate indices (n_cells,) for 3D data
            k_central: Number of neighbors on central slice
            k_adjacent: Number of neighbors on adjacent slices for 3D data
            n_adjacent_slices: Number of slices to search above/below for 3D data

        Returns:
            Tuple of (neighbor_indices, cell_sims, niche_sims) arrays
        """

        n_cells = len(coords)
        if cell_mask is None:
            cell_mask = np.ones(n_cells, dtype=bool)

        cell_indices = np.where(cell_mask)[0]
        n_masked = len(cell_indices)

        # Step 1: Find spatial neighbors (slice-aware if slice_ids provided)
        spatial_neighbors, neighbor_pool_per_slice = find_spatial_neighbors_with_slices(
            coords=coords,
            slice_ids=slice_ids,
            query_cell_mask=cell_mask,
            high_quality_cell_mask=high_quality_mask,
            k_central=k_central,
            k_adjacent=k_adjacent,
            n_adjacent_slices=n_adjacent_slices
        )

        # Log statistics about neighbor pool cells per slice if available
        if neighbor_pool_per_slice:
            for slice_id, neighbor_pool_indices in neighbor_pool_per_slice.items():
                logger.debug(f"Slice {slice_id}: {len(neighbor_pool_indices)} high quality cells available for neighbor search")

        # Step 2 & 3: Find anchors and homogeneous neighbors in batches
        logger.info(f"Finding anchors and homogeneous neighbors (batch size: {self.find_homogeneous_batch_size})...")

        # Convert embeddings to JAX arrays once (shared for both paths)
        # Note: float16 provides sufficient precision for normalized embeddings
        # emb_niche is guaranteed to exist now (may be dummy ones)
        all_emb_niche_norm_jax = jnp.array(emb_niche, dtype=jnp.float16)
        all_emb_indv_norm_jax = jnp.array(emb_indv, dtype=jnp.float16)

        # Get masked embeddings
        masked_cell_indices = np.where(cell_mask)[0]
        emb_niche_masked_jax = all_emb_niche_norm_jax[masked_cell_indices]
        emb_indv_masked_jax = all_emb_indv_norm_jax[masked_cell_indices]


        if self.config.fix_cross_slice_homogenous_neighbors:
            logger.info(f"Using 3D constrained selection (ensuring {self.config.homogeneous_neighbors} neighbors per slice)")

            # Use memory-efficient version that processes slices separately
            homogeneous_neighbors, homogeneous_cell_sims, homogeneous_niche_sims = _find_homogeneous_3d_memory_efficient(
                emb_niche_masked_jax=emb_niche_masked_jax,
                emb_indv_masked_jax=emb_indv_masked_jax,
                spatial_neighbors=spatial_neighbors,
                all_emb_niche_norm_jax=all_emb_niche_norm_jax,
                all_emb_indv_norm_jax=all_emb_indv_norm_jax,
                num_homogeneous_per_slice=self.config.homogeneous_neighbors,
                k_central=k_central,
                k_adjacent=k_adjacent,
                n_adjacent_slices=n_adjacent_slices,
                cell_embedding_similarity_threshold=self.config.cell_embedding_similarity_threshold,
                spatial_domain_similarity_threshold=self.config.spatial_domain_similarity_threshold,
                find_homogeneous_batch_size=self.find_homogeneous_batch_size
            )

        else:
            # Convert spatial_neighbors to JAX array for regular processing
            spatial_neighbors_jax = jnp.array(spatial_neighbors, dtype=jnp.int32)
            homogeneous_neighbors_list = []
            homogeneous_cell_sims_list = []
            homogeneous_niche_sims_list = []


            # Use the standard function (2D or 3D without fix_cross_slice_homogenous_neighbors)
            for batch_start in track(range(0, n_masked, self.find_homogeneous_batch_size), description="Finding homogeneous neighbors", transient=True):
                batch_end = min(batch_start + self.find_homogeneous_batch_size, n_masked)
                batch_indices = slice(batch_start, batch_end)

                # Get batch data directly from JAX arrays (no GPU movement)
                emb_niche_batch_norm = emb_niche_masked_jax[batch_indices]
                emb_indv_batch_norm = emb_indv_masked_jax[batch_indices]
                spatial_neighbors_batch = spatial_neighbors_jax[batch_indices]

                # Process batch with single JIT-compiled function
                homo_neighbors_batch, cell_sims_batch, niche_sims_batch = _find_homogeneous_spot_dual_embedding_batch_jit(
                    emb_niche_batch_norm=emb_niche_batch_norm,
                    emb_indv_batch_norm=emb_indv_batch_norm,
                    spatial_neighbors=spatial_neighbors_batch,
                    all_emb_niche_norm=all_emb_niche_norm_jax,
                    all_emb_indv_norm=all_emb_indv_norm_jax,
                    homogeneous_neighbors=self.config.total_homogeneous_neighbor_per_cell,
                    cell_embedding_similarity_threshold=self.config.cell_embedding_similarity_threshold,
                    spatial_domain_similarity_threshold=self.config.spatial_domain_similarity_threshold
            )

                # Convert back to numpy and append
                homogeneous_neighbors_list.append(np.array(homo_neighbors_batch))
                homogeneous_cell_sims_list.append(np.array(cell_sims_batch))
                homogeneous_niche_sims_list.append(np.array(niche_sims_batch))

            # Regular batched processing
            homogeneous_neighbors = np.vstack(homogeneous_neighbors_list)
            homogeneous_cell_sims = np.vstack(homogeneous_cell_sims_list)
            homogeneous_niche_sims = np.vstack(homogeneous_niche_sims_list)

        # Return dense format: (n_masked, num_homogeneous) arrays
        return homogeneous_neighbors, homogeneous_cell_sims, homogeneous_niche_sims
