causal-hub 0.0.6

use core::f64;
use std::cell::RefCell;

use ndarray::prelude::*;
use ndarray_stats::QuantileExt;
use rand::{
    Rng, RngExt, SeedableRng,
    distr::{Distribution, weighted::WeightedIndex},
};
use rand_distr::Exp;
use rayon::prelude::*;

use crate::{
    datasets::{CatSample, CatTable, CatTrj, CatType, GaussTable},
    models::{
        BN, CIM, CPD, CTBN, CatBN, CatCTBN, GaussBN, HasLabels, MixedBN, MixedCPD, MixedSample,
        MixedSupport, MixedTable,
    },
    samplers::{BNSampler, CTBNSampler, ParBNSampler, ParCTBNSampler},
    set,
    types::{EPSILON, Error, Map, Result, Set},
};

/// A forward sampler.
#[derive(Debug)]
pub struct ForwardSampler<'a, R, M> {
    rng: RefCell<&'a mut R>,
    model: &'a M,
}

impl<'a, R, M> ForwardSampler<'a, R, M> {
    /// Construct a new forward sampler.
    ///
    /// # Arguments
    ///
    /// * `rng` - A random number generator.
    /// * `model` - A reference to the model to sample from.
    ///
    /// # Returns
    ///
    /// A new `ForwardSampler` instance.
    ///
    #[inline]
    pub fn new(rng: &'a mut R, model: &'a M) -> Result<Self> {
        // Wrap the RNG in a RefCell to allow interior mutability.
        let rng = RefCell::new(rng);

        Ok(Self { rng, model })
    }
}

impl<R: Rng> BNSampler<CatBN> for ForwardSampler<'_, R, CatBN> {
    type Sample = <CatBN as BN>::Sample;
    type Samples = <CatBN as BN>::Samples;

    fn sample(&self) -> Result<Self::Sample> {
        // Get a mutable reference to the RNG.
        let mut rng = self.rng.borrow_mut();
        // Allocate the sample.
        let mut sample = Array::zeros(self.model.labels().len());

        // For each vertex in the topological order ...
        for &i in self.model.topological_order() {
            // Get the CPD.
            let cpd_i = &self.model.cpds()[i];
            // Compute the index on the parents to condition on.
            let pa_i = self.model.graph().parents(&set![i])?;
            let pa_i = pa_i.iter().map(|&z| sample[z]).collect();
            // Sample from the distribution.
            sample[i] = cpd_i.sample(&mut rng, &pa_i)?[0];
        }

        Ok(sample)
    }

    fn sample_n(&self, n: usize) -> Result<Self::Samples> {
        // Allocate the dataset.
        let mut dataset = Array::zeros((n, self.model.labels().len()));

        // For each sample, sample from the distribution using iterators.
        dataset
            .rows_mut()
            .into_iter()
            .try_for_each(|mut row| -> Result<_> {
                // Sample from the distribution.
                row.assign(&self.sample()?);

                Ok(())
            })?;

        // Construct the dataset.
        CatTable::new(self.model.support().clone(), dataset)
    }
}

impl<R: Rng + SeedableRng> ParBNSampler<CatBN> for ForwardSampler<'_, R, CatBN> {
    type Samples = <CatBN as BN>::Samples;

    fn par_sample_n(&self, n: usize) -> Result<Self::Samples> {
        // Get a mutable reference to the RNG.
        let rng = self.rng.borrow_mut();
        // Generate a random seed for each sample.
        let seeds: Vec<_> = rng.random_iter().take(n).collect();

        // Allocate the samples.
        let mut samples = Array::zeros((n, self.model.labels().len()));

        // Sample the samples in parallel.
        seeds
            .into_par_iter()
            .zip(samples.axis_iter_mut(Axis(0)))
            .try_for_each(|(seed, mut row)| -> Result<()> {
                // Create a new random number generator with the seed.
                let mut rng = R::seed_from_u64(seed);
                // Create a new sampler with the random number generator and model.
                let sampler = ForwardSampler::new(&mut rng, self.model)?;
                // Sample from the distribution.
                row.assign(&sampler.sample()?);

                Ok(())
            })?;

        // Construct the dataset.
        CatTable::new(self.model.support().clone(), samples)
    }
}

