causal-hub 0.0.5

use std::fmt::Display;

use itertools::Itertools;
use ndarray::prelude::*;

use crate::types::{Error, Labels, Result};

#[allow(clippy::upper_case_acronyms)]
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
#[repr(C)]
enum PKS {
    Unknown,
    Forbidden,
    Required,
}

impl PKS {
    #[inline]
    pub const fn is_unknown(&self) -> bool {
        matches!(self, Self::Unknown)
    }

    #[inline]
    pub const fn is_forbidden(&self) -> bool {
        matches!(self, Self::Forbidden)
    }

    #[inline]
    pub const fn is_required(&self) -> bool {
        matches!(self, Self::Required)
    }
}

impl Display for PKS {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Unknown => write!(f, "Unknown"),
            Self::Forbidden => write!(f, "Forbidden"),
            Self::Required => write!(f, "Required"),
        }
    }
}

/// A structure representing prior knowledge for structure learning.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PK {
    labels: Labels,
    adjacency_matrix: Array2<PKS>,
}

impl PK {
    /// Creates a new instance of prior knowledge.
    ///
    /// # Arguments
    ///
    /// * `labels` - The labels of the vertices.
    /// * `forbidden` - An iterator over forbidden edges.
    /// * `required` - An iterator over required edges.
    /// * `temporal_order` - An iterator over tiers of vertices, where each tier is an iterator of vertex indices.
    ///
    /// # Errors
    ///
    /// * If the labels are not sorted.
    /// * If any of the vertices in the forbidden, required, or temporal order are out of bounds.
    /// * If any edge is set to both forbidden and required.
    ///
    /// # Returns
    ///
    /// A new instance of prior knowledge.
    ///
    pub fn new<I, J, K, L>(
        labels: Labels,
        forbidden: I,
        required: J,
        temporal_order: K,
    ) -> Result<Self>
    where
        I: IntoIterator<Item = (usize, usize)>,
        J: IntoIterator<Item = (usize, usize)>,
        K: IntoIterator<Item = L>,
        L: IntoIterator<Item = usize>,
    {
        // Check if the labels are sorted.
        if !labels.iter().is_sorted() {
            return Err(Error::InvalidParameter("labels", "must be sorted"));
        }

        // Get the number of labels.
        let n = labels.len();
        // Initialize an adjacency matrix with `Unknown` state.
        let mut adjacency_matrix = Array::from_elem((n, n), PKS::Unknown);

        // Set the forbidden edges to `Forbidden`.
        forbidden.into_iter().try_for_each(|(i, j)| {
            // Check if the vertices are within bounds.
            if i >= n || j >= n {
                return Err(Error::IndexOutOfBounds(if i >= n { i } else { j }));
            }
            // Set the edge to `Forbidden`.
            adjacency_matrix[[i, j]] = PKS::Forbidden;
            Ok(())
        })?;

        // Set the required edges to `Required`.
        required.into_iter().try_for_each(|(i, j)| {
            // Check if the vertices are within bounds.
            if i >= n || j >= n {
                return Err(Error::IndexOutOfBounds(if i >= n { i } else { j }));
            }
            // Check that the edge is set to unknown.
            if !adjacency_matrix[[i, j]].is_unknown() {
                return Err(Error::PriorKnowledgeConflict(&format!(
                    "Edge ({i}, {j}) is already set to a non-unknown state: \n\
                    \t expected:    ({i}, {j}) set to 'Unknown', \n\
                    \t found:       ({i}, {j}) set to '{}'.",
                    adjacency_matrix[[i, j]]
                )));
            }
            // Set the edge to `Required`.
            adjacency_matrix[[i, j]] = PKS::Required;
            Ok(())
        })?;

        // Collect the tiered edges.
        let temporal_order: Vec<Vec<_>> = temporal_order
            .into_iter()
            .map(|tier| {
                let tier: Vec<_> = tier.into_iter().collect();
                // Check if the vertices are within bounds.
                tier.iter().try_for_each(|&i| {
                    if i >= n {
                        return Err(Error::IndexOutOfBounds(i));
                    }
                    Ok(())
                })?;
                Ok(tier)
            })
            .collect::<Result<_>>()?;
        // Edges from a vertex in a higher tier to a vertex in a lower tier are forbidden.
        temporal_order
            .iter()
            .enumerate()
            .try_for_each(|(t, tier)| {
                // Get the vertices in previous tiers.
                let previous_tiers = temporal_order[..t].iter().flatten();
                // For each vertex in the current tier, set edges to previous tiers as forbidden.
                tier.iter()
                    .cartesian_product(previous_tiers)
                    .try_for_each(|(&i, &j)| {
                        // Check that the edge is not required.
                        if adjacency_matrix[[i, j]].is_required() {
                            return Err(Error::PriorKnowledgeConflict(&format!(
                                "Edge ({i}, {j}) is already set to a 'Required' state: \n\
                            \t expected:    ({i}, {j}) set to 'Unknown' or 'Forbidden', \n\
                            \t found:       ({i}, {j}) set to '{}'.",
                                adjacency_matrix[[i, j]]
                            )));
                        }
                        // Set the edge to `Forbidden`.
                        adjacency_matrix[[i, j]] = PKS::Forbidden;
                        Ok(())
                    })
            })?;

        Ok(Self {
            labels,
            adjacency_matrix,
        })
    }

