from __future__ import annotations

from typing import Any, Dict

import jax
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from beartype import beartype
from numpyro.handlers import scale

from .RNA_utils import prepare_model_input, sample_prior
from .utils import solve_DE, sort_times_over_all_cells


@beartype
def mlp(params: dict, x: jnp.ndarray) -> jnp.ndarray:
    """
    Multilayer Perceptron (MLP) with residual connections.

    Args:
        params: Dictionary containing neural network parameters (weights and biases).
        x: Input data array.

    Returns:
        Output array after passing through the MLP.
    """

    # Block 1
    h = x @ params["W1"] + params["b1"]  # (n_cells, 512)
    h = jax.nn.softplus(h)
    h = h @ params["W2"] + params["b2"]  # (n_cells, 1000)
    h = jax.nn.softplus(h)
    h = h + x  # residual

    # Final head
    h2 = h @ params["W3"] + params["b3"]  # (n_cells, 512)
    h2 = jax.nn.softplus(h2)
    out = h2 @ params["W4"] + params["b4"]  # (n_cells, 2000)

    return out


# Define the ODE function using the MLP
@beartype
def dstate_dt(t, state, args):
    """
    Compute the derivative of the state vector for the coupled system using a neural network.

    Args:
        t: Time scalar.
        state: State vector [u_1, ..., u_G, s_1, ..., s_G].
        args: Tuple containing parameters (G, beta_g, gamma_g, nn_params, T_ON).

    Returns:
        Derivative of the state vector
    """

    alpha_0, beta_g, gamma_g, lamda, kappa, nn_params, tf_indices, T_ON = args
    u = state[..., 0]  # Unspliced counts for all genes.
    s = state[..., 1]  # Spliced counts for all genes.
    p = state[..., 2]  # Proteins for all genes.

    # Clip state values to increase numerical stability:
    u = jnp.clip(u, 0, 1e3)
    s = jnp.clip(s, 0, 1e3)
    p = jnp.clip(p, 0, 1e3)

    # Compute transcription rates using the neural network
    alpha = alpha_0 * jnp.clip(mlp(nn_params, p[:, tf_indices]), 0, 1e3)  # Shape: (G,)

    du_dt = alpha - beta_g * u
    ds_dt = beta_g * u - gamma_g * s
    # Compute dp_dt only for TF indices
    dp_dt = jnp.zeros_like(u)
    dp_dt = dp_dt.at[:, tf_indices].set(
        lamda * s[:, tf_indices] - kappa * p[:, tf_indices]
    )

    dstate_dt = jnp.stack([du_dt, ds_dt, dp_dt], axis=-1)  # Shape: (3G,)
    return dstate_dt


