causal-hub 0.0.6

use serde::{Deserialize, Serialize};

use crate::{
    datasets::GaussType,
    models::HasLabels,
    types::{Error, Labels, Result, Set},
};

/// Gaussian evidence type.
#[non_exhaustive]
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum GaussEvT {
    /// Certain positive evidence.
    CertainPositive {
        /// The observed event of the evidence.
        event: usize,
        /// The value of the evidence.
        value: GaussType,
    },
}

impl GaussEvT {
    /// Return the observed event of the evidence.
    ///
    /// # Returns
    ///
    /// The observed event of the evidence.
    ///
    pub const fn event(&self) -> usize {
        match self {
            Self::CertainPositive { event, .. } => *event,
        }
    }

    /// Set the observed event of the evidence.
    ///
    /// # Arguments
    ///
    /// * `event` - The new observed event of the evidence.
    ///
    pub const fn set_event(&mut self, event: usize) {
        match self {
            Self::CertainPositive {
                event: evidence, ..
            } => *evidence = event,
        }
    }
}

/// Gaussian evidence structure.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct GaussEv {
    labels: Labels,
    evidences: Vec<Option<GaussEvT>>,
}

impl HasLabels for GaussEv {
    #[inline]
    fn labels(&self) -> &Labels {
        &self.labels
    }
}

impl GaussEv {
    /// Create a new Gaussian evidence structure.
    ///
    /// # Arguments
    ///
    /// * `labels` - The labels of the evidence structure.
    /// * `values` - An iterator over the evidence values.
    ///
    /// # Returns
    ///
    /// A new Gaussian evidence structure.
    ///
    pub fn new<I>(mut labels: Labels, values: I) -> Result<Self>
    where
        I: IntoIterator<Item = GaussEvT>,
    {
        // Get shortened variable type.
        use GaussEvT as E;

        // Fill the evidences.
        let mut evidences = values.into_iter().try_fold(
            vec![None; labels.len()],
            |mut evidences, evidence| -> Result<_> {
                // Get the event of the evidence.
                let event = evidence.event();
                // Check if event is in bounds.
                if event >= evidences.len() {
                    return Err(Error::IndexOutOfBounds(event));
                }
                // Push the value into the variable events.
                evidences[event] = Some(evidence);

                Ok(evidences)
            },
        )?;

        // Sort labels, if necessary.
        if !labels.is_sorted() {
            // Clone the labels.
            let mut new_labels = labels.clone();
            // Sort the labels.
            new_labels.sort();

            // Sort the evidences.
            let new_evidences = evidences.into_iter().flatten().try_fold(
                vec![None; new_labels.len()],
                |mut new_evidences, evidence| -> Result<_> {
                    // Get the event of the evidence.
                    let event_name = labels
                        .get_index(evidence.event())
                        .ok_or_else(|| Error::IndexOutOfBounds(evidence.event()))?;
                    // Sort the event index.
                    let event = new_labels
                        .get_index_of(event_name)
                        .ok_or_else(|| Error::MissingLabel(event_name))?;

                    // Sort the variable events.
                    let evidence = match evidence {
                        E::CertainPositive { value, .. } => E::CertainPositive { event, value },
                    };

                    // Push the value into the variable events.
                    new_evidences[event] = Some(evidence);

                    Ok(new_evidences)
                },
            )?;

            // Update the labels.
            labels = new_labels;
            // Update the evidences.
            evidences = new_evidences;
        }

        Ok(Self { labels, evidences })
    }

    /// The evidences of the evidence.
    ///
    /// # Returns
    ///
    /// A reference to the evidences of the evidence.
    ///
    #[inline]
    pub const fn evidences(&self) -> &Vec<Option<GaussEvT>> {
        &self.evidences
    }

    /// Restrict the evidence to the specified variables.
    ///
    /// # Arguments
    ///
    /// * `x` - Set of variables to select.
    ///
    /// # Errors
    ///
    /// * If the set of variables is empty.
    /// * If any variable in the set is out of bounds.
    ///
    /// # Returns
    ///
    /// The evidence restricted to the specified variables.
    ///
    pub fn select(&self, x: &Set<usize>) -> Result<Self>
    where
        Self: Sized,
    {
        // Check that the variables are in bounds.
        x.iter().try_for_each(|&i| {
            if i >= self.labels.len() {
                return Err(Error::IndexOutOfBounds(i));
            }
            Ok(())
        })?;

        // Sort the indices.
        let mut x = x.clone();
        x.sort();

        // Get the new labels.
        let labels: Labels = x
            .iter()
            .map(|&i| {
                self.labels
                    .get_index(i)
                    .cloned()
                    .ok_or_else(|| Error::IndexOutOfBounds(i))
            })
            .collect::<Result<_>>()?;

        // Get the new values.
        let evidences = x.into_iter().enumerate().filter_map(|(i, x)| {
            self.evidences[x].clone().map(|mut evidence| {
                evidence.set_event(i);
                evidence
            })
        });

        // Create the new evidence.
        Self::new(labels, evidences)
    }
}