data-beans 0.11.127

use indicatif::ParallelProgressIterator;
use legume_numeric::matrix::dmatrix_io::*;
use legume_numeric::matrix::traits::*;
use log::info;
use rand::SeedableRng;
use rand_distr::{Distribution, Normal, Poisson, Uniform};
use rayon::prelude::*;

/// Sample topic proportions `[K, N]` with hard assignments softened by PVE.
///
/// Each cell is hard-assigned to a topic `k*`, then mixed with a uniform
/// background:
///   `θ[k*] = pve_topic + (1 - pve_topic)/K`,
///   `θ[k]  =             (1 - pve_topic)/K`   for `k ≠ k*`.
/// Columns sum to `1` for every value of `pve_topic ∈ [0, 1]`. At
/// `pve_topic = 0` every column is uniform `1/K` (no topic structure);
/// at `pve_topic = 1` every column is one-hot.
pub(crate) fn sample_theta_kn(
    kk: usize,
    nn: usize,
    pve_topic: f32,
    rng: &mut impl rand::Rng,
) -> anyhow::Result<DMatrix<f32>> {
    let runif = Uniform::new(0, kk)?;
    let k_membership: Vec<usize> = (0..nn).map(|_| runif.sample(rng)).collect();
    // One-hot over all `kk` topics, including any that no cell drew.
    let mut theta_kn = DMatrix::<f32>::zeros(kk, nn);
    for (j, &k) in k_membership.iter().enumerate() {
        theta_kn[(k, j)] = 1.0;
    }

    let pve_topic = pve_topic.clamp(0., 1.);
    if kk > 1 && pve_topic < 1. {
        let p_background = (1.0 - pve_topic) / kk as f32;
        let theta_null = DMatrix::<f32>::from_element(kk, nn, p_background);
        theta_kn = (theta_kn * pve_topic) + theta_null;
    }

    Ok(theta_kn)
}

/// Sample a log-normal dictionary `β[D, K]` with explicit topic-PVE
/// decomposition:
///
/// ```text
/// log β_{g,k} = σ_β · [√π_topic · u_{g,k} + √(1−π_topic) · v_g] − σ_β² / 2
/// ```
///
/// with `u_{g,k}, v_g ~ N(0, 1)` iid. By construction:
///   - `Var(log β_{g,k}) = σ_β²` independent of `π_topic`,
///   - `E[β_{g,k}] = 1` (multiplicative-perturbation interpretation, so a
///     downstream `(depth/G)·β·θ` makes `depth` an emergent library size),
///   - `π_topic = 0` ⇒ β purely per-gene (no topic structure),
///   - `π_topic = 1` ⇒ β purely per-(gene, topic) (no shared baseline).
pub(crate) fn sample_lognormal_dictionary(
    dd: usize,
    kk: usize,
    pve_topic: f32,
    beta_scale: f32,
    rng: &mut impl rand::Rng,
) -> DMatrix<f32> {
    let pve = pve_topic.clamp(0.0, 1.0);
    let a_topic = pve.sqrt();
    let a_invariant = (1.0 - pve).sqrt();
    let center = 0.5 * beta_scale * beta_scale;
    let normal = Normal::new(0.0_f32, 1.0_f32).expect("standard normal");

    let v: Vec<f32> = (0..dd).map(|_| normal.sample(rng)).collect();

    let mut beta = DMatrix::<f32>::zeros(dd, kk);
    for g in 0..dd {
        for k in 0..kk {
            let u: f32 = normal.sample(rng);
            let log_b = beta_scale * (a_topic * u + a_invariant * v[g]) - center;
            beta[(g, k)] = log_b.exp();
        }
    }
    beta
}