# Define the complete NumPyro model
@beartype
def model(
    data: Any,
    M_c: Any,
    batch_index: Any,
    tf_indices: Any,
    total_num_cells: int,
    n_batch: int,
    prior_time: Any,
    prior_timespan: Any,
    prior_path: Any,
    unknown_idx: Any,
    num_paths: int,
    return_alpha: bool = False,
    Tmax_alpha: float = 50.0,
    Tmax_beta: float = 1.0,
    splicing_rate_alpha_hyp_prior_alpha: float = 20.0,
    splicing_rate_alpha_hyp_prior_mean: float = 5.0,
    splicing_rate_mean_hyp_prior_alpha: float = 10.0,
    splicing_rate_mean_hyp_prior_mean: float = 1.0,
    degradation_rate_alpha_hyp_prior_alpha: float = 20.0,
    degradation_rate_alpha_hyp_prior_mean: float = 5.0,
    degradation_rate_mean_hyp_prior_alpha: float = 10.0,
    degradation_rate_mean_hyp_prior_mean: float = 1.0,
    transcription_rate_alpha_hyp_prior_alpha: float = 20.0,
    transcription_rate_alpha_hyp_prior_mean: float = 2.0,
    transcription_rate_mean_hyp_prior_alpha: float = 10.0,
    transcription_rate_mean_hyp_prior_mean: float = 5.0,
    lambda_alpha: float = 1.0,
    lambda_mean: float = 1.0,
    kappa_alpha: float = 1.0,
    kappa_mean: float = 1.0,
    detection_mean_hyp_prior_alpha: float = 1.0,
    detection_mean_hyp_prior_beta: float = 1.0,
    detection_hyp_prior_alpha: float = 10.0,
    detection_i_prior_alpha: float = 100.0,
    detection_gi_prior_alpha: float = 200.0,
    gene_add_alpha_hyp_prior_alpha: float = 9.0,
    gene_add_alpha_hyp_prior_beta: float = 3.0,
    gene_add_mean_hyp_prior_alpha: float = 1.0,
    gene_add_mean_hyp_prior_beta: float = 100.0,
    stochastic_v_ag_hyp_prior_alpha: float = 9.0,
    stochastic_v_ag_hyp_prior_beta: float = 3.0,
    sde_rng_key=0,
):
    """
    NumPyro model for coupled transcription and splicing dynamics.

    Args:
        data: Observed data array of shape (num_cells, num_genes, num_modalities).
        M_c: Number of cells in each metacell.
        batch_index: Array indicating batch assignments for each cell.
        tf_indices: Indices of genes that are TFs.
        total_num_cells: Number of cells in the full dataset.
        n_batch: Number of batches.
        Tmax_alpha: Alpha parameter for Tmax prior.
        Tmax_beta: Beta parameter for Tmax prior.
        ... (other hyperparameters for priors)
        key: Random number generator key.

    Returns:
        None. Defines the probabilistic model for inference.
    """

    num_cells = int(data.shape[0])
    num_genes = int(data.shape[1])
    num_modalities = int(data.shape[2])
    num_tfs = tf_indices.shape[0]
    batch_size = num_cells
    obs2sample = jax.nn.one_hot(
        batch_index, num_classes=n_batch
    )  # Shape: (num_cells, n_batch)

    # Splicing Rates for mRNA
    splicing_alpha = numpyro.sample(
        "splicing_alpha",
        dist.Gamma(
            splicing_rate_alpha_hyp_prior_alpha,
            splicing_rate_alpha_hyp_prior_alpha / splicing_rate_alpha_hyp_prior_mean,
            validate_args=True,
        ),
    )
    splicing_alpha = jnp.clip(splicing_alpha, a_min=1e-2, a_max=1e3)

    splicing_mean = numpyro.sample(
        "splicing_mean",
        dist.Gamma(
            splicing_rate_mean_hyp_prior_alpha,
            splicing_rate_mean_hyp_prior_alpha / splicing_rate_mean_hyp_prior_mean,
            validate_args=True,
        ),
    )
    splicing_mean = jnp.clip(splicing_mean, a_min=1e-2, a_max=1e3)

    beta_g = numpyro.sample(
        "beta_g",
        dist.Gamma(splicing_alpha, splicing_alpha / splicing_mean, validate_args=True)
        .expand([num_genes])
        .to_event(1),
    )

    # Degradation Rates for mRNA
    degradation_alpha = numpyro.sample(
        "degradation_alpha",
        dist.Gamma(
            degradation_rate_alpha_hyp_prior_alpha,
            degradation_rate_alpha_hyp_prior_alpha
            / degradation_rate_alpha_hyp_prior_mean,
            validate_args=True,
        ),
    )
    degradation_alpha = degradation_alpha + 0.001  # Prevent zero

    degradation_mean = numpyro.sample(
        "degradation_mean",
        dist.Gamma(
            degradation_rate_mean_hyp_prior_alpha,
            degradation_rate_mean_hyp_prior_alpha
            / degradation_rate_mean_hyp_prior_mean,
            validate_args=True,
        ),
    )
    degradation_mean = jnp.clip(degradation_mean, a_min=1e-2, a_max=1e3)

    gamma_g = numpyro.sample(
        "gamma_g",
        dist.Gamma(
            degradation_alpha, degradation_alpha / degradation_mean, validate_args=True
        )
        .expand([num_genes])
        .to_event(1),
    )

    # Translation rate for proteins
    lamda = numpyro.sample(
        "lambda",
        dist.Gamma(lambda_alpha, lambda_alpha / lambda_mean, validate_args=True)
        .expand([num_tfs])
        .to_event(1),
    )
    # Degradation rate for proteins
    kappa = numpyro.sample(
        "kappa",
        dist.Gamma(kappa_alpha, kappa_alpha / kappa_mean, validate_args=True)
        .expand([num_tfs])
        .to_event(1),
    )

    # Time Parameters
    prior_time_sd = prior_timespan / 6
    T_scaling = numpyro.sample("T_scaling", dist.Gamma(1.0, 1.0))
    if prior_time is None:
        with numpyro.plate("cells", batch_size):
            lower_limit = -prior_timespan / 2
            upper_limit = prior_timespan / 2
            t_c = numpyro.sample(
                "t_c",
                dist.TruncatedNormal(
                    low=lower_limit, high=upper_limit, loc=0, scale=prior_time_sd
                ),
            )
            T_c = numpyro.deterministic(
                "T_c", jnp.clip(t_c, min=lower_limit, max=upper_limit)
            )
    else:
        with numpyro.plate("cells", batch_size):
            lower_limit = jnp.min(prior_time) - 3 * prior_time_sd
            upper_limit = jnp.max(prior_time) + 3 * prior_time_sd
            t_c = numpyro.sample(
                "t_c",
                dist.TruncatedNormal(
                    low=-4 * prior_time_sd,
                    high=4 * prior_time_sd,
                    loc=0,
                    scale=prior_time_sd,
                ),
            )
            T_c = numpyro.deterministic("T_c", T_scaling * (t_c + prior_time))
    T_ON = lower_limit - 1

    # ============= Expression model =============== #

    # scale of alpha
    alpha_0 = numpyro.sample(
        "alpha_0",
        dist.Gamma(1.0, 1.0, validate_args=True).expand([num_genes]).to_event(1),
    )

    # Hyper-prior for the shared “center” of (unspliced, spliced)
    init_center_2d = numpyro.sample(
        "init_center_2d",
        dist.Gamma(2.0, 0.5).expand([num_genes, 2]).to_event(2),
    )

    # Hyper-prior for the shared center of TF-protein initials
    init_center_tf = numpyro.sample(
        "init_center_tf",
        dist.Gamma(2.0, 0.5).expand([len(tf_indices)]).to_event(1),
    )

    # Small Gaussian jitter around zero for each path
    σ_dev = 1.0
    jitter_2d = numpyro.sample(
        "jitter_2d",
        dist.Normal(0.0, σ_dev).expand([num_paths, num_genes, 2]).to_event(3),
    )
    jitter_tf = numpyro.sample(
        "jitter_tf",
        dist.Normal(0.0, σ_dev).expand([num_paths, len(tf_indices)]).to_event(2),
    )

    # Build each path’s initial state by adding jitter to the shared center
    initial_state_2d = init_center_2d[None, :, :] + jitter_2d
    # → shape (num_paths, num_genes, 2)

    # Now insert the TF-protein entries into the third dim
    third_dim = jnp.zeros((num_paths, num_genes))
    third_dim = third_dim.at[:, tf_indices].set(init_center_tf[None, :] + jitter_tf)
    initial_state = jnp.concatenate(
        [initial_state_2d, third_dim[..., None]],
        axis=-1,
    )

    initial_state = numpyro.deterministic(
        "initial_state",
        jnp.concatenate(
            [initial_state_2d, third_dim[..., None]],
            axis=-1,
        ),
    )

    # Neural Net parameters
    rng = jax.random.PRNGKey(0)
    P = numpyro.param

    nn_params = {
        # Block 1
        "W1": P("W1", jax.random.normal(rng, (num_tfs, 512)) * 0.01),
        "b1": P("b1", jnp.zeros((512,))),
        "W2": P("W2", jax.random.normal(rng, (512, num_tfs)) * 0.01),
        "b2": P("b2", jnp.zeros((num_tfs,))),
        # Final head
        "W3": P("W3", jax.random.normal(rng, (num_tfs, 512)) * 0.01),
        "b3": P("b3", jnp.zeros((512,))),
        "W4": P("W4", jax.random.normal(rng, (512, num_genes)) * 0.01),
        "b4": P("b4", jnp.zeros((num_genes,))),
    }

    # Prepare Parameters for ODE Solver
    params = (alpha_0, beta_g, gamma_g, lamda, kappa, nn_params, tf_indices, T_ON)

    # Get Ordered Time Vector
    all_times, time_indices, _ = sort_times_over_all_cells(T_c)

    # Solve the Coupled ODE
    predictions = solve_DE(
        ts=all_times.squeeze(),
        params=params,
        initial_state=initial_state,
        time_step=0.1,
        drift=dstate_dt,
        t0=lower_limit,
        t1=upper_limit,
    )

    predictions = numpyro.deterministic("predictions", predictions.astype(jnp.float32))

    predictions_rearranged = numpyro.deterministic(
        "predictions_rearranged", predictions[time_indices.ravel(), :]
    )

    # —— path_weights —— #

    # one-hot for the cells with known path (works even if prior_path is None)
    if prior_path is None:
        one_hot = jnp.zeros((batch_size, num_paths))
        unknown_idx = list(range(batch_size))
    else:
        one_hot = jax.nn.one_hot(prior_path, num_paths)  # (N, K)
        unknown_idx = unknown_idx  # comes in as a Python list

    n_unknown = len(unknown_idx)

    # sample only for unknown cells
    w_unknown = numpyro.sample(
        "path_weights_unknown",
        dist.Dirichlet(jnp.ones(num_paths)).expand([n_unknown]).to_event(1),
    )  # shape (n_unknown, K)

    # scatter them back into a full array
    path_weights = one_hot
    for j, i in enumerate(unknown_idx):
        path_weights = path_weights.at[i].set(w_unknown[j])

    numpyro.deterministic("path_weights", path_weights)  # (N, K)

    # AVERAGE THE predictions REARRANGED BY THESE WEIGHTS ---
    weighted_preds = jnp.sum(
        predictions_rearranged * path_weights[:, :, None, None], axis=1
    )  # → (num_cells, num_genes, 3)

    # now downstream we use weighted_preds as mu_expression
    mu_expression = jnp.clip(
        weighted_preds[..., :2],
        a_min=1e-5,
        a_max=1e5,
    )

    # if return_alpha:
    #     final_p_clipped = jnp.clip(predictions_rearranged[..., 2].squeeze(1), 0, 1e3)
    #     alpha_cg = jnp.clip(
    #         alpha_0 * mlp(nn_params, final_p_clipped[:, tf_indices]), 0, 1e3
    #     )
    #     alpha_cg = numpyro.deterministic("alpha_cg", alpha_cg)

    # ============= Detection efficiency of spliced and unspliced counts =============== #

    detection_mean_y_e = numpyro.sample(
        "detection_mean_y_e",
        dist.Beta(
            jnp.ones((1, 1)) * detection_mean_hyp_prior_alpha,
            jnp.ones((1, 1)) * detection_mean_hyp_prior_beta,
            validate_args=True,
        )
        .expand([n_batch, 1])
        .to_event(2),
    )

    beta = ((detection_hyp_prior_alpha / (obs2sample @ detection_mean_y_e)).T)[0, ...]

    with numpyro.plate("cells", batch_size):
        detection_y_c = numpyro.sample(
            "detection_y_c", dist.Gamma(detection_hyp_prior_alpha, beta)
        )

    detection_y_i = numpyro.sample(
        "detection_y_i",
        dist.Gamma(detection_i_prior_alpha, detection_i_prior_alpha, validate_args=True)
        .expand([1, 1, 2])
        .to_event(3),
    )

    detection_y_gi = numpyro.sample(
        "detection_y_gi",
        dist.Gamma(
            jnp.ones((1, 1)) * detection_gi_prior_alpha,
            jnp.ones((1, 1)) * detection_gi_prior_alpha,
            validate_args=True,
        )
        .expand([1, num_genes, 2])
        .to_event(3),
    )

    # ======= Gene-specific additive component (Ambient RNA/ "Soup") for spliced and unspliced counts ====== #

    s_g_gene_add_alpha_hyp = numpyro.sample(
        "s_g_gene_add_alpha_hyp",
        dist.Gamma(
            gene_add_alpha_hyp_prior_alpha,
            gene_add_alpha_hyp_prior_beta,
            validate_args=True,
        ),
        sample_shape=(2,),
    )

    s_g_gene_add_mean = numpyro.sample(
        "s_g_gene_add_mean",
        dist.Gamma(
            gene_add_mean_hyp_prior_alpha,
            gene_add_mean_hyp_prior_beta,
            validate_args=True,
        )
        .expand([n_batch, 1, 2])
        .to_event(3),
    )

    s_g_gene_add_alpha_e_inv = numpyro.sample(
        "s_g_gene_add_alpha_e_inv",
        dist.Exponential(s_g_gene_add_alpha_hyp).expand([n_batch, 1, 2]).to_event(3),
    )

    s_g_gene_add_alpha_e = jnp.ones((1, 1)) / s_g_gene_add_alpha_e_inv**2

    s_g_gene_add = numpyro.sample(
        "s_g_gene_add",
        dist.Gamma(
            s_g_gene_add_alpha_e,
            s_g_gene_add_alpha_e / s_g_gene_add_mean,
            validate_args=True,
        )
        .expand([n_batch, num_genes, 2])
        .to_event(3),
    )

    # ========= Gene-specific overdispersion of spliced and unspliced counts ============== #

    stochastic_v_ag_hyp = numpyro.sample(
        "stochastic_v_ag_hyp",
        dist.Gamma(
            stochastic_v_ag_hyp_prior_alpha,
            stochastic_v_ag_hyp_prior_beta,
            validate_args=True,
        )
        .expand([1, 2])
        .to_event(2),
    )

    stochastic_v_ag_inv = numpyro.sample(
        "stochastic_v_ag_inv",
        dist.Exponential(stochastic_v_ag_hyp).expand([1, num_genes, 2]).to_event(3),
    )

    stochastic_v_ag = jnp.ones((1, 1)) / stochastic_v_ag_inv**2

    # ===================== Expected expression ======================= #

    additive_term = numpyro.deterministic(
        "additive_term", jnp.einsum("cb,bgi->cgi", obs2sample, s_g_gene_add)
    )

    normalizing_term = numpyro.deterministic(
        "normalizing_term",
        detection_y_c[:, None, None] * detection_y_i * detection_y_gi * M_c,
    )

    mu = numpyro.deterministic("mu", (mu_expression + additive_term) * normalizing_term)

    # ===================== DATA likelihood ======================= #

    concentration = stochastic_v_ag * M_c
    rate = concentration / mu

    # Likelihood (sampling distribution)
    # We apply a scale factor so that the negative log-likelihood is multiplied
    # by (total_num_cells / batch_size). That ensures an unbiased gradient estimate.
    scale_factor = total_num_cells / batch_size
    with scale(scale=scale_factor):
        data_target = numpyro.sample(
            "data_target", dist.GammaPoisson(concentration, rate), obs=data
        )