impl<R: Rng> BNSampler<GaussBN> for ForwardSampler<'_, R, GaussBN> {
    type Sample = <GaussBN as BN>::Sample;
    type Samples = <GaussBN as BN>::Samples;

    fn sample(&self) -> Result<Self::Sample> {
        // Get a mutable reference to the RNG.
        let mut rng = self.rng.borrow_mut();
        // Allocate the sample.
        let mut sample = Array::zeros(self.model.labels().len());

        // For each vertex in the topological order ...
        for &i in self.model.topological_order() {
            // Get the CPD.
            let cpd_i = &self.model.cpds()[i];
            // Get the parents.
            let pa_i = self.model.graph().parents(&set![i])?;
            let pa_i = pa_i.iter().map(|&z| sample[z]).collect();
            // Compute the value of the variable.
            sample[i] = cpd_i.sample(&mut rng, &pa_i)?[0];
        }

        Ok(sample)
    }

    fn sample_n(&self, n: usize) -> Result<Self::Samples> {
        // Allocate the samples.
        let mut samples = Array::zeros((n, self.model.labels().len()));

        // For each sample, sample from the distribution using iterators.
        samples
            .rows_mut()
            .into_iter()
            .try_for_each(|mut row| -> Result<_> {
                // Sample from the distribution.
                row.assign(&self.sample()?);

                Ok(())
            })?;

        // Construct the dataset.
        GaussTable::new(self.model.labels().clone(), samples)
    }
}

impl<R: Rng + SeedableRng> ParBNSampler<GaussBN> for ForwardSampler<'_, R, GaussBN> {
    type Samples = <GaussBN as BN>::Samples;

    fn par_sample_n(&self, n: usize) -> Result<Self::Samples> {
        // Get a mutable reference to the RNG.
        let rng = self.rng.borrow_mut();
        // Generate a random seed for each sample.
        let seeds: Vec<_> = rng.random_iter().take(n).collect();

        // Allocate the samples.
        let mut samples = Array::zeros((n, self.model.labels().len()));

        // Sample the samples in parallel.
        seeds
            .into_par_iter()
            .zip(samples.axis_iter_mut(Axis(0)))
            .try_for_each(|(seed, mut row)| -> Result<()> {
                // Create a new random number generator with the seed.
                let mut rng = R::seed_from_u64(seed);
                // Create a new sampler with the random number generator and model.
                let sampler = ForwardSampler::new(&mut rng, self.model)?;
                // Sample from the distribution.
                row.assign(&sampler.sample()?);
                Ok(())
            })?;

        // Construct the dataset.
        GaussTable::new(self.model.labels().clone(), samples)
    }
}

// ── MixedBN sampler helpers & macro ──────────────────────────────────

macro_rules! mixed_pure_dispatch {
    ($self:expr, $cat:block, $gauss:block) => {{
        let ok = $self
            .model
            .cpds()
            .values()
            .all(|distribution| matches!(distribution, MixedCPD::Categorical(_)));
        if ok {
            $cat
        } else {
            let ok = $self
                .model
                .cpds()
                .values()
                .all(|distribution| matches!(distribution, MixedCPD::Gaussian(_)));
            if ok {
                $gauss
            } else {
                return Err(Error::InvalidParameter(
                    "model",
                    "mixed CPD types not yet supported for forward sampling",
                ));
            }
        }
    }};
}

fn extract_cat_support(
    support: &crate::types::Map<String, MixedSupport>,
) -> Map<String, Set<String>> {
    support
        .iter()
        .map(|(label, mixed_supp)| match mixed_supp {
            MixedSupport::Categorical(cpd_support) => {
                let states = cpd_support[label].clone();
                (label.clone(), states)
            }
            _ => unreachable!(), // A mixed-BN's support is always categorical.
        })
        .collect()
}