/// Sample batch-effect log-shifts `log δ ∈ ℝ^{D × B}` with explicit
/// log-space variance decomposition:
///
/// ```text
/// log δ_{g, b} = √π_batch · z_{g, b} + √(1 − π_batch) · w_g
/// ```
///
/// `z` iid N(0, 1) per (gene, batch), z-scored per column to unit variance,
/// scaled by √π_batch. `w` iid N(0, 1) per gene, z-scored to unit variance,
/// scaled by √(1 − π_batch). By construction `Var(log δ_{g, b}) = 1`
/// independent of `π_batch`. Caller exponentiates if mean-space δ is needed.
pub(crate) fn sample_log_batch_effects(
    dd: usize,
    bb: usize,
    pve_batch: f32,
    rng: &mut impl rand::Rng,
) -> DMatrix<f32> {
    let pve = pve_batch.clamp(0.0, 1.0);
    let normal = Normal::new(0.0_f32, 1.0_f32).expect("standard normal");

    let mut ln_delta_db = DMatrix::from_fn(dd, bb, |_, _| normal.sample(rng));
    ln_delta_db.scale_columns_inplace();
    ln_delta_db *= pve.sqrt();

    let mut ln_null_d = DMatrix::from_fn(dd, 1, |_, _| normal.sample(rng));
    ln_null_d.scale_columns_inplace();
    ln_null_d *= (1.0 - pve).sqrt();

    for col in 0..ln_delta_db.ncols() {
        let mut col_mut = ln_delta_db.column_mut(col);
        col_mut += &ln_null_d.column(0);
    }

    ln_delta_db
}

/// Inject housekeeping genes (first `n_housekeeping` rows) with high,
/// approximately-uniform-across-topics expression. Each housekeeping gene
/// `g` gets a single log-normal draw `β_g ~ LN(log(fold·median(β)), σ_hk²)`
/// shared across all `K` topics — the per-gene baseline is correlated, the
/// per-topic noise is gone (these are by-design topic-invariant genes).
pub(crate) fn inject_housekeeping(
    beta_dk: &mut DMatrix<f32>,
    n_housekeeping: usize,
    fold: f32,
    rng: &mut impl rand::Rng,
) {
    if n_housekeeping == 0 || n_housekeeping >= beta_dk.nrows() {
        return;
    }
    let kk = beta_dk.ncols();
    // Pivot: per-gene mean (across topics) — robust enough for a baseline.
    let mean_val = beta_dk.mean().max(1e-30);
    let hk_mean = mean_val * fold;
    const HK_SIGMA: f32 = 0.3;
    let log_mean = hk_mean.ln() - 0.5 * HK_SIGMA * HK_SIGMA;
    let normal = Normal::new(0.0_f32, 1.0_f32).expect("standard normal");
    for g in 0..n_housekeeping {
        let u: f32 = normal.sample(rng);
        let val = (log_mean + HK_SIGMA * u).exp();
        for k in 0..kk {
            beta_dk[(g, k)] = val;
        }
    }
    info!(
        "injected {} housekeeping genes (LN(log {:.4}, {:.2}), fold={:.1}× mean {:.4})",
        n_housekeeping, hk_mean, HK_SIGMA, fold, mean_val
    );
}

/// Sample Poisson count triplets from
/// `Y(g,j) ~ Poisson( λ_scale · δ(g,B(j)) · Σ_k β(g,k) θ(k,j) )`.
///
/// `lambda_scale` should make `E[Σ_g λ_{g,j}] ≈ depth` for the dictionary's
/// normalization (see [`library_scale`]). Library size is **emergent**, not
/// enforced — there is no per-cell rescaling.
pub(crate) fn sample_poisson_triplets(
    beta_dk: &DMatrix<f32>,
    theta_kn: &DMatrix<f32>,
    delta_db: Option<&DMatrix<f32>>,
    batch_membership: &[usize],
    lambda_scale: f32,
    rseed: u64,
    seed_offset: u64,
) -> Vec<(u64, u64, f32)> {
    let nn = theta_kn.ncols();
    let eps = 1e-8_f32;
    let threshold = 0.5_f32;
    let has_batch = delta_db.is_some();
    theta_kn
        .column_iter()
        .enumerate()
        .par_bridge()
        .progress_count(nn as u64)
        .map(|(j, theta_j)| {
            let mut rng = rand::rngs::StdRng::seed_from_u64(
                rseed.wrapping_add(seed_offset).wrapping_add(j as u64),
            );

            let lambda_j = if has_batch {
                let b = batch_membership[j];
                (beta_dk * theta_j).component_mul(&delta_db.unwrap().column(b))
            } else {
                beta_dk * theta_j
            };

            lambda_j
                .iter()
                .enumerate()
                .filter_map(|(i, &l_ij)| {
                    let l_ij = (l_ij * lambda_scale).max(eps);
                    if let Ok(rpois) = Poisson::new(l_ij) {
                        let y_ij: f32 = rpois.sample(&mut rng);
                        if y_ij > threshold {
                            Some((i as u64, j as u64, y_ij))
                        } else {
                            None
                        }
                    } else {
                        None
                    }
                })
                .collect::<Vec<_>>()
        })
        .flatten()
        .collect()
}

