use itertools::Itertools;
use std::{collections::HashMap, fmt::Display};
#[allow(non_upper_case_globals)]
const color_green: &str = "\x1B[92m";
#[allow(non_upper_case_globals)]
const color_yellow: &str = "\x1B[93m";
#[allow(non_upper_case_globals)]
const color_grey: &str = "\x1B[37m";
#[allow(non_upper_case_globals)]
pub const color_reset: &str = "\x1B[39m";
#[derive(Clone)]
pub struct Counts {
pub data: HashMap<String, i32>,
pub shots: usize,
}
impl Display for Counts {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut pairs = self
.data
.iter()
.map(|(x, y)| (x.clone(), *y))
.collect::<Vec<(String, i32)>>();
pairs.sort_by(|a, b| {
let ai = i32::from_str_radix(&a.0, 10).unwrap();
let bi = i32::from_str_radix(&b.0, 10).unwrap();
ai.cmp(&bi)
});
let joined = pairs
.iter()
.map(|(st, va)| {
let p = (*va) as f32 / self.shots as f32;
if p > 0.25 {
format!("{color_green}{}|{}⟩{color_reset}", p, st)
} else if p > 0.1 {
format!("{color_yellow}{}|{}⟩{color_reset}", p, st)
} else {
format!("{color_grey}{}|{}⟩{color_reset}", p, st)
}
})
.join(" + ");
write!(f, "|ψ⟩ = {}", joined)
}
}
impl Counts {
pub fn print_psi(&self) {
println!("{}", self);
}
}