// ── MixedBN sampler ──────────────────────────────────────────────────

impl<R: Rng> BNSampler<MixedBN> for ForwardSampler<'_, R, MixedBN> {
    type Sample = MixedSample;
    type Samples = MixedTable;

    fn sample(&self) -> Result<Self::Sample> {
        let mut rng = self.rng.borrow_mut();
        let n_vars = self.model.labels().len();

        mixed_pure_dispatch!(
            self,
            {
                let mut sample = Array1::<CatType>::zeros(n_vars);
                for &i in self.model.topological_order() {
                    let cpd_i = &self.model.cpds()[i];
                    let pa_i = self.model.graph().parents(&set![i])?;
                    let pa_vals: Array1<CatType> = pa_i.iter().map(|&z| sample[z]).collect();
                    let pa_sample = MixedSample::Categorical(pa_vals);
                    let result = cpd_i.sample(&mut rng, &pa_sample)?;
                    if let MixedSample::Categorical(v) = result {
                        sample[i] = v[0];
                    }
                }
                Ok(MixedSample::Categorical(sample))
            },
            {
                let mut sample = Array1::<f64>::zeros(n_vars);
                for &i in self.model.topological_order() {
                    let cpd_i = &self.model.cpds()[i];
                    let pa_i = self.model.graph().parents(&set![i])?;
                    let pa_vals: Array1<f64> = pa_i.iter().map(|&z| sample[z]).collect();
                    let pa_sample = MixedSample::Gaussian(pa_vals);
                    let result = cpd_i.sample(&mut rng, &pa_sample)?;
                    if let MixedSample::Gaussian(v) = result {
                        sample[i] = v[0];
                    }
                }
                Ok(MixedSample::Gaussian(sample))
            }
        )
    }

    fn sample_n(&self, n: usize) -> Result<Self::Samples> {
        mixed_pure_dispatch!(
            self,
            {
                let mut dataset = Array::zeros((n, self.model.labels().len()));
                dataset
                    .rows_mut()
                    .into_iter()
                    .try_for_each(|mut row| -> Result<_> {
                        if let MixedSample::Categorical(stats) = self.sample()? {
                            row.assign(&stats);
                        }
                        Ok(())
                    })?;
                let cat_support = extract_cat_support(&self.model.support());
                CatTable::new(cat_support, dataset).map(MixedTable::Categorical)
            },
            {
                let mut dataset = Array::zeros((n, self.model.labels().len()));
                dataset
                    .rows_mut()
                    .into_iter()
                    .try_for_each(|mut row| -> Result<_> {
                        if let MixedSample::Gaussian(stats) = self.sample()? {
                            row.assign(&stats);
                        }
                        Ok(())
                    })?;
                GaussTable::new(self.model.labels().clone(), dataset).map(MixedTable::Gaussian)
            }
        )
    }
}

impl<R: Rng + SeedableRng> ParBNSampler<MixedBN> for ForwardSampler<'_, R, MixedBN> {
    type Samples = MixedTable;

    fn par_sample_n(&self, n: usize) -> Result<Self::Samples> {
        let rng = self.rng.borrow_mut();
        let seeds: Vec<_> = rng.random_iter().take(n).collect();

        mixed_pure_dispatch!(
            self,
            {
                let mut samples = Array::zeros((n, self.model.labels().len()));
                seeds
                    .into_par_iter()
                    .zip(samples.axis_iter_mut(Axis(0)))
                    .try_for_each(|(seed, mut row)| -> Result<()> {
                        let mut rng = R::seed_from_u64(seed);
                        let sampler = ForwardSampler::new(&mut rng, self.model)?;
                        if let MixedSample::Categorical(stats) = sampler.sample()? {
                            row.assign(&stats);
                        }
                        Ok(())
                    })?;
                let cat_support = extract_cat_support(&self.model.support());
                CatTable::new(cat_support, samples).map(MixedTable::Categorical)
            },
            {
                let mut samples = Array::zeros((n, self.model.labels().len()));
                seeds
                    .into_par_iter()
                    .zip(samples.axis_iter_mut(Axis(0)))
                    .try_for_each(|(seed, mut row)| -> Result<()> {
                        let mut rng = R::seed_from_u64(seed);
                        let sampler = ForwardSampler::new(&mut rng, self.model)?;
                        if let MixedSample::Gaussian(stats) = sampler.sample()? {
                            row.assign(&stats);
                        }
                        Ok(())
                    })?;
                GaussTable::new(self.model.labels().clone(), samples).map(MixedTable::Gaussian)
            }
        )
    }
}