/// `λ_scale` that makes `depth` the expected library size of a cell.
///
/// The flat log-normal dictionary has `E[β_{g,k}] = 1`, so a topic's column
/// sums to about `G` and the scale is `depth / G`. The hierarchical dictionary
/// is a partition of one probability simplex — all `G × K` leaf entries sum to
/// 1, a leaf column to about `1/K` — so the scale is `depth · K`.
pub(crate) fn library_scale(depth: usize, dd: usize, kk: usize, hierarchical: bool) -> f32 {
    if hierarchical {
        depth as f32 * kk as f32
    } else {
        depth as f32 / dd as f32
    }
}

/// Stage-1 topic term of the reference-mode samplers: `log(β·θ)_g` z-scored
/// across features, for one cell. Features with no topic loading
/// (`(β·θ)_g = 0`, e.g. an all-zero dictionary row) carry no topic term: they
/// get 0 and are left out of the mean and SD, so they do not set the scale of
/// the features that do vary.
pub(crate) fn standardized_log_topic(bt: &DVector<f32>) -> DVector<f32> {
    let mut t_z = DVector::<f32>::zeros(bt.len());
    let mut n_pos = 0usize;
    let mut sum = 0.0_f32;
    for (t, &x) in t_z.iter_mut().zip(bt.iter()) {
        if x > 0.0 {
            *t = x.max(1e-30).ln();
            sum += *t;
            n_pos += 1;
        }
    }
    if n_pos == 0 {
        return t_z;
    }
    let m_t = sum / n_pos as f32;
    let mut s2 = 0.0_f32;
    for (&t, &x) in t_z.iter().zip(bt.iter()) {
        if x > 0.0 {
            let d = t - m_t;
            s2 += d * d;
        }
    }
    let s_t = (s2 / n_pos as f32).sqrt().max(1e-12);
    for (t, &x) in t_z.iter_mut().zip(bt.iter()) {
        if x > 0.0 {
            *t = (*t - m_t) / s_t;
        }
    }
    t_z
}

pub struct SimArgs {
    pub rows: usize,
    pub cols: usize,
    pub depth: usize,
    pub factors: usize,
    pub batches: usize,
    /// Log-normal scale parameter `σ_β` for the dictionary. Replaces the
    /// old `overdisp` knob. Higher = more variable expression across genes
    /// and topics. Default 1.0.
    pub beta_scale: f32,
    pub pve_topic: f32,
    pub pve_batch: f32,
    pub rseed: u64,
    /// If set, generate a hierarchical gene dictionary using stick-breaking
    /// gates on a binary tree of this depth. K = 2^(depth-1) leaf topics.
    /// Overrides `factors` when set.
    pub hierarchical_depth: Option<usize>,
    /// Number of housekeeping genes with high uniform expression across all
    /// topics. These genes get a single log-normal value broadcast across
    /// all topics, simulating ubiquitous genes (B2M, EEF1A1, ribosomal).
    pub n_housekeeping: usize,
    /// Expression fold-change of housekeeping genes relative to the mean
    /// topic-specific gene. Default: 10.0
    pub housekeeping_fold: f32,
}

