causal-hub 0.0.6

use std::borrow::Cow;

use approx::{AbsDiffEq, RelativeEq};
use ndarray::prelude::*;
use serde::{
    Deserialize, Deserializer, Serialize, Serializer,
    de::{MapAccess, Visitor},
    ser::SerializeMap,
};

use crate::{
    datasets::{CatSample, CatTrj, CatTrjs},
    impl_json_io,
    models::{BN, CIM, CTBN, CatBN, CatCIM, CatCPD, CatSupport, DiGraph, Graph, HasLabels},
    set,
    types::{Error, Labels, Map, Result, Set},
};

/// A categorical continuous time Bayesian network.
#[derive(Clone, Debug)]
pub struct CatCTBN {
    /// The name of the model.
    name: Option<String>,
    /// The description of the model.
    description: Option<String>,
    /// The labels of the variables.
    labels: Labels,
    /// The support of the variables.
    support: CatSupport,
    /// The shape of the variables.
    shape: Array1<usize>,
    /// The initial distribution.
    initial_distribution: CatBN,
    /// The underlying graph.
    graph: DiGraph,
    /// The conditional intensity matrices.
    cims: Map<String, CatCIM>,
}

impl CatCTBN {
    /// Returns the name of the model, if any.
    ///
    /// # Returns
    ///
    /// The name of the model, if it exists.
    ///
    #[inline]
    pub fn name(&self) -> Option<&str> {
        self.name.as_deref()
    }

    /// Returns the description of the model, if any.
    ///
    /// # Returns
    ///
    /// The description of the model, if it exists.
    ///
    #[inline]
    pub fn description(&self) -> Option<&str> {
        self.description.as_deref()
    }

    /// Returns the support of the variables.
    ///
    /// # Returns
    ///
    /// A reference to the support of the variables.
    ///
    #[inline]
    pub const fn support(&self) -> &CatSupport {
        self.initial_distribution.support()
    }
}

impl PartialEq for CatCTBN {
    fn eq(&self, other: &Self) -> bool {
        self.labels.eq(&other.labels)
            && self.support.eq(&other.support)
            && self.shape.eq(&other.shape)
            && self.initial_distribution.eq(&other.initial_distribution)
            && self.graph.eq(&other.graph)
            && self.cims.eq(&other.cims)
    }
}

impl AbsDiffEq for CatCTBN {
    type Epsilon = f64;

    fn default_epsilon() -> Self::Epsilon {
        Self::Epsilon::default_epsilon()
    }

    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
        self.labels.eq(&other.labels)
            && self.support.eq(&other.support)
            && self.shape.eq(&other.shape)
            && self.initial_distribution.eq(&other.initial_distribution)
            && self.graph.eq(&other.graph)
            && self.cims.iter().zip(&other.cims).all(
                |((label, intensity), (other_label, other_cim))| {
                    label.eq(other_label) && intensity.abs_diff_eq(other_cim, epsilon)
                },
            )
    }
}

impl RelativeEq for CatCTBN {
    fn default_max_relative() -> Self::Epsilon {
        Self::Epsilon::default_max_relative()
    }

    fn relative_eq(
        &self,
        other: &Self,
        epsilon: Self::Epsilon,
        max_relative: Self::Epsilon,
    ) -> bool {
        self.labels.eq(&other.labels)
            && self.support.eq(&other.support)
            && self.shape.eq(&other.shape)
            && self.initial_distribution.eq(&other.initial_distribution)
            && self.graph.eq(&other.graph)
            && self.cims.iter().zip(&other.cims).all(
                |((label, intensity), (other_label, other_cim))| {
                    label.eq(other_label) && intensity.relative_eq(other_cim, epsilon, max_relative)
                },
            )
    }
}

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

impl CTBN for CatCTBN {
    type CIM = CatCIM;
    type Support = CatSupport;
    type InitialDistribution = CatBN;
    type Event = (f64, CatSample);
    type Trajectory = CatTrj;
    type Trajectories = CatTrjs;

