from collections import OrderedDict
from typing import Any, Callable, Dict, List, Optional

import jax.numpy as jnp
import jax.random as random
import numpy as np
from anndata import AnnData
from beartype import beartype
from numpyro.infer import Predictive


@beartype
def sample_prior(
    model: Callable[..., Any],
    model_input: Dict[str, Any],
    num_samples: int = 5,
    seed: int = 42,
) -> Dict[str, Any]:
    """
    Runs prior predictive sampling for a given model.

    Parameters:
    - model (Callable[..., Any]): The probabilistic model to sample from.
    - model_input (Dict[str, Any]): Input arguments to be passed to the model.
    - num_samples (int): Number of prior predictive samples to generate.
    - seed (int): Random seed for reproducibility.

    Returns:
    - Dict[str, Any]: Prior predictive samples generated by the model.
    """
    rng_key = random.PRNGKey(seed)  # Generate a random seed
    predictive = Predictive(model, num_samples=num_samples)  # Create predictive object
    prior_samples = predictive(rng_key, **model_input)  # Generate samples
    return prior_samples


def convert_to_dense(array):
    """Convert sparse array to dense if necessary."""
    return array.toarray() if hasattr(array, "toarray") else np.array(array)


@beartype
def prepare_model_input(
    adata_rna: AnnData,
    tf_list: List[str],
    n_cells_col: str = "n_cells",
    batch_annotation: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Prepare RNA input data for the model by extracting spliced and unspliced counts,
    computing transcription factor indices, and extracting cell metadata including batch
    information from adata_rna.obs if provided.

    Parameters
    ----------
    adata_rna : AnnData
        AnnData object containing RNA expression data with 'spliced' and 'unspliced' layers.
    tf_list : List[str]
        List of transcription factor names.
    n_cells_col : str, optional
        Column name in `adata_rna.obs` representing the number of cells per metacell
        (default: "n_cells").
    batch_annotation : Optional[str], optional
        Column name in `adata_rna.obs` that contains batch information.
        If provided, cells are assigned batch indices based on this column;
        otherwise, all cells are assigned to batch 1.

    Returns
    -------
    Dict[str, Any]
        Dictionary containing:
          - `data`: JAX array of shape (cells, features, 2) with stacked unspliced and spliced RNA counts.
          - `M_c`: JAX array of shape (cells, 1, 1) holding the per-cell metacell size.
          - `batch_index`: JAX array with batch indices for each cell.
          - `tf_indices`: JAX array of transcription factor indices.
    """
    # Convert RNA layers to dense format
    spliced_dense = convert_to_dense(adata_rna.layers["spliced"])
    unspliced_dense = convert_to_dense(adata_rna.layers["unspliced"])

    # Stack RNA data into shape (cells, features, 2)
    data = np.stack([unspliced_dense, spliced_dense], axis=-1).astype(np.float32)
    data = jnp.array(data)

    # Extract M_c (number of cells per metacell)
    M_c = jnp.array(
        np.expand_dims(np.expand_dims(adata_rna.obs[n_cells_col].to_numpy(), -1), -1)
    )
    total_num_cells = M_c.shape[0]

    # Build the batch index from adata_rna.obs if a batch annotation is provided
    if batch_annotation is not None:
        # Get the batch information from the observation dataframe.
        batch_series = adata_rna.obs[batch_annotation]
        # Map each unique batch label to an integer.
        unique_batches = batch_series.unique()
        batch_map = {batch: i for i, batch in enumerate(unique_batches)}
        batch_index = jnp.array(batch_series.map(batch_map).to_numpy())
    else:
        # If no batch annotation is provided, assign all cells to batch 1.
        batch_index = jnp.zeros(adata_rna.n_obs, dtype=int)

    # Identify transcription factor indices in adata_rna.var_names.
    is_tf = adata_rna.var_names.isin(tf_list)
    tf_indices = jnp.array(np.where(is_tf)[0])

    return OrderedDict(
        [
            ("data", data),
            ("M_c", M_c),
            ("batch_index", batch_index),
            ("tf_indices", tf_indices),
            ("total_num_cells", total_num_cells),
        ]
    )