pub struct SimOut {
    pub ln_delta_db: DMatrix<f32>,
    pub beta_dk: DMatrix<f32>,
    pub theta_kn: DMatrix<f32>,
    pub batch_membership: Vec<usize>,
    pub triplets: Vec<(u64, u64, f32)>,
    /// If hierarchical: all node probabilities [D, num_nodes] (1-indexed by column)
    pub hierarchy_node_probs: Option<DMatrix<f32>>,
}

/// Generate a simulated dataset with a log-normal factored model.
///
/// ```text
/// log β_{g,k} = σ_β · [√π_topic · u_{g,k} + √(1−π_topic) · v_g] − σ_β² / 2
/// log δ_{g,b} = √π_batch · z_{g,b} + √(1−π_batch) · w_g
/// λ_{g,j}    = λ_scale · δ_{g, B(j)} · Σ_k β_{g,k} θ_{k,j}
/// y_{g,j}    ~ Poisson(λ_{g,j})
/// ```
///
/// `λ_scale` comes from [`library_scale`], so `depth` is the **expected**
/// library size (emergent — no per-cell rescaling). `pve_topic` and `pve_batch` are independent variance shares
/// in log space — both can be 1 simultaneously.
pub fn generate_factored_poisson_gamma_data(args: &SimArgs) -> anyhow::Result<SimOut> {
    let nn = args.cols;
    let dd = args.rows;
    let kk = if let Some(depth) = args.hierarchical_depth {
        1usize << (depth - 1)
    } else {
        args.factors
    };
    let bb = args.batches;
    let rseed = args.rseed;
    let beta_scale = args.beta_scale;
    let pve_topic = args.pve_topic.clamp(0., 1.);
    let pve_batch = args.pve_batch.clamp(0., 1.);

    let mut rng = rand::rngs::StdRng::seed_from_u64(rseed);

    // 1. batch membership matrix
    let runif = Uniform::new(0, bb).expect("unif [0 .. bb)");
    let batch_membership: Vec<usize> = (0..nn).map(|_| runif.sample(&mut rng)).collect();

    let ln_delta_db = sample_log_batch_effects(dd, bb, pve_batch, &mut rng);
    let delta_db = ln_delta_db.map(|x| x.exp());
    info!("simulated batch effects (log-variance decomposition)");

    // 3. dictionary β
    let (mut beta_dk, hierarchy_node_probs) = if let Some(tree_depth) = args.hierarchical_depth {
        let (beta, node_probs) =
            generate_hierarchical_dictionary(dd, tree_depth, beta_scale, &mut rng);
        info!(
            "generated hierarchical dictionary: depth={}, K={} leaves, {} nodes",
            tree_depth,
            beta.ncols(),
            node_probs.ncols()
        );
        (beta, Some(node_probs))
    } else {
        (
            sample_lognormal_dictionary(dd, kk, pve_topic, beta_scale, &mut rng),
            None,
        )
    };

    // 4. housekeeping injection (log-normal, topic-invariant per gene)
    inject_housekeeping(
        &mut beta_dk,
        args.n_housekeeping,
        args.housekeeping_fold,
        &mut rng,
    );

    let theta_kn = sample_theta_kn(kk, nn, pve_topic, &mut rng)?;

    // 5. emergent-library Poisson sampling
    let lambda_scale = library_scale(args.depth, dd, kk, args.hierarchical_depth.is_some());
    let delta_ref = if bb > 1 { Some(&delta_db) } else { None };
    let triplets = sample_poisson_triplets(
        &beta_dk,
        &theta_kn,
        delta_ref,
        &batch_membership,
        lambda_scale,
        rseed,
        0,
    );

    info!(
        "sampled Poisson data with {} non-zero elements (λ_scale = {:.4})",
        triplets.len(),
        lambda_scale,
    );

    Ok(SimOut {
        ln_delta_db,
        beta_dk,
        theta_kn,
        batch_membership,
        triplets,
        hierarchy_node_probs,
    })
}

