use anyhow::Result;
use petgraph::stable_graph::{EdgeIndex, NodeIndex};
use std::collections::{BTreeMap, HashMap};
use crate::algorithms::equal_angle::{Pt, assign_coordinates_to_nodes};
use crate::nexus::network_writer::{leaf_label, node_degree};
use crate::phylo::phylo_graph::PhyloGraph;
pub type GraphEdge = (usize, usize, usize, i32, f64);
#[derive(Debug, Clone, Default)]
pub struct PhyloSplitsGraph {
pub base: PhyloGraph,
edge_split_map: Option<HashMap<EdgeIndex, i32>>,
edge_angle_map: Option<HashMap<EdgeIndex, f64>>,
pub rotation: HashMap<NodeIndex, Vec<EdgeIndex>>,
pub node_ids: BTreeMap<NodeIndex, usize>,
}
impl PhyloSplitsGraph {
pub fn new() -> Self {
Self::default()
}
pub fn clear(&mut self) {
self.base.clear();
self.edge_split_map = None;
self.edge_angle_map = None;
self.rotation.clear();
}
fn splits_mut(&mut self) -> &mut HashMap<EdgeIndex, i32> {
self.edge_split_map.get_or_insert_with(HashMap::new)
}
fn angles_mut(&mut self) -> &mut HashMap<EdgeIndex, f64> {
self.edge_angle_map.get_or_insert_with(HashMap::new)
}
pub fn set_split(&mut self, e: EdgeIndex, id: i32) {
self.splits_mut().insert(e, id);
}
pub fn get_split(&self, e: EdgeIndex) -> i32 {
self.edge_split_map
.as_ref()
.and_then(|m| m.get(&e).copied())
.unwrap_or(0)
}
pub fn set_angle(&mut self, e: EdgeIndex, val: f64) {
self.angles_mut().insert(e, val);
}
pub fn get_angle(&self, e: EdgeIndex) -> f64 {
self.edge_angle_map
.as_ref()
.and_then(|m| m.get(&e).copied())
.unwrap_or(0.0)
}
pub fn create_node_ids(&mut self) {
self.node_ids.clear();
let mut next = 1usize;
for v in self.base.graph.node_indices() {
self.node_ids.insert(v, next);
next += 1;
}
}
pub fn count_nodes(&self) -> usize {
self.base.graph.node_count()
}
pub fn count_edges(&self) -> usize {
self.base.graph.edge_count()
}
pub fn max_split_id(&self) -> i32 {
let mut max_id = 0;
for e in self.base.graph.edge_indices() {
let id = self.get_split(e);
if id > max_id {
max_id = id;
}
}
max_id
}
pub fn new_edge(&mut self, u: NodeIndex, v: NodeIndex) -> Result<EdgeIndex> {
let e = self.base.new_edge(u, v)?;
self.rot_append(u, e);
self.rot_append(v, e);
Ok(e)
}
pub fn new_edge_after(
&mut self,
u: NodeIndex,
v: NodeIndex,
f0_at_v: EdgeIndex,
) -> Result<EdgeIndex> {
let e = self.base.new_edge(u, v)?;
self.rot_append(u, e);
self.rot_insert_after(v, f0_at_v, e);
Ok(e)
}
pub fn remove_edge(&mut self, e: EdgeIndex) -> bool {
if let Some((u, v)) = self.base.graph.edge_endpoints(e) {
self.rot_remove(u, e);
self.rot_remove(v, e);
self.base.graph.remove_edge(e).is_some()
} else {
false
}
}
pub fn remove_node_and_cleanup(&mut self, v: NodeIndex) {
if let Some(rs) = self.rotation.remove(&v) {
for e in rs {
if let Some((a, b)) = self.base.graph.edge_endpoints(e) {
let other = if a == v { b } else { a };
self.rot_remove(other, e);
self.base.graph.remove_edge(e);
}
}
}
self.base.remove_node_and_cleanup(v);
}
#[inline]
pub fn rot_mut(&mut self, v: NodeIndex) -> &mut Vec<EdgeIndex> {
self.rotation.entry(v).or_default()
}
#[inline]
pub fn rot(&self, v: NodeIndex) -> &[EdgeIndex] {
self.rotation.get(&v).map(|r| r.as_slice()).unwrap_or(&[])
}
#[inline]
pub fn rot_insert_after(&mut self, v: NodeIndex, after: EdgeIndex, e_new: EdgeIndex) {
let r = self.rot_mut(v);
if let Some(i) = r.iter().position(|&x| x == after) {
r.insert(i + 1, e_new);
} else {
r.push(e_new);
}
}
#[inline]
pub fn rot_append(&mut self, v: NodeIndex, e_new: EdgeIndex) {
self.rot_mut(v).push(e_new);
}
#[inline]
pub fn rot_remove(&mut self, v: NodeIndex, e: EdgeIndex) {
if let Some(r) = self.rotation.get_mut(&v) {
if let Some(i) = r.iter().position(|&x| x == e) {
r.remove(i);
}
if r.is_empty() {
self.rotation.remove(&v);
}
}
}
#[inline]
pub fn first_adjacent_edge(&self, v: NodeIndex) -> Option<EdgeIndex> {
self.rot(v).first().copied()
}
pub fn next_adjacent_edge_cyclic(&self, v: NodeIndex, e: EdgeIndex) -> Option<EdgeIndex> {
let r = self.rot(v);
if r.is_empty() {
return None;
}
if let Some(i) = r.iter().position(|&x| x == e) {
Some(r[(i + 1) % r.len()])
} else {
r.first().copied()
}
}
pub fn is_leaf_edge(&self, e: EdgeIndex) -> bool {
if let Some((a, b)) = self.base.graph.edge_endpoints(e) {
let da = self.base.graph.neighbors(a).count();
let db = self.base.graph.neighbors(b).count();
da == 1 || db == 1
} else {
false
}
}
pub fn get_node_translations(
&self,
taxa_labels_1based: &[String],
) -> Result<Vec<(usize, String)>> {
let mut translations = Vec::new();
for v in self.base.graph.node_indices() {
if node_degree(&self.base, v) == 1
&& let Some(lbl) = leaf_label(&self.base, v, taxa_labels_1based)
{
translations.push((self.node_ids[&v], lbl));
}
}
Ok(translations)
}
pub fn get_node_positions(&self) -> Result<Vec<(usize, f64, f64)>> {
let coords = assign_coordinates_to_nodes(true, self, 1, 0);
let mut positions = Vec::new();
for v in self.base.graph.node_indices() {
let id = self.node_ids[&v];
let pt = coords.get(&v).copied().unwrap_or(Pt(0.0, 0.0));
positions.push((id, pt.0, pt.1));
}
Ok(positions)
}
pub fn get_graph_edges(&self) -> Result<Vec<GraphEdge>> {
let mut eid = 1usize;
let mut edges = Vec::new();
for e in self.base.graph.edge_indices() {
let (u, v) = self.base.graph.edge_endpoints(e).expect("valid endpoints");
let su = self.node_ids[&u];
let sv = self.node_ids[&v];
let sid = self.get_split(e);
let wgt = self.base.weight(e);
edges.push((eid, su, sv, sid, wgt));
eid += 1;
}
Ok(edges)
}
}