    #[inline]
    fn support(&self) -> Cow<'_, Self::Support> {
        Cow::Borrowed(&self.support)
    }

    fn new<I>(graph: DiGraph, cims: I) -> Result<Self>
    where
        I: IntoIterator<Item = Self::CIM>,
    {
        // Collect the CIMs into a map.
        let mut cims: Map<_, _> = cims
            .into_iter()
            .map(|x| {
                if x.labels().len() != 1 {
                    return Err(Error::InvalidParameter(
                        "cim",
                        "CIM must contain exactly one label.",
                    ));
                }
                Ok((x.labels()[0].to_owned(), x))
            })
            .collect::<Result<_>>()?;
        // Sort the CIMs by their labels.
        cims.sort_keys();

        // Allocate the support of the variables.
        let mut support: CatSupport = Default::default();
        // Insert the support of the variables into the map to check if they are the same.
        cims.values().try_for_each(|intensity| {
            intensity
                .support()
                .iter()
                .chain(intensity.conditioning_support())
                .try_for_each(|(l, stats)| {
                    // Check if the support are already in the map.
                    if let Some(existing_states) = support.get(l) {
                        // Check if the support are the same.
                        if existing_states != stats {
                            return Err(Error::InvalidParameter(
                                "cims",
                                &format!("CatSupport of `{l}` must be the same across CIMs."),
                            ));
                        }
                    } else {
                        // Insert the support into the map.
                        support.insert(l.to_owned(), stats.clone());
                    }
                    Ok(())
                })
        })?;
        // Sort the support of the variables.
        support.sort_keys();

        // Get the labels of the variables.
        let labels: Labels = support.keys().cloned().collect();
        // Get the shape of the variables.
        let shape = Array::from_iter(support.values().map(Set::len));

        // Check same number of graph labels and CIMs.
        if !graph.labels().iter().eq(cims.keys()) {
            return Err(Error::LabelMismatch("graph labels", "distributions labels"));
        }

        // Check if all vertices have the same labels as their parents.
        graph.vertices().iter().try_for_each(|&i| {
            // Get the parents of the vertex.
            let pa_i = graph.parents(&set![i])?.into_iter();
            let pa_i: &Labels = &pa_i.map(|j| labels[j].to_owned()).collect(); // FIXME: Use references to avoid clones
            // Get the conditioning labels of the CIM.
            let pa_j = cims[&labels[i]].conditioning_labels();
            // Check they are the same.
            if pa_i != pa_j {
                return Err(Error::LabelMismatch(
                    &format!("{pa_i:?}"),
                    &format!("{pa_j:?}"),
                ));
            }
            Ok(())
        })?;

        // Initialize an empty graph for the uniform initial distribution.
        let initial_graph = DiGraph::empty(graph.labels())?;
        // Initialize the CPDs as uniform distributions.
        let initial_cpds: Vec<_> = cims
            .values()
            .map(|intensity| {
                // Get label and support of the CIM.
                let support = intensity.support().clone();
                // Set empty conditioning support.
                let conditioning_support = CatSupport::default();
                // Set uniform parameters.
                let alpha = intensity.shape().product();
                let parameters = Array::from_vec(vec![1. / alpha as f64; alpha]);
                let parameters = parameters.insert_axis(Axis(0));
                // Construct the CPD.
                CatCPD::new(support, conditioning_support, parameters)
            })
            .collect::<Result<_>>()?;
        // Initialize a uniform initial distribution.
        let initial_distribution = CatBN::new(initial_graph, initial_cpds)?;

        Ok(Self {
            name: None,
            description: None,
            labels,
            support,
            shape,
            initial_distribution,
            graph,
            cims,
        })
    }

    fn initial_distribution(&self) -> &Self::InitialDistribution {
        &self.initial_distribution
    }

    fn graph(&self) -> &DiGraph {
        &self.graph
    }

    fn cims(&self) -> &Map<String, Self::CIM> {
        &self.cims
    }

    fn parameters_size(&self) -> usize {
        // Parameters size of the initial distribution.
        self.initial_distribution.parameters_size()
            // Parameters size of the CIMs.
            + self
                .cims
                .values()
                .map(|x| x.parameters_size())
                .sum::<usize>()
    }

    fn with_optionals<I>(
        name: Option<String>,
        description: Option<String>,
        initial_distribution: Self::InitialDistribution,
        graph: DiGraph,
        cims: I,
    ) -> Result<Self>
    where
        I: IntoIterator<Item = Self::CIM>,
    {
        // Check name is not empty string.
        if let Some(name) = &name
            && name.is_empty()
        {
            return Err(Error::InvalidParameter("name", "cannot be empty"));
        }
        // Check description is not empty string.
        if let Some(description) = &description
            && description.is_empty()
        {
            return Err(Error::InvalidParameter("description", "cannot be empty"));
        }

        // Construct the categorical CTBN.
        let mut continuous_time_bayesian_network = Self::new(graph, cims)?;

        // Check the initial distribution has same labels.
        if !initial_distribution
            .labels()
            .eq(continuous_time_bayesian_network.labels())
        {
            return Err(Error::LabelMismatch(
                "initial distribution labels",
                "cims labels",
            ));
        }
        // Check the initial distribution has same support.
        if !initial_distribution
            .cpds()
            .into_iter()
            .zip(continuous_time_bayesian_network.cims())
            .all(|((_, distribution), (_, intensity))| {
                distribution.support().eq(intensity.support())
            })
        {
            return Err(Error::InvalidParameter(
                "initial distribution",
                "Initial distribution support must be the same as the CIMs support.",
            ));
        }

        // Set the optional fields.
        continuous_time_bayesian_network.name = name;
        continuous_time_bayesian_network.description = description;
        continuous_time_bayesian_network.initial_distribution = initial_distribution;

        Ok(continuous_time_bayesian_network)
    }
}