impl<R: Rng> ForwardSampler<'_, R, CatCTBN> {
    /// Sample transition time for variable `X_i` with state `x_i`.
    fn sample_time(&self, event: &CatSample, i: usize) -> Result<f64> {
        // Cast the state to usize.
        let x = event[i] as usize;
        // Get the CIM.
        let cim_i = &self.model.cims()[i];
        // Compute the index on the parents to condition on.
        let pa_i = self.model.graph().parents(&set![i])?;
        let pa_i = pa_i.iter().map(|&z| event[z] as usize);
        let pa_i = cim_i.conditioning_multi_index().ravel(pa_i);
        // Get the distribution of the vertex.
        let q_i_x = -cim_i.parameters()[[pa_i, x, x]];
        // Initialize the exponential distribution.
        let exp_i_x = Exp::new(q_i_x)
            .map_err(|evidence| Error::RandDistr(&format!("Invalid lambda: {}", evidence)))?;
        // Sample the transition time.
        Ok(exp_i_x.sample(&mut self.rng.borrow_mut()))
    }
}

impl<R: Rng> CTBNSampler<CatCTBN> for ForwardSampler<'_, R, CatCTBN> {
    type Sample = <CatCTBN as CTBN>::Trajectory;
    type Samples = <CatCTBN as CTBN>::Trajectories;

    #[inline]
    fn sample_by_length(&self, max_length: usize) -> Result<Self::Sample> {
        // Delegate to generic function.
        self.sample_by_length_or_time(max_length, f64::MAX)
    }

    #[inline]
    fn sample_by_time(&self, max_time: f64) -> Result<Self::Sample> {
        // Delegate to generic function.
        self.sample_by_length_or_time(usize::MAX, max_time)
    }

    fn sample_by_length_or_time(&self, max_length: usize, max_time: f64) -> Result<Self::Sample> {
        // Check length is positive.
        if max_length == 0 {
            return Err(Error::InvalidParameter(
                "max_length",
                "The maximum length of the trajectory must be strictly positive.",
            ));
        }
        // Check time is positive.
        if max_time <= 0. {
            return Err(Error::InvalidParameter("max_time", "must be positive"));
        }

        // Allocate the trajectory components.
        let mut sample_events = Vec::new();
        let mut sample_times = Vec::new();

        // Sample the initial support.
        let mut event = {
            let mut rng = self.rng.borrow_mut();
            let initial = self.model.initial_distribution();
            let initial = ForwardSampler::new(&mut rng, initial)?;
            initial.sample()?
        };
        // Append the initial state to the trajectory.
        sample_events.push(event.clone());
        sample_times.push(0.);

        // Sample the transition time.
        let mut times: Array1<_> = (0..event.len())
            .map(|i| self.sample_time(&event, i))
            .collect::<Result<_>>()?;

        // Get the variable that transitions first.
        let mut i = times
            .argmin()
            .map_err(|evidence| Error::Stats(&format!("Failed to find min time: {}", evidence)))?;
        // Set global time.
        let mut time = times[i];

        // While:
        //  1. the length of the trajectory is less than max_length, and ...
        //  2. the time is less than max_time ...
        while sample_events.len() < max_length && time < max_time {
            // Cast the state to usize.
            let x = event[i] as usize;
            // Get the CIM.
            let cim_i = &self.model.cims()[i];
            // Compute the index on the parents to condition on.
            let pa_i = self.model.graph().parents(&set![i])?;
            let pa_i = pa_i.iter().map(|&z| event[z] as usize);
            let pa_i = cim_i.conditioning_multi_index().ravel(pa_i);
            // Get the distribution of the vertex.
            let mut q_i_zx = cim_i.parameters().slice(s![pa_i, x, ..]).to_owned();
            // Set the diagonal element to zero.
            q_i_zx[x] = 0.;
            // Normalize the probabilities.
            q_i_zx /= q_i_zx.sum();
            // Initialize a weighted index sampler.
            let s_i_zx = WeightedIndex::new(&q_i_zx).map_err(|evidence| {
                Error::RandDistr(&format!("Invalid probabilities: {}", evidence))
            })?;
            // Sample the next event.
            event[i] = s_i_zx.sample(&mut self.rng.borrow_mut()) as CatType;
            // Append the event to the trajectory.
            sample_events.push(event.clone());
            sample_times.push(time);
            // Update the transition times for { X } U Ch(X).
            for j in std::iter::once(i).chain(self.model.graph().children(&set![i])?) {
                // Sample the transition time.
                times[j] = time + self.sample_time(&event, j)?;
            }
            // Add a small epsilon to avoid zero transition times.
            times += EPSILON;
            // Get the variable to transition first.
            i = times.argmin().map_err(|evidence| {
                Error::Stats(&format!("Failed to find min time: {}", evidence))
            })?;
            // Update the global time.
            time = times[i];
        }

        // Get the support of the CIMs.
        let support = self.model.support().clone();

        // Convert the events to a 2D array.
        let shape = (sample_events.len(), sample_events[0].len());
        let sample_events = Array::from_iter(sample_events.into_iter().flatten())
            .into_shape_with_order(shape)
            .map_err(|evidence| Error::Shape(&evidence.to_string()))?;
        // Convert the times to a 1D array.
        let sample_times = Array::from_iter(sample_times);

        // Return the trajectory.
        CatTrj::new(support, sample_events, sample_times)
    }

    #[inline]
    fn sample_n_by_length(&self, max_length: usize, n: usize) -> Result<Self::Samples> {
        (0..n).map(|_| self.sample_by_length(max_length)).collect()
    }

    #[inline]
    fn sample_n_by_time(&self, max_time: f64, n: usize) -> Result<Self::Samples> {
        (0..n).map(|_| self.sample_by_time(max_time)).collect()
    }

    #[inline]
    fn sample_n_by_length_or_time(
        &self,
        max_length: usize,
        max_time: f64,
        n: usize,
    ) -> Result<Self::Samples> {
        (0..n)
            .map(|_| self.sample_by_length_or_time(max_length, max_time))
            .collect()
    }
}