/// Generate a hierarchical gene dictionary using stick-breaking gates on a
/// binary tree, mirroring the `HierarchicalTopicDecoder` structure.
///
/// * `dd` - number of genes (D)
/// * `tree_depth` - tree depth (>= 2), K = 2^(depth-1) leaf topics
/// * `beta_scale` - log-normal scale `σ_β` for the root distribution; gate
///   logits use a fixed `N(0, 1)` scale (sharper subtrees come from larger
///   `σ_β` driving the root, not from this gate scale).
/// * `rng` - random number generator
///
/// Returns `(beta_dk, node_probs_d_by_numnodes)`:
/// - `beta_dk`: [D, K] leaf dictionary (each row sums to ~1)
/// - `node_probs`: [D, num_nodes] all node probabilities (col 0 = root, etc.)
///
/// Note: the tree's stick-breaking already encodes topic structure by
/// construction, so `pve_topic` does NOT additionally blend β here — only
/// θ blends. (Flat mode applies the explicit log-space topic-PVE blend.)
pub(crate) fn generate_hierarchical_dictionary(
    dd: usize,
    tree_depth: usize,
    beta_scale: f32,
    rng: &mut impl rand::Rng,
) -> (DMatrix<f32>, DMatrix<f32>) {
    assert!(tree_depth >= 2, "Tree depth must be at least 2");

    let num_leaves = 1usize << (tree_depth - 1); // K = 2^(depth-1)
    let num_nodes = (1usize << tree_depth) - 1; // 2^depth - 1
    let num_internal = num_leaves - 1;

    let normal = Normal::new(0.0_f32, 1.0_f32).expect("standard normal");

    // 1. Root distribution: log-normal then normalized to probability simplex
    let center = 0.5 * beta_scale * beta_scale;
    let root_raw: Vec<f32> = (0..dd)
        .map(|_| (beta_scale * normal.sample(rng) - center).exp())
        .collect();
    let root_sum: f32 = root_raw.iter().sum::<f32>().max(1e-30);
    let root_prob: Vec<f32> = root_raw.iter().map(|&x| x / root_sum).collect();

    // 2. Gate logits ~ N(0, 1) → sigmoid gives gates in (0, 1).
    let gate_logits: Vec<Vec<f32>> = (0..num_internal)
        .map(|_| (0..dd).map(|_| normal.sample(rng)).collect())
        .collect();

    // 3. Top-down propagation: compute node probabilities (1-indexed: 0 = unused).
    let mut node_probs: Vec<Vec<f32>> = vec![vec![0.0; dd]; num_nodes + 1];
    node_probs[1] = root_prob;

    for h in 1..num_leaves {
        let gate_idx = h - 1;
        let left_child = 2 * h;
        let right_child = 2 * h + 1;

        let mut left = vec![0.0_f32; dd];
        let mut right = vec![0.0_f32; dd];

        for d in 0..dd {
            let gate = sigmoid(gate_logits[gate_idx][d]);
            left[d] = node_probs[h][d] * gate;
            right[d] = node_probs[h][d] * (1.0 - gate);
        }

        node_probs[left_child] = left;
        node_probs[right_child] = right;
    }

    // 4. Extract leaf dictionary [D, K]
    let mut beta_dk = DMatrix::<f32>::zeros(dd, num_leaves);
    for k in 0..num_leaves {
        let leaf_node = num_leaves + k;
        for d in 0..dd {
            beta_dk[(d, k)] = node_probs[leaf_node][d];
        }
    }

    // 5. Pack all node probabilities into [D, num_nodes]
    let mut all_node_probs = DMatrix::<f32>::zeros(dd, num_nodes);
    for h in 0..num_nodes {
        for d in 0..dd {
            all_node_probs[(d, h)] = node_probs[h + 1][d];
        }
    }

    (beta_dk, all_node_probs)
}

