p3l/stdlib/src/output/mod.rs

#![allow(clippy::new_without_default)]

use crate::int;
use indexmap::IndexMap;
use std::{
	error::Error,
	fmt::{self, Display, Formatter},
	fs::File,
	io::{self, stdout, Write},
	ops::{Deref, DerefMut},
	path::PathBuf,
	str::FromStr
};
use svgwriter::{tags::TagWithGlobalEventAttributes, Graphic, Value};

mod bar;
mod bubble;
mod list;

pub use bar::BarChartPrinter;
pub use bubble::BubbleChartPrinter;
pub use list::ListPrinter;

#[derive(Debug)]
pub struct NoSuchVariant(String);

impl Display for NoSuchVariant {
	fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
		write!(f, "No such variant: {}", self.0)
	}
}

impl Error for NoSuchVariant {}

/// Used for argument parsing
#[derive(Clone, Copy, Debug, Default)]
pub enum OutputPrinterAny {
	#[default]
	List,
	Bar
}

impl FromStr for OutputPrinterAny {
	type Err = NoSuchVariant;

	fn from_str(s: &str) -> Result<Self, Self::Err> {
		match s {
			"list" => Ok(Self::List),
			"bar" => Ok(Self::Bar),
			_ => Err(NoSuchVariant(s.to_owned()))
		}
	}
}

/// Used for argument parsing
#[derive(Clone, Copy, Debug, Default)]
pub enum OutputPrinterTwoVars {
	#[default]
	List,
	Bar,
	Bubble
}

impl FromStr for OutputPrinterTwoVars {
	type Err = NoSuchVariant;

	fn from_str(s: &str) -> Result<Self, Self::Err> {
		match s {
			"list" => Ok(Self::List),
			"bar" => Ok(Self::Bar),
			"bubble" => Ok(Self::Bubble),
			_ => Err(NoSuchVariant(s.to_owned()))
		}
	}
}

