causasv 0.7.0

use rand::RngExt;
use rand::SeedableRng;
use rand::rngs::StdRng;

use crate::graph::{Dag, NodeId};

/// Configuration for approximate ASV sampling.
pub struct SamplingConfig {
    pub n_samples: usize,
    pub seed: Option<u64>,
    /// When set, enables batched value-function evaluation: collect this many samples,
    /// deduplicate coalitions, call `value_fn_batch` once, then process IS weights.
    /// Reduces Python GIL acquisition overhead for large models.
    pub batch_size: Option<usize>,
    /// When true, use Rayon parallel sampling. Seeded + parallel → per-worker seeds via
    /// splitmix64 bijection so results are deterministic regardless of thread count.
    pub parallel: bool,
    /// Number of worker threads for seeded parallel sampling. None = rayon default.
    pub num_threads: Option<usize>,
}

impl SamplingConfig {
    pub fn new(n_samples: usize) -> Self {
        Self {
            n_samples,
            seed: None,
            batch_size: None,
            parallel: false,
            num_threads: None,
        }
    }

    pub fn with_seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }

    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = Some(batch_size);
        self
    }

    pub fn with_parallel(mut self, parallel: bool) -> Self {
        self.parallel = parallel;
        self
    }

    pub fn with_num_threads(mut self, num_threads: usize) -> Self {
        self.num_threads = Some(num_threads);
        self
    }
}

/// splitmix64 bijection: maps (global_seed, worker_index) → a deterministic per-worker seed.
///
/// Each worker gets a distinct seed that is a deterministic function of the global seed and
/// its index, so `seed=42, parallel=true` always produces the same aggregate result.
pub(crate) fn worker_seed(global: u64, k: usize) -> u64 {
    let mut x = global
        .wrapping_add(k as u64)
        .wrapping_mul(0x9e3779b97f4a7c15);
    x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
    x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb);
    x ^ (x >> 31)
}

/// Configuration for adaptive approximate ASV sampling.
///
/// Sampling runs in batches and stops when the per-node ASV estimates change by
/// less than `rel_tol` relative to their current values *and* the ESS ratio
/// exceeds `ess_ratio_min`. Falls back to `max_samples` if convergence is never
/// reached.
pub struct AdaptiveSamplingConfig {
    /// Minimum samples before convergence is checked. Default: 1 000.
    pub min_samples: usize,
    /// Hard upper bound on total samples. Default: 100 000.
    pub max_samples: usize,
    /// Samples per batch. Default: 1 000.
    pub batch_size: usize,
    /// Relative tolerance on per-node value change between batches. Default: 0.01.
    pub rel_tol: f64,
    /// Minimum required ESS / n_samples ratio. Default: 0.10.
    pub ess_ratio_min: f64,
    pub seed: Option<u64>,
}

impl Default for AdaptiveSamplingConfig {
    fn default() -> Self {
        Self {
            min_samples: 1_000,
            max_samples: 100_000,
            batch_size: 1_000,
            rel_tol: 0.01,
            ess_ratio_min: 0.10,
            seed: None,
        }
    }
}

impl AdaptiveSamplingConfig {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }

    pub fn with_rel_tol(mut self, tol: f64) -> Self {
        self.rel_tol = tol;
        self
    }

    pub fn with_max_samples(mut self, max: usize) -> Self {
        self.max_samples = max;
        self
    }
}

pub(crate) struct SampledOrdering {
    pub ordering: Vec<NodeId>,
    /// log q(π) under the frontier sampler.
    /// Used by the IS estimator in approx.rs — do not remove.
    pub log_q: f64,
}

pub(crate) fn sample_one(dag: &Dag, rng: &mut StdRng) -> SampledOrdering {
    let n = dag.node_count();
    let mut in_deg = dag.in_degrees();
    let mut ordering = Vec::with_capacity(n);
    let mut frontier = Vec::with_capacity(n);
    let mut log_q = 0.0f64;

    for _ in 0..n {
        // Frontier: nodes with in_deg == 0 (placed nodes have in_deg = MAX)
        frontier.clear();
        frontier.extend((0..n).filter(|&i| in_deg[i] == 0));
        log_q -= (frontier.len() as f64).ln();
        let idx = rng.random_range(0..frontier.len());
        let node = NodeId(frontier[idx] as u32);
        ordering.push(node);
        in_deg[node.0 as usize] = usize::MAX;
        for &child in dag.children_raw(node) {
            in_deg[child.0 as usize] -= 1;
        }
    }

    SampledOrdering { ordering, log_q }
}

pub(crate) fn make_rng(seed: Option<u64>) -> StdRng {
    match seed {
        Some(s) => StdRng::seed_from_u64(s),
        None => StdRng::from_rng(&mut rand::rng()),
    }
}

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

    fn chain_2() -> Dag {
        let mut dag = Dag::new();
        let a = dag.add_node("a");
        let b = dag.add_node("b");
        dag.add_edge(a, b).unwrap();
        dag
    }

    /// Chain a→b: each step has exactly one frontier node → log_q = -ln(1) - ln(1) = 0.
    #[test]
    fn test_chain_log_q_is_zero() {
        let dag = chain_2();
        let s = sample_one(&dag, &mut make_rng(Some(0)));
        assert!((s.log_q - 0.0).abs() < 1e-12);
    }

    /// Two independent nodes: first pick from {a,b} (prob 0.5), second forced → log_q = -ln(2).
    #[test]
    fn test_two_independent_log_q() {
        let mut dag = Dag::new();
        dag.add_node("a");
        dag.add_node("b");
        let s = sample_one(&dag, &mut make_rng(Some(0)));
        let expected = -(2.0f64.ln());
        assert!((s.log_q - expected).abs() < 1e-12);
    }

    /// For a chain, every sample must have the parent before its child.
    #[test]
    fn test_chain_ordering_always_valid() {
        let dag = chain_2();
        let mut rng = make_rng(Some(1));
        for _ in 0..200 {
            let s = sample_one(&dag, &mut rng);
            assert_eq!(s.ordering.len(), 2);
            let pos: Vec<usize> = s.ordering.iter().map(|id| id.0 as usize).collect();
            assert_eq!(pos[0], 0, "a (NodeId 0) must come first in chain");
        }
    }

    /// Every sample must contain all n nodes exactly once.
    #[test]
    fn test_no_missing_no_duplicate_nodes() {
        let mut dag = Dag::new();
        for i in 0..5usize {
            dag.add_node(&format!("n{i}"));
        }
        let mut rng = make_rng(Some(2));
        for _ in 0..100 {
            let s = sample_one(&dag, &mut rng);
            let mut sorted = s.ordering.clone();
            sorted.sort_unstable();
            let expected: Vec<NodeId> = (0..5).map(|i| NodeId(i as u32)).collect();
            assert_eq!(sorted, expected);
        }
    }

    /// Same seed must produce identical ordering and log_q.
    #[test]
    fn test_same_seed_same_result() {
        let dag = chain_2();
        let s1 = sample_one(&dag, &mut make_rng(Some(42)));
        let s2 = sample_one(&dag, &mut make_rng(Some(42)));
        assert_eq!(s1.ordering, s2.ordering);
        assert_eq!(s1.log_q, s2.log_q);
    }
}