PronounsToday/src/weighted_table.rs

use crate::{
	user_preferences::ParseError,
	util::{pcg64, pcg64_seed},
};

use std::collections::HashMap;

use time::{Date, Month, OffsetDateTime};

/// The start of the COVID-19 lockdowns
///
/// This is used as an epoch in order to convert from a given date to an integer seed.  This is
/// specified as part of the algorithm for randomly selecting from a weighted list.
pub const COVID_EPOCH: Date = match Date::from_calendar_date(2020, Month::January, 26) {
	Ok(d) => d,
	Err(_) => Date::MIN, // This never runs, but we can't unwrap, so this is what we're stuck with
};

/// A list of pronouns and their associated weights, used for random selection
///
/// Weights are typically representative of a user's preference towards a pronoun.  A pronoun with
/// a weight of 10 is twice as likely to be selected as a pronoun with a weight of 5.
///
/// This struct is use to represent these weights before they are used to randomly select a
/// pronoun.  Additional methods are provided to perform this random selection on a weighted list,
/// using as a seed both an arbitrary string of bytes and a Date.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct WeightedTable<Loot: std::cmp::Eq + std::hash::Hash>(HashMap<Loot, u8>);

impl<Loot: std::cmp::Eq + std::hash::Hash + Copy> WeightedTable<Loot> {

	/// Attempt to generate a new table
	///
	/// Accepts an iterator of tuples.  Each tupple represents one outcome.  The left
	/// element of the tupple is the value that outcome produces (the table loot, if you
	/// will), and the right element of the tuple is the weight associated with that
	/// outcome.
	///
	/// Performs a safety check to prevent creating an empty roll table.  Raises a parse
	/// error if the safety check fails.  The original tabel can be accessed as a
	/// [`HashMap`] via the [`WeightedTable::weights()`] method
	pub fn new(table: impl IntoIterator<Item=(Loot, u8)>) -> Result<Self, ParseError> {
		let table: HashMap<Loot, u8> = table.into_iter()
			.filter(|(_, weight)| *weight > 0)
			.collect();

		if table.is_empty() {
			Err(ParseError::EmptyWeightedTable)
		} else {
			Ok(WeightedTable(table))
		}
	}

	/// Access the orginal table of weights underlying this table
	///
	/// Returns a [`HashMap`] mapping pronouns to their respective weights.
	pub fn weights(&self) -> &HashMap<Loot, u8> {
		&self.0
	}

	/// Produces a table of cumulative weights for the roll table
	///
	/// Each entry has the weight of itself as well as all of the previous entries.  This
	/// is provided in a [`Vec`] ordered from lowest weight to highest.  This table can be
	/// used to efficiently perform a custom roll.
	///
	/// Two values are returned:  The table itself, and the highest value in the table
	/// (that is, the sum of all of the weights)
	pub fn rollable_table(&self) -> (Vec<(u64, Loot)>, u64) {
		let mut cumulative_weights: Vec<(u64, Loot)> = Vec::with_capacity(self.0.len());
		let mut weight_so_far = 0;
		for (pronoun, weight) in &self.0 {
			weight_so_far += *weight as u64;
			cumulative_weights.push((weight_so_far, *pronoun));
		}
		(cumulative_weights, weight_so_far)
	}

	/// A shorthand for calling [`WeightedTable::select_on_date()`] with today's date
	///
	/// The date is generated for the system's time and timezone
	pub fn select_today(&self, seed: &[u8]) -> Loot {
		self.select_on_date(
			seed,
			OffsetDateTime::now_local()
				.unwrap_or_else(|_| OffsetDateTime::now_utc())
				.date()
		)
	}

	/// Randomly select a pronoun set for a given date and name.
	///
	/// Is a wrapper for calling [`WeightedTable::select`] with the given date mixed into the seed.
	pub fn select_on_date(&self, seed: &[u8], date: Date) -> Loot {
		let mut new_seed: Vec<u8> = Vec::with_capacity(seed.len() + 4);
		new_seed.extend(
			(
				(date - COVID_EPOCH)
				.whole_days()
				as u32
			).to_le_bytes()
		);
		new_seed.extend(seed);
		self.select(seed.as_ref())
	}

	/// Randomly select a pronoun set for a given seed
	///
	/// This function is *pure*, and any randomness is produced internally using PRNG seeded with
	/// the given date and seed.  That is to say, for any given seed, this table must always
	/// produce the same pronoun set.
	pub fn select(&self, seed: &[u8]) -> Loot {
		let (rollable_table, sum_weights) = self.rollable_table();

		let mut generator: u128 = 131213121312;
		pcg64_seed(&mut generator, seed);
		let random = pcg64(&mut generator) % sum_weights;

		rollable_table.iter()
			.filter(|(weight, _)| random < *weight)
			.next()
			.expect("A table was generated with zero entries.  This should be impossible.")
			.1
	}

}