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

import jax
import jax.numpy as jnp
import jax.random as random
import numpy as np
import pandas as pd
from anndata import AnnData
from beartype import beartype
from numpyro.infer import Predictive
from scipy import sparse


@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",
    prior_time_col: Optional[str] = None,
    batch_annotation: Optional[str] = None,
    prior_path_col: 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. In addition, reads a hard prior on path
    assignment from `adata_rna.obs[prior_path_col]`, where integer entries 1..P indicate
    known path membership (converted to 0-based), and other values (None/NaN) are treated
    as unknown (-1) for inference.

    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").
    prior_time_col : Optional[str], optional
        Column name in `adata_rna.obs` containing prior pseudotimes; if provided, used
        to center the prior over t_c (default: None).
    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.
    prior_path_col : str, optional
        Column name in `adata_rna.obs` that contains a hard prior assignment (1-based)
        to path indices; missing or invalid entries become -1 (no prior).

    Returns
    -------
    Dict[str, Any]
        Dictionary containing:
          - `data`: JAX array of shape (cells, genes, 2) with stacked unspliced and spliced counts.
          - `M_c`: JAX array of shape (cells, 1, 1) holding the per-cell metacell size.
          - `batch_index`: JAX array of shape (cells,) with batch indices.
          - `tf_indices`: JAX array of TF gene indices.
          - `total_num_cells`: int, number of cells.
          - `n_batch`: int, number of batches.
          - `prior_time`: JAX array of prior times or None.
          - `prior_timespan`: float, span of prior_time or default.
          - `prior_path`: JAX int array of shape (cells,) with –1 or 0-based path ids.
          - `unknown_idx`: Python list of ints, indices of cells with unknown prior.
    """
    # 1) stack raw data

    def _to_dense_f32(layer):
        if sparse.issparse(layer):
            # cast to float32 and densify in one go
            return layer.astype(np.float32).toarray()
        else:
            # if it’s already an ndarray, just ensure dtype
            return np.asarray(layer, dtype=np.float32)

    # pull out each layer
    spliced_layer = adata_rna.layers["spliced"]
    unspliced_layer = adata_rna.layers["unspliced"]

    # densify (only if needed) and get float32 arrays
    spliced = _to_dense_f32(spliced_layer)
    unspliced = _to_dense_f32(unspliced_layer)

    # allocate final array of shape (..., 2)
    #   last axis 0=unspliced, 1=spliced
    data_np = np.empty(spliced.shape + (2,), dtype=np.float32)
    data_np[..., 0] = unspliced
    data_np[..., 1] = spliced

    # convert once into JAX
    data = jnp.array(data_np)

    # 2) metacell sizes
    M_c = jnp.array(
        np.expand_dims(np.expand_dims(adata_rna.obs[n_cells_col].to_numpy(), -1), -1),
        dtype=jnp.float32,
    )
    total_num_cells = int(data.shape[0])

    # 3) prior time
    if prior_time_col is not None:
        prior_time = jnp.array(
            adata_rna.obs[prior_time_col].to_numpy(), dtype=jnp.float32
        )
        prior_timespan = float(prior_time.max() - prior_time.min())
        T_min = prior_time.min() - prior_timespan / 6.0
        T_max = prior_time.max() + prior_timespan / 6.0
    else:
        prior_time = None
        T_min = -prior_timespan / 2.0
        T_max = +prior_timespan / 2.0

    # 4) batch index
    if batch_annotation is not None:
        batches = adata_rna.obs[batch_annotation].astype(str)
        uniq = batches.unique()
        mapping = {b: i for i, b in enumerate(uniq)}
        batch_index = jnp.array(batches.map(mapping).to_numpy(), dtype=jnp.int32)
    else:
        batch_index = jnp.zeros(total_num_cells, dtype=jnp.int32)
    n_batch = int(jnp.unique(batch_index).shape[0])

    obs2sample = jax.nn.one_hot(batch_index, num_classes=n_batch, dtype=jnp.float32)

    # 5) TF indices
    is_tf = np.isin(adata_rna.var_names, tf_list)
    tf_indices = jnp.array(np.where(is_tf)[0], dtype=jnp.int32)

    # 6) prior path
    if prior_path_col is not None:
        raw_prior = adata_rna.obs[prior_path_col].to_numpy()
        prior_path = np.full(raw_prior.shape, -1, dtype=np.int32)
        for i, v in enumerate(raw_prior):
            if pd.notna(v):
                try:
                    prior_path[i] = int(v) - 1
                except:
                    prior_path[i] = -1
        prior_path = jnp.array(prior_path, dtype=jnp.int32)
    else:
        # no column → everyone unknown
        prior_path = jnp.full((total_num_cells,), -1, dtype=jnp.int32)

    # 7) unknown indices
    unknown_idx = [int(i) for i, p in enumerate(prior_path) if int(p) < 0]

    return OrderedDict(
        [
            ("data", data),
            ("M_c", M_c),
            ("obs2sample", obs2sample),
            ("batch_index", batch_index),
            ("tf_indices", tf_indices),
            ("total_num_cells", total_num_cells),
            ("n_batch", n_batch),
            ("prior_time", prior_time),
            ("prior_timespan", prior_timespan),
            ("unknown_idx", unknown_idx),
            ("T_limits", (T_min, T_max)),
            ("unknown_idx", unknown_idx),
        ]
    )
