use anyhow::Result;
use petgraph::Undirected;
use petgraph::stable_graph::{EdgeIndex, NodeIndex, StableGraph};
use std::collections::HashMap;
pub const DEFAULT_WEIGHT: f64 = 1.0;
#[derive(Debug, Clone, Default)]
pub struct NodeData {
pub label: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct EdgeData {
pub label: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct PhyloGraph {
pub graph: StableGraph<NodeData, EdgeData, Undirected>,
edge_weights: Option<HashMap<EdgeIndex, f64>>,
taxon2node: Option<HashMap<usize, NodeIndex>>,
node2taxa: Option<HashMap<NodeIndex, Vec<usize>>>,
}
impl PhyloGraph {
pub fn new() -> Self {
Self::default()
}
pub fn clear(&mut self) {
self.graph = StableGraph::default();
self.edge_weights = None;
self.taxon2node = None;
self.node2taxa = None;
}
pub fn new_node(&mut self) -> NodeIndex {
self.graph.add_node(NodeData::default())
}
pub fn new_node_with_label<S: Into<String>>(&mut self, label: S) -> NodeIndex {
self.graph.add_node(NodeData {
label: Some(label.into()),
})
}
pub fn set_node_label<S: Into<String>>(&mut self, v: NodeIndex, label: S) {
if let Some(nd) = self.graph.node_weight_mut(v) {
nd.label = Some(label.into());
}
}
pub fn node_label(&self, v: NodeIndex) -> Option<&str> {
self.graph.node_weight(v).and_then(|nd| nd.label.as_deref())
}
pub fn new_edge(&mut self, u: NodeIndex, v: NodeIndex) -> Result<EdgeIndex> {
if u == v {
return Err(anyhow::anyhow!("Illegal self-edge"));
}
Ok(self.graph.add_edge(u, v, EdgeData::default()))
}
pub fn set_edge_label<S: Into<String>>(&mut self, e: EdgeIndex, label: S) {
if let Some(ed) = self.graph.edge_weight_mut(e) {
ed.label = Some(label.into());
}
}
pub fn edge_label(&self, e: EdgeIndex) -> Option<&str> {
self.graph.edge_weight(e).and_then(|ed| ed.label.as_deref())
}
pub fn remove_node_and_cleanup(&mut self, v: NodeIndex) -> bool {
self.clear_taxa_for_node(v);
self.graph.remove_node(v).is_some()
}
fn weights_mut(&mut self) -> &mut HashMap<EdgeIndex, f64> {
self.edge_weights.get_or_insert_with(HashMap::new)
}
pub fn set_weight(&mut self, e: EdgeIndex, val: f64) {
if val == DEFAULT_WEIGHT && self.edge_weights.is_none() {
return;
}
self.weights_mut().insert(e, val);
}
pub fn weight(&self, e: EdgeIndex) -> f64 {
self.edge_weights
.as_ref()
.and_then(|m| m.get(&e).copied())
.unwrap_or(DEFAULT_WEIGHT)
}
fn taxon2node_mut(&mut self) -> &mut HashMap<usize, NodeIndex> {
self.taxon2node.get_or_insert_with(HashMap::new)
}
fn node2taxa_mut(&mut self) -> &mut HashMap<NodeIndex, Vec<usize>> {
self.node2taxa.get_or_insert_with(HashMap::new)
}
pub fn taxon2node(&self) -> Option<&HashMap<usize, NodeIndex>> {
self.taxon2node.as_ref()
}
pub fn number_of_taxa(&self) -> usize {
self.taxon2node.as_ref().map(|m| m.len()).unwrap_or(0)
}
pub fn has_edge_weights(&self) -> bool {
self.edge_weights.is_some()
}
pub fn node2taxa(&self) -> Option<&HashMap<NodeIndex, Vec<usize>>> {
self.node2taxa.as_ref()
}
pub fn get_taxon_node(&self, taxon: usize) -> Option<NodeIndex> {
self.taxon2node.as_ref()?.get(&taxon).copied()
}
pub fn get_node_taxon(&self, node: NodeIndex) -> Option<&[usize]> {
self.node2taxa.as_ref()?.get(&node).map(|v| v.as_slice())
}
pub fn add_taxon(&mut self, v: NodeIndex, taxon: usize) {
self.taxon2node_mut().insert(taxon, v);
let vec = self.node2taxa_mut().entry(v).or_default();
if !vec.contains(&taxon) {
vec.push(taxon);
}
}
pub fn clear_taxa_for_node(&mut self, v: NodeIndex) {
if let (Some(t2n), Some(n2t)) = (self.taxon2node.as_mut(), self.node2taxa.as_mut()) {
if let Some(list) = n2t.get(&v) {
for &t in list {
if t2n.get(&t).copied() == Some(v) {
t2n.remove(&t);
}
}
}
n2t.remove(&v);
}
}
pub fn get_opposite(&self, v: NodeIndex, e: EdgeIndex) -> NodeIndex {
let endpoints = self.graph.edge_endpoints(e).expect("valid edge endpoints");
if endpoints.0 == v {
endpoints.1
} else {
endpoints.0
}
}
}