/// Used for output printing.
pub trait OutputPrinter {
	fn print<I, E, W>(self, list: OutputList<I>, out: W) -> io::Result<()>
	where
		I: IntoIterator<Item = (E, usize)>,
		E: IntoIterator<Item = (&'static str, int)>,
		W: Write;
}

/// Used for output printing.
pub struct OutputList<I> {
	/// An iterator over all samples. Must yield pairs of (values, qty),
	/// where values is an iterator over (variable, value).
	pub samples: I,
	/// The total quantity of all samples.
	pub n: usize,
	/// The minimum of each output variable.
	min: IndexMap<&'static str, int>,
	/// The maximum of each output variable.
	max: IndexMap<&'static str, int>,
	/// The mean of each output variable.
	mean: IndexMap<&'static str, f64>,
	/// The variance of each output variable.
	variance: IndexMap<&'static str, f64>
}

#[derive(Debug)]
pub struct Stats {
	pub min: int,
	pub max: int,
	pub mean: f64,
	pub variance: f64,
	pub standard_deviation: f64
}

impl<I, E> OutputList<I>
where
	I: IntoIterator<Item = (E, usize)>,
	for<'a> &'a I: IntoIterator<Item = (&'a E, &'a usize)>,
	E: IntoIterator<Item = (&'static str, int)>,
	for<'a> &'a E: IntoIterator<Item = (&'static str, int)>
{
	pub fn new(samples: I, n: usize) -> Self {
		let mut min = IndexMap::new();
		let mut max = IndexMap::new();
		let mut mean = IndexMap::new();
		let mut qty = 0;
		for (sample, sample_qty) in &samples {
			let new_qty = (qty + sample_qty) as f64;
			for (var, value) in sample {
				if min.get(var).map(|v| *v > value).unwrap_or(true) {
					min.insert(var, value);
				}
				if max.get(var).map(|v| *v < value).unwrap_or(true) {
					max.insert(var, value);
				}
				mean.insert(
					var,
					mean.get(var).copied().unwrap_or(0.0) * (qty as f64 / new_qty)
						+ value as f64 * (*sample_qty as f64 / new_qty)
				);
			}
			qty += sample_qty;
		}
		debug_assert_eq!(qty, n);

		let mut variance = IndexMap::new();
		let mut qty;
		let mut samples_iter = (&samples).into_iter();
		let (first_sample, first_sample_qty) =
			samples_iter.next().expect("I need at least one sample");
		if *first_sample_qty > 1 {
			qty = *first_sample_qty;
			for (var, value) in first_sample {
				variance.insert(
					var,
					(value as f64 - mean[var]).powi(2) * (qty as f64 / (qty - 1) as f64)
				);
			}
		} else {
			let (second_sample, second_sample_qty) =
				samples_iter.next().expect("I need at least two samples");
			let first_sample: IndexMap<&'static str, int> =
				first_sample.into_iter().collect();
			qty = first_sample_qty + second_sample_qty;
			for (var, second_value) in second_sample {
				let first_value = first_sample[var];
				variance.insert(
					var,
					(first_value as f64 - mean[var]).powi(2)
						* (*first_sample_qty as f64 / (qty - 1) as f64)
						+ (second_value as f64 - mean[var]).powi(2)
							* (*second_sample_qty as f64 / (qty - 1) as f64)
				);
			}
		}
		for (sample, sample_qty) in samples_iter {
			let new_qty = (qty + sample_qty) as f64;
			for (var, value) in sample {
				variance.insert(
					var,
					variance[var] * ((qty - 1) as f64 / (new_qty - 1.0))
						+ (value as f64 - mean[var]).powi(2)
							* (*sample_qty as f64 / (new_qty - 1.0))
				);
			}
			qty += sample_qty;
		}
		debug_assert_eq!(qty, n);

		Self {
			samples,
			n,
			min,
			max,
			mean,
			variance
		}
	}
}

impl<I> OutputList<I> {
	pub fn min(&self, var: &'static str) -> int {
		self.min.get(var).copied().unwrap_or(0)
	}

	pub fn max(&self, var: &'static str) -> int {
		self.max.get(var).copied().unwrap_or(0)
	}

	pub fn mean(&self, var: &'static str) -> f64 {
		self.mean.get(var).copied().unwrap_or(0.0)
	}

	pub fn variance(&self, var: &'static str) -> f64 {
		self.variance.get(var).copied().unwrap_or(f64::INFINITY)
	}

	pub fn stats(&self) -> impl Iterator<Item = (&'static str, Stats)> + '_ {
		self.mean.iter().map(|(key, mean)| {
			let variance = self.variance(key);
			(*key, Stats {
				min: self.min(key),
				max: self.max(key),
				mean: *mean,
				variance,
				standard_deviation: variance.sqrt()
			})
		})
	}
}

/// Used for output printing.
pub trait OutputTarget {
	fn print<I, E>(self, list: OutputList<I>)
	where
		I: IntoIterator<Item = (E, usize)>,
		E: IntoIterator<Item = (&'static str, int)>;
}

fn print_with<I, E, P>(list: OutputList<I>, printer: P, path: Option<PathBuf>)
where
	I: IntoIterator<Item = (E, usize)>,
	E: IntoIterator<Item = (&'static str, int)>,
	P: OutputPrinter
{
	match path {
		Some(path) => {
			let file = File::create(&path).expect("Failed to create output file");
			printer.print(list, file).expect("Failed to print output");
			eprintln!("Wrote output to {}", path.display());
		},
		None => {
			printer
				.print(list, stdout())
				.expect("Failed to print output");
		}
	}
}

impl OutputTarget for (OutputPrinterAny, Option<PathBuf>) {
	fn print<I, E>(self, list: OutputList<I>)
	where
		I: IntoIterator<Item = (E, usize)>,
		E: IntoIterator<Item = (&'static str, int)>
	{
		match self.0 {
			OutputPrinterAny::List => print_with(list, ListPrinter, self.1),
			OutputPrinterAny::Bar => {
				print_with(list, BarChartPrinter::new(self.1.is_none()), self.1)
			},
		}
	}
}

impl OutputTarget for (OutputPrinterTwoVars, Option<PathBuf>) {
	fn print<I, E>(self, list: OutputList<I>)
	where
		I: IntoIterator<Item = (E, usize)>,
		E: IntoIterator<Item = (&'static str, int)>
	{
		match self.0 {
			OutputPrinterTwoVars::List => print_with(list, ListPrinter, self.1),
			OutputPrinterTwoVars::Bar => {
				print_with(list, BarChartPrinter::new(self.1.is_none()), self.1)
			},
			OutputPrinterTwoVars::Bubble => {
				print_with(list, BubbleChartPrinter::new(self.1.is_none()), self.1)
			},
		}
	}
}

/// Helper for creating svgs
struct Svg(Graphic);

impl Deref for Svg {
	type Target = Graphic;

	fn deref(&self) -> &Self::Target {
		&self.0
	}
}

impl DerefMut for Svg {
	fn deref_mut(&mut self) -> &mut Self::Target {
		&mut self.0
	}
}

impl Svg {
	fn new(width: i32, height: i32) -> Self {
		let mut svg = Graphic::new();
		svg.set_view_box(format!("0 0 {width} {height}"));
		svg.set_height("20cm");
		Self(svg)
	}

	fn with_onload<V>(mut self, onload: V) -> Self
	where
		V: Value + 'static
	{
		self.0.set_onload(onload);
		self
	}
}