impl<R: Rng + SeedableRng> ParCTBNSampler<CatCTBN> for ForwardSampler<'_, R, CatCTBN> {
    type Samples = <CatCTBN as CTBN>::Trajectories;

    #[inline]
    fn par_sample_n_by_length(&self, max_length: usize, n: usize) -> Result<Self::Samples> {
        self.par_sample_n_by_length_or_time(max_length, f64::MAX, n)
    }

    #[inline]
    fn par_sample_n_by_time(&self, max_time: f64, n: usize) -> Result<Self::Samples> {
        self.par_sample_n_by_length_or_time(usize::MAX, max_time, n)
    }

    fn par_sample_n_by_length_or_time(
        &self,
        max_length: usize,
        max_time: f64,
        n: usize,
    ) -> Result<Self::Samples> {
        // Get a mutable reference to the RNG.
        let rng = self.rng.borrow_mut();
        // Generate a random seed for each trajectory.
        let seeds: Vec<_> = rng.random_iter().take(n).collect();
        // Sample the trajectories in parallel.
        seeds
            .into_par_iter()
            .map(|seed| {
                // Create a new random number generator with the seed.
                let mut rng = R::seed_from_u64(seed);
                // Create a new sampler with the random number generator and model.
                let sampler = ForwardSampler::new(&mut rng, self.model)?;
                // Sample the trajectory.
                sampler.sample_by_length_or_time(max_length, max_time)
            })
            .collect()
    }
}