from __future__ import annotations

from typing import Literal, Optional, Tuple

import numpy as np
from scipy import sparse
from sklearn.neighbors import NearestNeighbors


def _to_dense_float32(x) -> np.ndarray:
    if sparse.issparse(x):
        x = x.toarray()
    return np.asarray(x, dtype=np.float32)


def make_stage_bins(
    pseudotime: np.ndarray,
    *,
    n_bins: int = 20,
    strategy: Literal["quantile", "uniform"] = "quantile",
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Discretize pseudotime into ordered stage bins.

    Returns:
      bin_id: int array [n_cells] in [0, n_bins_eff-1]
      edges: float array of (possibly non-unique) bin edges (for plotting/debugging)
    """
    pt = np.asarray(pseudotime, dtype=np.float64)
    if pt.ndim != 1:
        raise ValueError("pseudotime must be 1D")
    if not np.isfinite(pt).all():
        raise ValueError("pseudotime contains NaN/Inf")
    if n_bins < 2:
        raise ValueError("n_bins must be >= 2")

    if strategy == "quantile":
        qs = np.linspace(0.0, 1.0, n_bins + 1)
        edges = np.quantile(pt, qs)
    elif strategy == "uniform":
        edges = np.linspace(float(pt.min()), float(pt.max()), n_bins + 1)
    else:
        raise ValueError("strategy must be 'quantile' or 'uniform'")

    # If many ties exist, quantile edges can repeat. We still use them for digitize,
    # then remap to contiguous non-empty bins.
    bin_id = np.digitize(pt, edges[1:-1], right=False).astype(np.int64)
    present = np.unique(bin_id)
    remap = {old: new for new, old in enumerate(present.tolist())}
    bin_id = np.vectorize(remap.get, otypes=[np.int64])(bin_id)
    return bin_id, edges


def compute_future_target_knn(
    adata,
    embedding: np.ndarray,
    *,
    time_key: str = "dpt_pseudotime",
    n_neighbors: int = 30,
    mode: Literal["average", "nearest"] = "average",
    epsilon: float = 1e-3,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    TIVelo-style pairing: KNN in embedding space, then keep only neighbors with later pseudotime.

    Returns:
      future_X: array [n_valid, n_genes]
      valid_indices: int array [n_valid] (indices into original adata)
    """
    pseudotime = np.asarray(adata.obs[time_key].values, dtype=np.float64)
    X = _to_dense_float32(adata.X)

    n_cells = int(X.shape[0])
    if embedding.shape[0] != n_cells:
        raise ValueError("embedding and adata must have same n_cells")

    search_k = min(n_cells, max(n_neighbors * 2, n_neighbors + 1))
    nbrs = NearestNeighbors(n_neighbors=search_k, algorithm="auto").fit(embedding)
    _, indices = nbrs.kneighbors(embedding)

    future_X_list = []
    valid_indices = []
    for i in range(n_cells):
        t_i = pseudotime[i]
        neigh = indices[i]
        future = neigh[pseudotime[neigh] > (t_i + epsilon)]
        if future.size == 0:
            continue

        if mode == "average":
            future = future[: min(n_neighbors, future.size)]
            target_expr = X[future].mean(axis=0)
        elif mode == "nearest":
            target_expr = X[int(future[0])]
        else:
            raise ValueError("mode must be 'average' or 'nearest'")

        future_X_list.append(target_expr)
        valid_indices.append(i)

    return np.asarray(future_X_list, dtype=np.float32), np.asarray(valid_indices, dtype=np.int64)


def compute_future_target_stage_binned(
    adata,
    embedding: np.ndarray,
    *,
    time_key: str = "dpt_pseudotime",
    n_neighbors: int = 30,
    n_bins: int = 20,
    binning: Literal["quantile", "uniform"] = "quantile",
    step: int = 1,
    mode: Literal["average", "nearest", "sample"] = "average",
    epsilon: float = 1e-3,
    random_state: Optional[int] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Stage-binned pairing: bin pseudotime into stages, then pair each cell only to cells in the next stage.

    This reduces "dt" variance and makes the pairing distribution more controllable than global KNN+time filter.

    Returns:
      future_X: array [n_valid, n_genes]
      valid_indices: int array [n_valid] (indices into original adata)
    """
    if step < 1:
        raise ValueError("step must be >= 1")

    pseudotime = np.asarray(adata.obs[time_key].values, dtype=np.float64)
    X = _to_dense_float32(adata.X)

    n_cells = int(X.shape[0])
    if embedding.shape[0] != n_cells:
        raise ValueError("embedding and adata must have same n_cells")

    bin_id, _ = make_stage_bins(pseudotime, n_bins=n_bins, strategy=binning)
    n_bins_eff = int(bin_id.max()) + 1

    rng = np.random.default_rng(random_state)

    future_X_list = []
    valid_indices = []

    for b in range(0, n_bins_eff - step):
        src = np.flatnonzero(bin_id == b)
        tgt = np.flatnonzero(bin_id == (b + step))
        if src.size == 0 or tgt.size == 0:
            continue

        k = min(int(n_neighbors), int(tgt.size))
        if k <= 0:
            continue

        # Fit KNN only on the next-stage cells, then query for current-stage cells.
        nbrs = NearestNeighbors(n_neighbors=k, algorithm="auto").fit(embedding[tgt])
        _, nn = nbrs.kneighbors(embedding[src])
        # Map neighbor indices back to original cell indices
        nn_global = tgt[nn]  # shape [len(src), k]

        for row, i in enumerate(src.tolist()):
            # Optional time guard inside bins (handles bin-edge ties/noise)
            choices = nn_global[row]
            choices = choices[pseudotime[choices] > (pseudotime[i] + epsilon)]
            if choices.size == 0:
                continue

            if mode == "average":
                target_expr = X[choices].mean(axis=0)
            elif mode == "nearest":
                target_expr = X[int(choices[0])]
            elif mode == "sample":
                target_expr = X[int(rng.choice(choices))]
            else:
                raise ValueError("mode must be 'average', 'nearest', or 'sample'")

            future_X_list.append(target_expr)
            valid_indices.append(i)

    return np.asarray(future_X_list, dtype=np.float32), np.asarray(valid_indices, dtype=np.int64)

