open-hypergraphs 0.3.1

use crate::array::{Array, ArrayKind, NaturalArray};
use crate::category::*;
use crate::finite_function::{coequalizer_universal, FiniteFunction};
use crate::indexed_coproduct::*;
use crate::operations::Operations;
use crate::semifinite::*;
#[cfg(feature = "experimental")]
use crate::strict::hypergraph::arrow::HypergraphArrow;

use core::fmt::Debug;
use core::ops::Add;
use num_traits::Zero;

#[derive(Debug)]
pub enum InvalidHypergraph<K: ArrayKind> {
    SourcesCount(K::I, K::I),
    TargetsCount(K::I, K::I),
    SourcesSet(K::I, K::I),
    TargetsSet(K::I, K::I),
}

pub struct Hypergraph<K: ArrayKind, O, A> {
    pub s: IndexedCoproduct<K, FiniteFunction<K>>,
    pub t: IndexedCoproduct<K, FiniteFunction<K>>,
    pub w: SemifiniteFunction<K, O>,
    pub x: SemifiniteFunction<K, A>,
}

impl<K: ArrayKind, O, A> Hypergraph<K, O, A>
where
    K::Type<K::I>: AsRef<K::Index>,
    K::Type<K::I>: NaturalArray<K>,
    K::Type<O>: Array<K, O>,
    K::Type<A>: Array<K, A>,
{
    /// Safely create a Hypergraph, ensuring its data is valid.
    pub fn new(
        s: IndexedCoproduct<K, FiniteFunction<K>>,
        t: IndexedCoproduct<K, FiniteFunction<K>>,
        w: SemifiniteFunction<K, O>,
        x: SemifiniteFunction<K, A>,
    ) -> Result<Hypergraph<K, O, A>, InvalidHypergraph<K>> {
        let h = Hypergraph { s, t, w, x };
        h.validate()
    }

    /// A hypergraph is valid when for both sources and targets segmented arrays:
    ///
    /// 1. Number of segments is equal to number of operations (`x.len()`)
    /// 2. Values of segmented array are indices in set `w.source()`
    ///
    pub fn validate(self) -> Result<Self, InvalidHypergraph<K>> {
        // num ops, wires
        let n_x = self.x.len();
        let n_w = self.w.len();

        // Sources segmented array has as many segments as operations
        if self.s.len() != self.x.len() {
            return Err(InvalidHypergraph::SourcesCount(self.s.len(), n_x));
        }

        // Targets segmented array has as many segments as operations
        if self.t.len() != self.x.len() {
            return Err(InvalidHypergraph::TargetsCount(self.t.len(), n_x));
        }

        // *values* of sources segmented array should index into operations
        let n_s = self.s.values.target();
        if n_s != n_w {
            return Err(InvalidHypergraph::SourcesSet(n_s, n_w));
        }

        // *values* of targets segmented array should index into operations
        let n_t = self.t.values.target();
        if n_t != n_w {
            return Err(InvalidHypergraph::TargetsSet(n_t, n_w));
        }

        Ok(self)
    }

    // TODO: This is the unit object - put inside category interface?
    /// Construct the empty hypergraph with no nodes and no hyperedges.
    pub fn empty() -> Hypergraph<K, O, A> {
        Hypergraph {
            s: IndexedCoproduct::initial(K::I::zero()),
            t: IndexedCoproduct::initial(K::I::zero()),
            w: SemifiniteFunction::zero(),
            x: SemifiniteFunction::zero(),
        }
    }

    /// The discrete hypergraph, consisting of hypernodes labeled in `O`.
    pub fn discrete(w: SemifiniteFunction<K, O>) -> Hypergraph<K, O, A> {
        Hypergraph {
            s: IndexedCoproduct::initial(w.len()),
            t: IndexedCoproduct::initial(w.len()),
            w,
            x: SemifiniteFunction::zero(),
        }
    }

    pub fn is_discrete(&self) -> bool {
        self.s.is_empty() && self.t.is_empty() && self.x.0.is_empty()
    }

    pub fn coproduct(&self, other: &Self) -> Self {
        Hypergraph {
            s: self.s.tensor(&other.s),
            t: self.t.tensor(&other.t),
            w: self.w.coproduct(&other.w),
            x: self.x.coproduct(&other.x),
        }
    }

    pub fn tensor_operations(Operations { x, a, b }: Operations<K, O, A>) -> Hypergraph<K, O, A> {
        // NOTE: the validity of the result assumes validity of `operations`.
        let inj0 = FiniteFunction::inj0(a.values.len(), b.values.len());
        let inj1 = FiniteFunction::inj1(a.values.len(), b.values.len());
        let s = IndexedCoproduct::new(a.sources, inj0).expect("invalid Operations?");
        let t = IndexedCoproduct::new(b.sources, inj1).expect("invalid Operations?");
        let w = a.values + b.values;
        Hypergraph { s, t, w, x }
    }
}