#[inline]
fn sigmoid(x: f32) -> f32 {
    1.0 / (1.0 + (-x).exp())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn sim_args() -> SimArgs {
        SimArgs {
            rows: 400,
            cols: 300,
            depth: 1000,
            factors: 4,
            batches: 1,
            beta_scale: 1.0,
            pve_topic: 1.0,
            pve_batch: 1.0,
            rseed: 11,
            hierarchical_depth: None,
            n_housekeeping: 0,
            housekeeping_fold: 10.0,
        }
    }

    fn mean_library(triplets: &[(u64, u64, f32)], n_cells: usize) -> f32 {
        triplets.iter().map(|&(_, _, v)| v).sum::<f32>() / n_cells as f32
    }

    /// θ keeps one row per topic even when far fewer cells than topics are
    /// drawn, so some topics are never anyone's dominant topic.
    #[test]
    fn theta_has_a_row_for_every_topic() -> anyhow::Result<()> {
        for seed in 0..20u64 {
            for &pve in &[1.0_f32, 0.5] {
                let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
                let theta = sample_theta_kn(8, 3, pve, &mut rng)?;
                assert_eq!(theta.shape(), (8, 3), "seed {seed}, pve {pve}");
                for j in 0..3 {
                    let s: f32 = theta.column(j).sum();
                    assert!((s - 1.0).abs() < 1e-5, "column {j} sums to {s}");
                }
            }
        }
        Ok(())
    }

    /// Few cells, many topics: the last topic is routinely undrawn and the
    /// run still completes with a full-width θ.
    #[test]
    fn few_cells_many_topics_simulates() -> anyhow::Result<()> {
        for seed in 0..10u64 {
            let args = SimArgs {
                rows: 50,
                cols: 4,
                factors: 6,
                pve_topic: 0.7,
                rseed: seed,
                ..sim_args()
            };
            let out = generate_factored_poisson_gamma_data(&args)?;
            assert_eq!(out.theta_kn.shape(), (6, 4));
        }
        Ok(())
    }

    /// `depth` is the expected library size in both dictionary modes.
    #[test]
    fn depth_is_the_expected_library_size() -> anyhow::Result<()> {
        for hierarchical_depth in [None, Some(3)] {
            let args = SimArgs {
                hierarchical_depth,
                ..sim_args()
            };
            let out = generate_factored_poisson_gamma_data(&args)?;
            let lib = mean_library(&out.triplets, args.cols);
            assert!(
                (lib - 1000.0).abs() < 200.0,
                "hierarchical_depth {hierarchical_depth:?}: mean library {lib}"
            );
        }
        Ok(())
    }

    /// A seed at the top of the u64 range wraps instead of overflowing.
    #[test]
    fn per_cell_seed_wraps() {
        let beta = DMatrix::<f32>::from_element(5, 2, 1.0);
        let theta = DMatrix::<f32>::from_element(2, 3, 0.5);
        let trips = sample_poisson_triplets(&beta, &theta, None, &[0, 0, 0], 3.0, u64::MAX, 7);
        assert!(!trips.is_empty());
    }

    /// Features with a zero rate get no topic term and do not enter the
    /// z-score of the features that have one.
    #[test]
    fn log_topic_standardizes_over_loaded_features_only() {
        let loaded = [0.2_f32, 0.5, 1.0, 4.0];
        let mut padded = loaded.to_vec();
        padded.extend([0.0; 12]);
        let dense = standardized_log_topic(&DVector::from_row_slice(&loaded));
        let sparse = standardized_log_topic(&DVector::from_row_slice(&padded));
        for g in 0..loaded.len() {
            assert!((dense[g] - sparse[g]).abs() < 1e-6, "feature {g}");
        }
        assert!(sparse.iter().skip(loaded.len()).all(|&t| t == 0.0));
        let var: f32 = dense.iter().map(|t| t * t).sum::<f32>() / loaded.len() as f32;
        assert!((var - 1.0).abs() < 1e-4, "unit variance, got {var}");
        assert!(standardized_log_topic(&DVector::zeros(3))
            .iter()
            .all(|&t| t == 0.0));
    }
}