    /// Returns a reference to the labels of the prior knowledge.
    ///
    /// # Returns
    ///
    /// A reference to the labels.
    ///
    #[inline]
    pub const fn labels(&self) -> &Labels {
        &self.labels
    }

    /// Checks if an edge is unknown.
    ///
    /// # Arguments
    ///
    /// * `x` - The index of the first vertex.
    /// * `y` - The index of the second vertex.
    ///
    /// # Returns
    ///
    /// `true` if the edge is unknown, `false` otherwise.
    ///
    #[inline]
    pub fn is_unknown(&self, x: usize, y: usize) -> bool {
        self.adjacency_matrix[[x, y]].is_unknown()
    }

    /// Returns the unknown edges.
    ///
    /// # Returns
    ///
    /// A vector of tuples representing the indices of the unknown edges.
    ///
    pub fn unknown_edges(&self) -> Vec<(usize, usize)> {
        self.adjacency_matrix
            .indexed_iter()
            .filter_map(|((i, j), &state)| {
                if state.is_unknown() {
                    Some((i, j))
                } else {
                    None
                }
            })
            .collect()
    }

    /// Checks if an edge is forbidden.
    ///
    /// # Arguments
    ///
    /// * `x` - The index of the first vertex.
    /// * `y` - The index of the second vertex.
    ///
    /// # Returns
    ///
    /// `true` if the edge is forbidden, `false` otherwise.
    ///
    #[inline]
    pub fn is_forbidden(&self, x: usize, y: usize) -> bool {
        self.adjacency_matrix[[x, y]].is_forbidden()
    }

    /// Returns the forbidden edges.
    ///
    /// # Returns
    ///
    /// A vector of tuples representing the indices of the forbidden edges.
    ///
    pub fn forbidden_edges(&self) -> Vec<(usize, usize)> {
        self.adjacency_matrix
            .indexed_iter()
            .filter_map(|((i, j), &state)| {
                if state.is_forbidden() {
                    Some((i, j))
                } else {
                    None
                }
            })
            .collect()
    }

    /// Checks if an edge is required.
    ///
    /// # Arguments
    ///
    /// * `x` - The index of the first vertex.
    /// * `y` - The index of the second vertex.
    ///
    /// # Returns
    ///
    /// `true` if the edge is required, `false` otherwise.
    ///
    #[inline]
    pub fn is_required(&self, x: usize, y: usize) -> bool {
        self.adjacency_matrix[[x, y]].is_required()
    }

    /// Returns the required edges.
    ///
    /// # Returns
    ///
    /// A vector of tuples representing the indices of the required edges.
    ///
    pub fn required_edges(&self) -> Vec<(usize, usize)> {
        self.adjacency_matrix
            .indexed_iter()
            .filter_map(|((i, j), &state)| {
                if state.is_required() {
                    Some((i, j))
                } else {
                    None
                }
            })
            .collect()
    }
}