impl<K: ArrayKind, O, A> Hypergraph<K, O, A>
where
    K::Type<K::I>: NaturalArray<K>,
    K::Type<O>: Array<K, O>,
{
    /// The number of occurrences of `node` as a target across all hyperedges.
    pub fn in_degree(&self, node: K::I) -> K::I {
        let node = node.clone();
        assert!(node < self.w.len(), "node id {:?} is out of bounds", node);
        let counts = (self.t.values.table.as_ref() as &K::Type<K::I>).bincount(self.w.len());
        counts.get(node)
    }

    /// The number of occurrences of `node` as a source across all hyperedges.
    pub fn out_degree(&self, node: K::I) -> K::I {
        let node = node.clone();
        assert!(node < self.w.len(), "node id {:?} is out of bounds", node);
        let counts = (self.s.values.table.as_ref() as &K::Type<K::I>).bincount(self.w.len());
        counts.get(node)
    }
}

impl<K: ArrayKind, O, A> Hypergraph<K, O, A>
where
    K::Type<K::I>: AsRef<K::Index>,
    K::Type<K::I>: NaturalArray<K>,
    K::Type<O>: Array<K, O> + PartialEq,
    K::Type<A>: Array<K, A>,
{
    pub fn coequalize_vertices(&self, q: &FiniteFunction<K>) -> Option<Hypergraph<K, O, A>> {
        // TODO: wrap coequalizers in a newtype!
        let s = self.s.map_values(q)?;
        let t = self.t.map_values(q)?;
        let w = SemifiniteFunction(coequalizer_universal(q, &self.w.0)?);
        let x = self.x.clone();
        Some(Hypergraph { s, t, w, x })
    }

    // Compute the pushout of a span of wire maps into hypergraphs.
    #[cfg(feature = "experimental")]
    pub(crate) fn pushout_along_span(
        left: &Hypergraph<K, O, A>,
        right: &Hypergraph<K, O, A>,
        f: &FiniteFunction<K>,
        g: &FiniteFunction<K>,
    ) -> Option<(
        Hypergraph<K, O, A>,
        HypergraphArrow<K, O, A>,
        HypergraphArrow<K, O, A>,
    )>
    where
        K::Type<K::I>: NaturalArray<K>,
        K::Type<O>: Array<K, O> + PartialEq,
        K::Type<A>: Array<K, A> + PartialEq,
    {
        if f.target() != left.w.len() || g.target() != right.w.len() {
            return None;
        }

        // Lift span legs into the shared coproduct codomain left.w + right.w
        // and coequalize to obtain the wire identifications.
        let f_lift = f.inject0(right.w.len());
        let g_lift = g.inject1(left.w.len());
        let q = f_lift.coequalizer(&g_lift)?;
        // Form the coproduct hypergraph and quotient its vertices by the coequalizer.
        let coproduct = left + right;
        let target = coproduct.coequalize_vertices(&q)?;

        // Build the induced arrows from each side of the span into the pushout.
        let w_left = FiniteFunction::inj0(left.w.len(), right.w.len()).compose(&q)?;
        let w_right = FiniteFunction::inj1(left.w.len(), right.w.len()).compose(&q)?;
        let x_left = FiniteFunction::inj0(left.x.len(), right.x.len());
        let x_right = FiniteFunction::inj1(left.x.len(), right.x.len());

        let left_arrow = HypergraphArrow::new(left.clone(), target.clone(), w_left, x_left).ok()?;
        let right_arrow =
            HypergraphArrow::new(right.clone(), target.clone(), w_right, x_right).ok()?;

        Some((target, left_arrow, right_arrow))
    }
}

impl<K: ArrayKind, O, A> Add<&Hypergraph<K, O, A>> for &Hypergraph<K, O, A>
where
    K::Type<K::I>: NaturalArray<K>,
    K::Type<O>: Array<K, O>,
    K::Type<A>: Array<K, A>,
{
    type Output = Hypergraph<K, O, A>;

    fn add(self, rhs: &Hypergraph<K, O, A>) -> Self::Output {
        self.coproduct(rhs)
    }
}

impl<K: ArrayKind, O, A> Clone for Hypergraph<K, O, A>
where
    K::Type<K::I>: Clone,
    K::Type<O>: Clone,
    K::Type<A>: Clone,
{
    fn clone(&self) -> Self {
        Self {
            s: self.s.clone(),
            t: self.t.clone(),
            w: self.w.clone(),
            x: self.x.clone(),
        }
    }
}

// NOTE: manual Debug required because we need to specify array bounds.
impl<K: ArrayKind, O: Debug, A: Debug> Debug for Hypergraph<K, O, A>
where
    K::Index: Debug,
    K::Type<K::I>: Debug,
    K::Type<O>: Debug,
    K::Type<A>: Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Hypergraph")
            .field("s", &self.s)
            .field("t", &self.t)
            .field("w", &self.w)
            .field("x", &self.x)
            .finish()
    }
}