unobtanium-graph-algorithms 0.1.0

// SPDX-FileCopyrightText: 2025 Slatian 2025
//
// SPDX-License-Identifier: LGPL-3.0-only

use core::hash::Hash;
use std::collections::HashMap;
use std::iter::FusedIterator;

use crate::graph::Graph;
use crate::node::NodeId;

/// Holds the result of a ranking algorithm that was run on a [Graph].
pub struct RankedNodes<'a, T: Eq + Hash + Clone> {
	pub(crate) graph: &'a Graph<T>,
	pub(crate) scores: Vec<f32>,
}

/// Unsorted iterator over the [RankedNodes] structure.
pub struct UnsortedRankedNodesIter<'a, T: Eq + Hash + Clone> {
	inner: &'a RankedNodes<'a, T>,
	index: usize,
}

/// Iterator sorted by best score first over the elements in a [RankedNodes] structure.
pub struct SortedRankedNodesIter<'a, T: Eq + Hash + Clone> {
	inner: &'a RankedNodes<'a, T>,
	sorted_node_ids: Vec<(NodeId, f32)>,
	index: usize,
}

impl<'a, T: Eq + Hash + Clone> RankedNodes<'a, T> {
	/// Returns a hash map that maps node data to the resulting score
	pub fn hash_map(&self) -> HashMap<T, f32> {
		self.unsorted_iter()
			.map(|(data, score)| (data.clone(), score))
			.collect()
	}

	/// Returns an unsorted iterator over all nodes and their scores `(&data, score)`.
	///
	/// This is useful for building maps and filtering by data aspects before sorting.
	pub fn unsorted_iter(&'a self) -> UnsortedRankedNodesIter<'a, T> {
		UnsortedRankedNodesIter {
			inner: self,
			index: 0,
		}
	}

	/// Returns a sorted iterator over all nodes and their scores `(&data, score)`.
	///
	/// The highest scores are sorted first.
	///
	/// This is useful for getting the top *n* results of the ranking.
	pub fn sorted_iter(&'a self) -> SortedRankedNodesIter<'a, T> {
		let mut sorted_list: Vec<(NodeId, f32)> = self
			.scores
			.iter()
			.enumerate()
			.filter(|(_, score)| !score.is_nan())
			.map(|(n, score)| (NodeId(n), *score))
			.collect();
		sorted_list.sort_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap());
		SortedRankedNodesIter {
			inner: self,
			sorted_node_ids: sorted_list,
			index: 0,
		}
	}

	/// Convenience function to look the score of a single data point
	///
	/// This is useful if one only needs the scores for a few specific data points.
	pub fn get_score_for(&self, data: &T) -> Option<f32> {
		let node = self.graph.nodes.get2(data)?;
		self.scores.get(node.id.0).copied()
	}

	/// Returns the number of scores in this result set.
	///
	/// This is usually the number of nodes in the graph.
	pub fn len(&self) -> usize {
		self.scores.len()
	}

	/// Wheter the result set is empty.
	pub fn is_empty(&self) -> bool {
		self.scores.is_empty()
	}
}

impl<'a, T: Eq + Hash + Clone> Iterator for UnsortedRankedNodesIter<'a, T> {
	type Item = (&'a T, f32);

	fn next(&mut self) -> Option<Self::Item> {
		self.inner.scores.get(self.index).map(|score| {
			let data = &self
				.inner
				.graph
				.nodes
				.get1(&NodeId(self.index))
				.unwrap()
				.data;
			self.index += 1;
			(data, *score)
		})
	}
}

impl<T: Eq + Hash + Clone> FusedIterator for UnsortedRankedNodesIter<'_, T> {}

impl<'a, T: Eq + Hash + Clone> Iterator for SortedRankedNodesIter<'a, T> {
	type Item = (&'a T, f32);

	fn next(&mut self) -> Option<Self::Item> {
		self.sorted_node_ids
			.get(self.index)
			.map(|(node_id, score)| {
				let data = &self.inner.graph.nodes.get1(node_id).unwrap().data;
				self.index += 1;
				(data, *score)
			})
	}
}

impl<T: Eq + Hash + Clone> FusedIterator for SortedRankedNodesIter<'_, T> {}