impl Serialize for CatCTBN {
    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        // Count the elements to serialize.
        let mut size = 4;
        size += self.name.is_some() as usize;
        size += self.description.is_some() as usize;

        // Allocate the map.
        let mut map = serializer.serialize_map(Some(size))?;

        // Convert the CIMs to a flat format.
        let cims: Vec<_> = self.cims.values().cloned().collect();

        // Serialize name, if any.
        if let Some(name) = &self.name {
            map.serialize_entry("name", name)?;
        }
        // Serialize description, if any.
        if let Some(description) = &self.description {
            map.serialize_entry("description", description)?;
        }
        // Serialize initial distribution.
        map.serialize_entry("initial_distribution", &self.initial_distribution)?;
        // Serialize graph.
        map.serialize_entry("graph", &self.graph)?;
        // Serialize CIMs.
        map.serialize_entry("cims", &cims)?;
        // Serialize type.
        map.serialize_entry("type", "catctbn")?;

        // Finalize the map serialization.
        map.end()
    }
}

impl<'de> Deserialize<'de> for CatCTBN {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(field_identifier, rename_all = "snake_case")]
        enum Field {
            Name,
            Description,
            InitialDistribution,
            Graph,
            Cims,
            Type,
        }

        struct CatCTBNVisitor;

        impl<'de> Visitor<'de> for CatCTBNVisitor {
            type Value = CatCTBN;

            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
                formatter.write_str("struct CatCTBN")
            }

            fn visit_map<V>(self, mut map: V) -> std::result::Result<CatCTBN, V::Error>
            where
                V: MapAccess<'de>,
            {
                use serde::de::Error as E;

                // Allocate fields
                let mut name = None;
                let mut description = None;
                let mut initial_distribution = None;
                let mut graph = None;
                let mut cims = None;
                let mut type_ = None;

                // Parse the map.
                while let Some(key) = map.next_key()? {
                    match key {
                        Field::Name => {
                            if name.is_some() {
                                return Err(E::duplicate_field("name"));
                            }
                            name = Some(map.next_value()?);
                        }
                        Field::Description => {
                            if description.is_some() {
                                return Err(E::duplicate_field("description"));
                            }
                            description = Some(map.next_value()?);
                        }
                        Field::InitialDistribution => {
                            if initial_distribution.is_some() {
                                return Err(E::duplicate_field("initial_distribution"));
                            }
                            initial_distribution = Some(map.next_value()?);
                        }
                        Field::Graph => {
                            if graph.is_some() {
                                return Err(E::duplicate_field("graph"));
                            }
                            graph = Some(map.next_value()?);
                        }
                        Field::Cims => {
                            if cims.is_some() {
                                return Err(E::duplicate_field("cims"));
                            }
                            cims = Some(map.next_value()?);
                        }
                        Field::Type => {
                            if type_.is_some() {
                                return Err(E::duplicate_field("type"));
                            }
                            type_ = Some(map.next_value()?);
                        }
                    }
                }

                // Check required fields.
                let initial_distribution =
                    initial_distribution.ok_or_else(|| E::missing_field("initial_distribution"))?;
                let graph = graph.ok_or_else(|| E::missing_field("graph"))?;
                let cims = cims.ok_or_else(|| E::missing_field("cims"))?;

                // Check type is correct.
                let type_: String = type_.ok_or_else(|| E::missing_field("type"))?;
                if type_ != "catctbn" {
                    return Err(E::custom(format!(
                        "Invalid type for CatCTBN: expected 'catctbn', found '{type_}'"
                    )));
                }

                // Set helper types.
                let cims: Vec<_> = cims;

                CatCTBN::with_optionals(name, description, initial_distribution, graph, cims)
                    .map_err(serde::de::Error::custom)
            }
        }

        const FIELDS: &[&str] = &[
            "name",
            "description",
            "initial_distribution",
            "graph",
            "cims",
            "type",
        ];

        deserializer.deserialize_struct("CatCTBN", FIELDS, CatCTBNVisitor)
    }
}

// Implement `JsonIO` for `CatCTBN`.
impl_json_io!(CatCTBN);