Refactor a huge amount of code!

Signed-off-by: prescientmoon <git@moonythm.dev>

src/arcaea/jacket.rs

@@ -4,10 +4,10 @@ use image::{imageops::FilterType, GenericImageView, Rgba};
 use num::Integer;
 
 use crate::{
+	arcaea::chart::{Difficulty, Jacket, SongCache},
 	assets::{get_assets_dir, should_skip_jacket_art},
-	chart::{Difficulty, Jacket, SongCache},
 	context::Error,
-	score::guess_chart_name,
+	recognition::fuzzy_song_name::guess_chart_name,
 };
 
 /// How many sub-segments to split each side into
@@ -78,7 +78,7 @@ pub struct JacketCache {
 }
 
 impl JacketCache {
-	// {{{ Generate tree
+	// {{{ Generate
 	// This is a bit inefficient (using a hash set), but only runs once
 	pub fn new(data_dir: &PathBuf, song_cache: &mut SongCache) -> Result<Self, Error> {
 		let jacket_dir = data_dir.join("jackets");

src/arcaea/mod.rs

@@ -0,0 +1,4 @@
+pub mod chart;
+pub mod jacket;
+pub mod play;
+pub mod score;

src/arcaea/play.rs

@@ -0,0 +1,371 @@
+use std::str::FromStr;
+
+use num::traits::Euclid;
+use poise::serenity_prelude::{
+	Attachment, AttachmentId, CreateAttachment, CreateEmbed, CreateEmbedAuthor, Timestamp,
+};
+use sqlx::{query_as, SqlitePool};
+
+use crate::arcaea::chart::{Chart, Song};
+use crate::context::{Error, UserContext};
+use crate::user::User;
+
+use super::score::Score;
+
+// {{{ Create play
+#[derive(Debug, Clone)]
+pub struct CreatePlay {
+	chart_id: u32,
+	user_id: u32,
+	discord_attachment_id: Option<AttachmentId>,
+
+	// Actual score data
+	score: Score,
+	zeta_score: Score,
+
+	// Optional score details
+	max_recall: Option<u32>,
+	far_notes: Option<u32>,
+
+	// Creation data
+	creation_ptt: Option<u32>,
+	creation_zeta_ptt: Option<u32>,
+}
+
+impl CreatePlay {
+	#[inline]
+	pub fn new(score: Score, chart: &Chart, user: &User) -> Self {
+		Self {
+			chart_id: chart.id,
+			user_id: user.id,
+			discord_attachment_id: None,
+			score,
+			zeta_score: score.to_zeta(chart.note_count as u32),
+			max_recall: None,
+			far_notes: None,
+			// TODO: populate these
+			creation_ptt: None,
+			creation_zeta_ptt: None,
+		}
+	}
+
+	#[inline]
+	pub fn with_attachment(mut self, attachment: &Attachment) -> Self {
+		self.discord_attachment_id = Some(attachment.id);
+		self
+	}
+
+	#[inline]
+	pub fn with_fars(mut self, far_count: Option<u32>) -> Self {
+		self.far_notes = far_count;
+		self
+	}
+
+	#[inline]
+	pub fn with_max_recall(mut self, max_recall: Option<u32>) -> Self {
+		self.max_recall = max_recall;
+		self
+	}
+
+	// {{{ Save
+	pub async fn save(self, ctx: &UserContext) -> Result<Play, Error> {
+		let attachment_id = self.discord_attachment_id.map(|i| i.get() as i64);
+		let play = sqlx::query!(
+			"
+        INSERT INTO plays(
+        user_id,chart_id,discord_attachment_id,
+        score,zeta_score,max_recall,far_notes
+        )
+        VALUES(?,?,?,?,?,?,?)
+        RETURNING id, created_at
+      ",
+			self.user_id,
+			self.chart_id,
+			attachment_id,
+			self.score.0,
+			self.zeta_score.0,
+			self.max_recall,
+			self.far_notes
+		)
+		.fetch_one(&ctx.db)
+		.await?;
+
+		Ok(Play {
+			id: play.id as u32,
+			created_at: play.created_at,
+			chart_id: self.chart_id,
+			user_id: self.user_id,
+			discord_attachment_id: self.discord_attachment_id,
+			score: self.score,
+			zeta_score: self.zeta_score,
+			max_recall: self.max_recall,
+			far_notes: self.far_notes,
+			creation_ptt: self.creation_ptt,
+			creation_zeta_ptt: self.creation_zeta_ptt,
+		})
+	}
+	// }}}
+}
+// }}}
+// {{{ DbPlay
+/// Version of `Play` matching the format sqlx expects
+#[derive(Debug, Clone, sqlx::FromRow)]
+pub struct DbPlay {
+	pub id: i64,
+	pub chart_id: i64,
+	pub user_id: i64,
+	pub discord_attachment_id: Option<String>,
+	pub score: i64,
+	pub zeta_score: i64,
+	pub max_recall: Option<i64>,
+	pub far_notes: Option<i64>,
+	pub created_at: chrono::NaiveDateTime,
+	pub creation_ptt: Option<i64>,
+	pub creation_zeta_ptt: Option<i64>,
+}
+
+impl DbPlay {
+	#[inline]
+	pub fn to_play(self) -> Play {
+		Play {
+			id: self.id as u32,
+			chart_id: self.chart_id as u32,
+			user_id: self.user_id as u32,
+			score: Score(self.score as u32),
+			zeta_score: Score(self.zeta_score as u32),
+			max_recall: self.max_recall.map(|r| r as u32),
+			far_notes: self.far_notes.map(|r| r as u32),
+			created_at: self.created_at,
+			discord_attachment_id: self
+				.discord_attachment_id
+				.and_then(|s| AttachmentId::from_str(&s).ok()),
+			creation_ptt: self.creation_ptt.map(|r| r as u32),
+			creation_zeta_ptt: self.creation_zeta_ptt.map(|r| r as u32),
+		}
+	}
+}
+// }}}
+// {{{ Play
+#[derive(Debug, Clone)]
+pub struct Play {
+	pub id: u32,
+	pub chart_id: u32,
+	pub user_id: u32,
+
+	#[allow(unused)]
+	pub discord_attachment_id: Option<AttachmentId>,
+
+	// Actual score data
+	pub score: Score,
+	pub zeta_score: Score,
+
+	// Optional score details
+	pub max_recall: Option<u32>,
+	pub far_notes: Option<u32>,
+
+	// Creation data
+	pub created_at: chrono::NaiveDateTime,
+
+	#[allow(unused)]
+	pub creation_ptt: Option<u32>,
+
+	#[allow(unused)]
+	pub creation_zeta_ptt: Option<u32>,
+}
+
+impl Play {
+	// {{{ Play => distribution
+	pub fn distribution(&self, note_count: u32) -> Option<(u32, u32, u32, u32)> {
+		if let Some(fars) = self.far_notes {
+			let (_, shinies, units) = self.score.analyse(note_count);
+			let (pures, rem) = units.checked_sub(fars)?.div_rem_euclid(&2);
+			if rem == 1 {
+				println!("The impossible happened: got an invalid amount of far notes!");
+				return None;
+			}
+
+			let lost = note_count.checked_sub(fars + pures)?;
+			let non_max_pures = pures.checked_sub(shinies)?;
+			Some((shinies, non_max_pures, fars, lost))
+		} else {
+			None
+		}
+	}
+	// }}}
+	// {{{ Play => status
+	#[inline]
+	pub fn status(&self, chart: &Chart) -> Option<String> {
+		let score = self.score.0;
+		if score >= 10_000_000 {
+			if score > chart.note_count + 10_000_000 {
+				return None;
+			}
+
+			let non_max_pures = (chart.note_count + 10_000_000).checked_sub(score)?;
+			if non_max_pures == 0 {
+				Some("MPM".to_string())
+			} else {
+				Some(format!("PM (-{})", non_max_pures))
+			}
+		} else if let Some(distribution) = self.distribution(chart.note_count) {
+			// if no lost notes...
+			if distribution.3 == 0 {
+				Some(format!("FR (-{}/-{})", distribution.1, distribution.2))
+			} else {
+				Some(format!(
+					"C (-{}/-{}/-{})",
+					distribution.1, distribution.2, distribution.3
+				))
+			}
+		} else {
+			None
+		}
+	}
+
+	#[inline]
+	pub fn short_status(&self, chart: &Chart) -> Option<char> {
+		let score = self.score.0;
+		if score >= 10_000_000 {
+			let non_max_pures = (chart.note_count + 10_000_000).checked_sub(score)?;
+			if non_max_pures == 0 {
+				Some('M')
+			} else {
+				Some('P')
+			}
+		} else if let Some(distribution) = self.distribution(chart.note_count)
+			&& distribution.3 == 0
+		{
+			Some('F')
+		} else {
+			Some('C')
+		}
+	}
+	// }}}
+	// {{{ Play to embed
+	/// Creates a discord embed for this play.
+	///
+	/// The `index` variable is only used to create distinct filenames.
+	pub async fn to_embed(
+		&self,
+		db: &SqlitePool,
+		user: &User,
+		song: &Song,
+		chart: &Chart,
+		index: usize,
+		author: Option<&poise::serenity_prelude::User>,
+	) -> Result<(CreateEmbed, Option<CreateAttachment>), Error> {
+		// {{{ Get previously best score
+		let previously_best = query_as!(
+			DbPlay,
+			"
+        SELECT * FROM plays
+        WHERE user_id=?
+        AND chart_id=?
+        AND created_at<?
+        ORDER BY score DESC
+    ",
+			user.id,
+			chart.id,
+			self.created_at
+		)
+		.fetch_optional(db)
+		.await
+		.map_err(|_| {
+			format!(
+				"Could not find any scores for {} [{:?}]",
+				song.title, chart.difficulty
+			)
+		})?
+		.map(|p| p.to_play());
+		// }}}
+
+		let attachement_name = format!("{:?}-{:?}-{:?}.png", song.id, self.score.0, index);
+		let icon_attachement = match chart.cached_jacket.as_ref() {
+			Some(jacket) => Some(CreateAttachment::bytes(jacket.raw, &attachement_name)),
+			None => None,
+		};
+
+		let mut embed = CreateEmbed::default()
+			.title(format!(
+				"{} [{:?} {}]",
+				&song.title, chart.difficulty, chart.level
+			))
+			.field("Score", format!("{} (+?)", self.score), true)
+			.field(
+				"Rating",
+				format!(
+					"{:.2} (+?)",
+					self.score.play_rating_f32(chart.chart_constant)
+				),
+				true,
+			)
+			.field("Grade", format!("{}", self.score.grade()), true)
+			.field("ξ-Score", format!("{} (+?)", self.zeta_score), true)
+			// {{{ ξ-Rating
+			.field(
+				"ξ-Rating",
+				{
+					let play_rating = self.zeta_score.play_rating_f32(chart.chart_constant);
+					if let Some(previous) = previously_best {
+						let previous_play_rating =
+							previous.zeta_score.play_rating_f32(chart.chart_constant);
+
+						if play_rating >= previous_play_rating {
+							format!(
+								"{:.2} (+{})",
+								play_rating,
+								play_rating - previous_play_rating
+							)
+						} else {
+							format!(
+								"{:.2} (-{})",
+								play_rating,
+								play_rating - previous_play_rating
+							)
+						}
+					} else {
+						format!("{:.2}", play_rating)
+					}
+				},
+				true,
+			)
+			// }}}
+			.field("ξ-Grade", format!("{}", self.zeta_score.grade()), true)
+			.field(
+				"Status",
+				self.status(chart).unwrap_or("-".to_string()),
+				true,
+			)
+			.field(
+				"Max recall",
+				if let Some(max_recall) = self.max_recall {
+					format!("{}", max_recall)
+				} else {
+					format!("-")
+				},
+				true,
+			)
+			.field("ID", format!("{}", self.id), true);
+
+		if icon_attachement.is_some() {
+			embed = embed.thumbnail(format!("attachment://{}", &attachement_name));
+		}
+
+		if let Some(user) = author {
+			let mut embed_author = CreateEmbedAuthor::new(&user.name);
+			if let Some(url) = user.avatar_url() {
+				embed_author = embed_author.icon_url(url);
+			}
+
+			embed = embed
+				.timestamp(Timestamp::from_millis(
+					self.created_at.and_utc().timestamp_millis(),
+				)?)
+				.author(embed_author);
+		}
+
+		Ok((embed, icon_attachement))
+	}
+	// }}}
+}
+// }}}

src/arcaea/score.rs

@@ -0,0 +1,348 @@
+use std::fmt::Display;
+
+use num::Rational64;
+
+use crate::context::Error;
+
+// {{{ Grade
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub enum Grade {
+	EXP,
+	EX,
+	AA,
+	A,
+	B,
+	C,
+	D,
+}
+
+impl Grade {
+	pub const GRADE_STRINGS: [&'static str; 7] = ["EX+", "EX", "AA", "A", "B", "C", "D"];
+	pub const GRADE_SHORTHANDS: [&'static str; 7] = ["exp", "ex", "aa", "a", "b", "c", "d"];
+
+	#[inline]
+	pub fn to_index(self) -> usize {
+		self as usize
+	}
+}
+
+impl Display for Grade {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "{}", Self::GRADE_STRINGS[self.to_index()])
+	}
+}
+// }}}
+// {{{ Score
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub struct Score(pub u32);
+
+impl Score {
+	// {{{ Score analysis
+	// {{{ Mini getters
+	#[inline]
+	pub fn to_zeta(self, note_count: u32) -> Score {
+		self.analyse(note_count).0
+	}
+
+	#[inline]
+	pub fn shinies(self, note_count: u32) -> u32 {
+		self.analyse(note_count).1
+	}
+
+	#[inline]
+	pub fn units(self, note_count: u32) -> u32 {
+		self.analyse(note_count).2
+	}
+	// }}}
+
+	#[inline]
+	pub fn increment(note_count: u32) -> Rational64 {
+		Rational64::new_raw(5_000_000, note_count as i64).reduced()
+	}
+
+	/// Remove the contribution made by shinies to a score.
+	#[inline]
+	pub fn forget_shinies(self, note_count: u32) -> Self {
+		Self(
+			(Self::increment(note_count) * Rational64::from_integer(self.units(note_count) as i64))
+				.floor()
+				.to_integer() as u32,
+		)
+	}
+
+	/// Compute a score without making a distinction between shinies and pures. That is, the given
+	/// value for `pures` must refer to the sum of `pure` and `shiny` notes.
+	///
+	/// This is the simplest way to compute a score, and is useful for error analysis.
+	#[inline]
+	pub fn compute_naive(note_count: u32, pures: u32, fars: u32) -> Self {
+		Self(
+			(Self::increment(note_count) * Rational64::from_integer((2 * pures + fars) as i64))
+				.floor()
+				.to_integer() as u32,
+		)
+	}
+
+	/// Returns the zeta score, the number of shinies, and the number of score units.
+	///
+	/// Pure (and higher) notes reward two score units, far notes reward one, and lost notes reward
+	/// none.
+	pub fn analyse(self, note_count: u32) -> (Score, u32, u32) {
+		// Smallest possible difference between (zeta-)scores
+		let increment = Self::increment(note_count);
+		let zeta_increment = Rational64::new_raw(2_000_000, note_count as i64).reduced();
+
+		let score = Rational64::from_integer(self.0 as i64);
+		let score_units = (score / increment).floor();
+
+		let non_shiny_score = (score_units * increment).floor();
+		let shinies = score - non_shiny_score;
+
+		let zeta_score_units = Rational64::from_integer(2) * score_units + shinies;
+		let zeta_score = Score((zeta_increment * zeta_score_units).floor().to_integer() as u32);
+
+		(
+			zeta_score,
+			shinies.to_integer() as u32,
+			score_units.to_integer() as u32,
+		)
+	}
+	// }}}
+	// {{{ Score => Play rating
+	#[inline]
+	pub fn play_rating(self, chart_constant: u32) -> i32 {
+		chart_constant as i32
+			+ if self.0 >= 10_000_000 {
+				200
+			} else if self.0 >= 9_800_000 {
+				100 + (self.0 as i32 - 9_800_000) / 2_000
+			} else {
+				(self.0 as i32 - 9_500_000) / 3_000
+			}
+	}
+
+	#[inline]
+	pub fn play_rating_f32(self, chart_constant: u32) -> f32 {
+		(self.play_rating(chart_constant)) as f32 / 100.0
+	}
+	// }}}
+	// {{{ Score => grade
+	#[inline]
+	// TODO: Perhaps make an enum for this
+	pub fn grade(self) -> Grade {
+		let score = self.0;
+		if score > 9900000 {
+			Grade::EXP
+		} else if score > 9800000 {
+			Grade::EX
+		} else if score > 9500000 {
+			Grade::AA
+		} else if score > 9200000 {
+			Grade::A
+		} else if score > 8900000 {
+			Grade::B
+		} else if score > 8600000 {
+			Grade::C
+		} else {
+			Grade::D
+		}
+	}
+	// }}}
+	// {{{ Scores & Distribution => score
+	pub fn resolve_ambiguities(
+		scores: Vec<Score>,
+		read_distribution: Option<(u32, u32, u32)>,
+		note_count: u32,
+	) -> Result<(Score, Option<u32>, Option<&'static str>), Error> {
+		if scores.len() == 0 {
+			return Err("No scores in list to disambiguate from.")?;
+		}
+
+		let mut no_shiny_scores: Vec<_> = scores
+			.iter()
+			.map(|score| score.forget_shinies(note_count))
+			.collect();
+		no_shiny_scores.sort();
+		no_shiny_scores.dedup();
+
+		if let Some(read_distribution) = read_distribution {
+			let pures = read_distribution.0;
+			let fars = read_distribution.1;
+			let losts = read_distribution.2;
+
+			// Compute score from note breakdown subpairs
+			let pf_score = Score::compute_naive(note_count, pures, fars);
+			let fl_score = Score::compute_naive(
+				note_count,
+				note_count.checked_sub(losts + fars).unwrap_or(0),
+				fars,
+			);
+			let lp_score = Score::compute_naive(
+				note_count,
+				pures,
+				note_count.checked_sub(losts + pures).unwrap_or(0),
+			);
+
+			if no_shiny_scores.len() == 1 {
+				// {{{ Score is fixed, gotta figure out the exact distribution
+				let score = *scores.iter().max().unwrap();
+
+				// {{{ Look for consensus among recomputed scores
+				// Lemma: if two computed scores agree, then so will the third
+				let consensus_fars = if pf_score == fl_score {
+					Some(fars)
+				} else {
+					// Due to the above lemma, we know all three scores must be distinct by
+					// this point.
+					//
+					// Our strategy is to check which of the three scores agrees with the real
+					// score, and to then trust the `far` value that contributed to that pair.
+					let no_shiny_score = score.forget_shinies(note_count);
+					let pf_appears = no_shiny_score == pf_score;
+					let fl_appears = no_shiny_score == fl_score;
+					let lp_appears = no_shiny_score == lp_score;
+
+					match (pf_appears, fl_appears, lp_appears) {
+						(true, false, false) => Some(fars),
+						(false, true, false) => Some(fars),
+						(false, false, true) => Some(note_count - pures - losts),
+						_ => None,
+					}
+				};
+				// }}}
+
+				if scores.len() == 1 {
+					Ok((score, consensus_fars, None))
+				} else {
+					Ok((score, consensus_fars, Some("Due to a reading error, I could not make sure the shiny-amount I calculated is accurate!")))
+				}
+
+			// }}}
+			} else {
+				// {{{ Score is not fixed, gotta figure out everything at once
+				// Some of the values in the note distribution are likely wrong (due to reading
+				// errors). To get around this, we take each pair from the triplet, compute the score
+				// it induces, and figure out if there's any consensus as to which value in the
+				// provided score list is the real one.
+				//
+				// Note that sometimes the note distribution cannot resolve any of the issues. This is
+				// usually the case when the disagreement comes from the number of shinies.
+
+				// {{{ Look for consensus among recomputed scores
+				// Lemma: if two computed scores agree, then so will the third
+				let (trusted_pure_count, consensus_computed_score, consensus_fars) = if pf_score
+					== fl_score
+				{
+					(true, pf_score, fars)
+				} else {
+					// Due to the above lemma, we know all three scores must be distinct by
+					// this point.
+					//
+					// Our strategy is to check which of the three scores appear in the
+					// provided score list.
+					let pf_appears = no_shiny_scores.contains(&pf_score);
+					let fl_appears = no_shiny_scores.contains(&fl_score);
+					let lp_appears = no_shiny_scores.contains(&lp_score);
+
+					match (pf_appears, fl_appears, lp_appears) {
+                        (true, false, false) => (true, pf_score, fars),
+                        (false, true, false) => (false, fl_score, fars),
+                        (false, false, true) => (true, lp_score, note_count - pures - losts),
+                        _ => Err(format!("Cannot disambiguate scores {:?}. Multiple disjoint note breakdown subpair scores appear on the possibility list", scores))?
+                    }
+				};
+				// }}}
+				// {{{ Collect all scores that agree with the consensus score.
+				let agreement: Vec<_> = scores
+					.iter()
+					.filter(|score| score.forget_shinies(note_count) == consensus_computed_score)
+					.filter(|score| {
+						let shinies = score.shinies(note_count);
+						shinies <= note_count && (!trusted_pure_count || shinies <= pures)
+					})
+					.map(|v| *v)
+					.collect();
+				// }}}
+				// {{{ Case 1: Disagreement in the amount of shinies!
+				if agreement.len() > 1 {
+					let agreement_shiny_amounts: Vec<_> =
+						agreement.iter().map(|v| v.shinies(note_count)).collect();
+
+					println!(
+						"Shiny count disagreement. Possible scores: {:?}. Possible shiny amounts: {:?}, Read distribution: {:?}",
+						scores, agreement_shiny_amounts, read_distribution
+					);
+
+					let msg = Some(
+                            "Due to a reading error, I could not make sure the shiny-amount I calculated is accurate!"
+                            );
+
+					Ok((
+						agreement.into_iter().max().unwrap(),
+						Some(consensus_fars),
+						msg,
+					))
+				// }}}
+				// {{{ Case 2: Total agreement!
+				} else if agreement.len() == 1 {
+					Ok((agreement[0], Some(consensus_fars), None))
+				// }}}
+				// {{{ Case 3: No agreement!
+				} else {
+					Err(format!("Could not disambiguate between possible scores {:?}. Note distribution does not agree with any possibility, leading to a score of {:?}.", scores, consensus_computed_score))?
+				}
+				// }}}
+				// }}}
+			}
+		} else {
+			if no_shiny_scores.len() == 1 {
+				if scores.len() == 1 {
+					Ok((scores[0], None, None))
+				} else {
+					Ok((scores.into_iter().max().unwrap(), None, Some("Due to a reading error, I could not make sure the shiny-amount I calculated is accurate!")))
+				}
+			} else {
+				Err("Cannot disambiguate between more than one score without a note distribution.")?
+			}
+		}
+	}
+	// }}}
+}
+
+impl Display for Score {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		let score = self.0;
+		write!(
+			f,
+			"{}'{:0>3}'{:0>3}",
+			score / 1000000,
+			(score / 1000) % 1000,
+			score % 1000
+		)
+	}
+}
+// }}}
+// {{{ Tests
+#[cfg(test)]
+mod score_tests {
+	use super::*;
+
+	#[test]
+	fn zeta_score_consistent_with_pms() {
+		// note counts
+		for note_count in 200..=2000 {
+			for shiny_count in 0..=note_count {
+				let score = Score(10000000 + shiny_count);
+				let zeta_score_units = 4 * (note_count - shiny_count) + 5 * shiny_count;
+				let (zeta_score, computed_shiny_count, units) = score.analyse(note_count);
+				let expected_zeta_score = Rational64::from_integer(zeta_score_units as i64)
+					* Rational64::new_raw(2000000, note_count as i64).reduced();
+
+				assert_eq!(zeta_score, Score(expected_zeta_score.to_integer() as u32));
+				assert_eq!(computed_shiny_count, shiny_count);
+				assert_eq!(units, 2 * note_count);
+			}
+		}
+	}
+}
+// }}}

src/assets.rs

@@ -4,7 +4,7 @@ use std::{cell::RefCell, env::var, path::PathBuf, str::FromStr, sync::OnceLock};
 use freetype::{Face, Library};
 use image::{imageops::FilterType, ImageBuffer, Rgb, Rgba};
 
-use crate::chart::Difficulty;
+use crate::arcaea::chart::Difficulty;
 
 #[inline]
 pub fn get_data_dir() -> PathBuf {

src/bitmap.rs

@@ -34,7 +34,7 @@ impl Color {
 
 	#[inline]
 	pub const fn from_bytes(bytes: [u8; 4]) -> Self {
-		Self(bytes[0], bytes[1], bytes[1], bytes[3])
+		Self(bytes[0], bytes[1], bytes[2], bytes[3])
 	}
 
 	#[inline]

src/commands/chart.rs

@@ -2,9 +2,9 @@ use poise::serenity_prelude::{CreateAttachment, CreateEmbed, CreateMessage};
 use sqlx::query;
 
 use crate::{
-	chart::Side,
+	arcaea::chart::Side,
 	context::{Context, Error},
-	score::guess_song_and_chart,
+	recognition::fuzzy_song_name::guess_song_and_chart,
 };
 
 // {{{ Chart

src/commands/mod.rs

@@ -3,6 +3,7 @@ use crate::context::{Context, Error};
 pub mod chart;
 pub mod score;
 pub mod stats;
+mod utils;
 
 // {{{ Help
 /// Show this help menu

src/commands/score.rs

@@ -1,10 +1,10 @@
-use std::fmt::Display;
-
+use crate::arcaea::play::{CreatePlay, Play};
+use crate::arcaea::score::Score;
 use crate::context::{Context, Error};
-use crate::score::{CreatePlay, ImageCropper, Play, Score, ScoreKind};
+use crate::recognition::recognize::{ImageAnalyzer, ScoreKind};
 use crate::user::{discord_it_to_discord_user, User};
-use image::imageops::FilterType;
-use poise::serenity_prelude::{CreateAttachment, CreateEmbed, CreateMessage};
+use crate::{edit_reply, get_user};
+use poise::serenity_prelude::CreateMessage;
 use poise::{serenity_prelude as serenity, CreateReply};
 use sqlx::query;
 
@@ -21,46 +21,13 @@ pub async fn score(_ctx: Context<'_>) -> Result<(), Error> {
 }
 // }}}
 // {{{ Score magic
-// {{{ Send error embed with image
-async fn error_with_image(
-	ctx: Context<'_>,
-	bytes: &[u8],
-	filename: &str,
-	message: &str,
-	err: impl Display,
-) -> Result<(), Error> {
-	let error_attachement = CreateAttachment::bytes(bytes, filename);
-	let msg = CreateMessage::default().embed(
-		CreateEmbed::default()
-			.title(message)
-			.attachment(filename)
-			.description(format!("{}", err)),
-	);
-
-	ctx.channel_id()
-		.send_files(ctx.http(), [error_attachement], msg)
-		.await?;
-
-	Ok(())
-}
-// }}}
-
 /// Identify scores from attached images.
 #[poise::command(prefix_command, slash_command)]
 pub async fn magic(
 	ctx: Context<'_>,
 	#[description = "Images containing scores"] files: Vec<serenity::Attachment>,
 ) -> Result<(), Error> {
-	let user = match User::from_context(&ctx).await {
-		Ok(user) => user,
-		Err(_) => {
-			ctx.say("You are not an user in my database, sorry!")
-				.await?;
-			return Ok(());
-		}
-	};
-
-	println!("Handling command from user {:?}", user.discord_id);
+	let user = get_user!(&ctx);
 
 	if files.len() == 0 {
 		ctx.reply("No images found attached to message").await?;
@@ -71,246 +38,103 @@ pub async fn magic(
 			.reply(format!("Processed 0/{} scores", files.len()))
 			.await?;
 
+		let mut analyzer = ImageAnalyzer::default();
+
 		for (i, file) in files.iter().enumerate() {
 			if let Some(_) = file.dimensions() {
-				// {{{ Image pre-processing
 				let bytes = file.download().await?;
 				let mut image = image::load_from_memory(&bytes)?;
 				// image = image.resize(1024, 1024, FilterType::Nearest);
-				// }}}
-				// {{{ Detection
-				// Create cropper and run OCR
-				let mut cropper = ImageCropper::default();
 
-				let edited = CreateReply::default()
-					.reply(true)
-					.content(format!("Image {}: reading jacket", i + 1));
-				handle.edit(ctx, edited).await?;
+				let result: Result<(), Error> = try {
+					// {{{ Detection
+					// This makes OCR more likely to work
+					let mut ocr_image = image.grayscale().blur(1.);
 
-				// This makes OCR more likely to work
-				let mut ocr_image = image.grayscale().blur(1.);
+					edit_reply!(ctx, handle, "Image {}: reading kind", i + 1).await?;
+					let kind = analyzer.read_score_kind(ctx.data(), &ocr_image)?;
 
-				// {{{ Kind
-				let edited = CreateReply::default()
-					.reply(true)
-					.content(format!("Image {}: reading kind", i + 1));
-				handle.edit(ctx, edited).await?;
+					edit_reply!(ctx, handle, "Image {}: reading difficulty", i + 1).await?;
+					// Do not use `ocr_image` because this reads the colors
+					let difficulty = analyzer.read_difficulty(ctx.data(), &image, kind)?;
 
-				let kind = match cropper.read_score_kind(ctx.data(), &ocr_image) {
-					// {{{ OCR error handling
-					Err(err) => {
-						error_with_image(
-							ctx,
-							&cropper.bytes,
-							&file.filename,
-							"Could not read kind from picture",
-							&err,
-						)
+					edit_reply!(ctx, handle, "Image {}: reading jacket", i + 1).await?;
+					let (song, chart) = analyzer
+						.read_jacket(ctx.data(), &mut image, kind, difficulty)
 						.await?;
 
-						continue;
-					}
-					// }}}
-					Ok(k) => k,
-				};
-				// }}}
-				// {{{ Difficulty
-				let edited = CreateReply::default()
-					.reply(true)
-					.content(format!("Image {}: reading difficulty", i + 1));
-				handle.edit(ctx, edited).await?;
+					ocr_image.invert();
 
-				// Do not use `ocr_image` because this reads the colors
-				let difficulty = match cropper.read_difficulty(ctx.data(), &image, kind) {
-					// {{{ OCR error handling
-					Err(err) => {
-						error_with_image(
-							ctx,
-							&cropper.bytes,
-							&file.filename,
-							"Could not read difficulty from picture",
-							&err,
-						)
-						.await?;
+					let (note_distribution, max_recall) = match kind {
+						ScoreKind::ScoreScreen => {
+							edit_reply!(ctx, handle, "Image {}: reading distribution", i + 1)
+								.await?;
+							let note_distribution =
+								Some(analyzer.read_distribution(ctx.data(), &image)?);
 
-						continue;
-					}
-					// }}}
-					Ok(d) => d,
-				};
+							edit_reply!(ctx, handle, "Image {}: reading max recall", i + 1).await?;
+							let max_recall = Some(analyzer.read_max_recall(ctx.data(), &image)?);
 
-				println!("{difficulty:?}");
-				// }}}
-				// {{{ Jacket & distribution
-				let mut jacket_rect = None;
-				let song_by_jacket = cropper
-					.read_jacket(ctx.data(), &mut image, kind, difficulty, &mut jacket_rect)
-					.await;
-				// image.invert();
-				ocr_image.invert();
-				let note_distribution = match kind {
-					ScoreKind::ScoreScreen => Some(cropper.read_distribution(ctx.data(), &image)?),
-					ScoreKind::SongSelect => None,
-				};
-				// }}}
-				// {{{ Title
-				let edited = CreateReply::default()
-					.reply(true)
-					.content(format!("Image {}: reading title", i + 1));
-				handle.edit(ctx, edited).await?;
-
-				let song_by_name = match kind {
-					ScoreKind::SongSelect => None,
-					ScoreKind::ScoreScreen => {
-						Some(cropper.read_song(ctx.data(), &ocr_image, difficulty))
-					}
-				};
-
-				let (song, chart) = match (song_by_jacket, song_by_name) {
-					// {{{ Only name succeeded
-					(Err(err_jacket), Some(Ok(by_name))) => {
-						println!("Could not recognise jacket with error: {}", err_jacket);
-						by_name
-					}
-					// }}}
-					// {{{ Both succeeded
-					(Ok(by_jacket), Some(Ok(by_name))) => {
-						if by_name.0.id != by_jacket.0.id {
-							println!(
-								"Got diverging choices between '{}' and '{}'",
-								by_jacket.0.title, by_name.0.title
-							);
-						};
-
-						by_jacket
-					} // }}}
-					// {{{ Only jacket succeeded
-					(Ok(by_jacket), err_name) => {
-						if let Some(err) = err_name {
-							println!("Could not read name with error: {:?}", err.unwrap_err());
+							(note_distribution, max_recall)
 						}
+						ScoreKind::SongSelect => (None, None),
+					};
 
-						by_jacket
-					}
-					// }}}
-					// {{{ Both errors
-					(Err(err_jacket), err_name) => {
-						if let Some(rect) = jacket_rect {
-							cropper.crop_image_to_bytes(&image, rect)?;
-							error_with_image(
-							ctx,
-							&cropper.bytes,
-							&file.filename,
-							"Hey! I could not read the score in the provided picture.",
-							&format!(
-                                "This can mean one of three things:
-1. The image you provided is *not that of an Arcaea score
-2. The image you provided contains a newly added chart that is not in my database yet
-3. The image you provided contains character art that covers the chart name. When this happens, I try to make use of the jacket art in order to determine the chart. Contact `@prescientmoon` on discord to try and resolve the issue!
+					edit_reply!(ctx, handle, "Image {}: reading score", i + 1).await?;
+					let score_possibilities = analyzer.read_score(
+						ctx.data(),
+						Some(chart.note_count),
+						&ocr_image,
+						kind,
+					)?;
 
-Nerdy info:
-```
-Jacket error: {}
-Title error: {:?}
-```" ,
-								err_jacket, err_name
-							),
-						)
-						.await?;
-						} else {
-							ctx.reply(format!(
-								"This is a weird error that should never happen...
-Nerdy info:
-```
-Jacket error: {}
-Title error: {:?}
-```",
-								err_jacket, err_name
-							))
-							.await?;
-						}
-						continue;
-					} // }}}
-				};
-
-				println!("{}", song.title);
-				// }}}
-				// {{{ Score
-				let edited = CreateReply::default()
-					.reply(true)
-					.content(format!("Image {}: reading score", i + 1));
-				handle.edit(ctx, edited).await?;
-
-				let score_possibilities = match cropper.read_score(
-					ctx.data(),
-					Some(chart.note_count),
-					&ocr_image,
-					kind,
-				) {
-					// {{{ OCR error handling
-					Err(err) => {
-						error_with_image(
-							ctx,
-							&cropper.bytes,
-							&file.filename,
-							"Could not read score from picture",
-							&err,
-						)
-						.await?;
-
-						continue;
-					}
-					// }}}
-					Ok(scores) => scores,
-				};
-				// }}}
-				// {{{ Build play
-				let (score, maybe_fars, score_warning) = Score::resolve_ambiguities(
-					score_possibilities,
-					note_distribution,
-					chart.note_count,
-				)
-				.map_err(|err| {
-					format!(
-						"Error occurred when disambiguating scores for '{}' [{:?}] by {}: {}",
-						song.title, difficulty, song.artist, err
+					// {{{ Build play
+					let (score, maybe_fars, score_warning) = Score::resolve_ambiguities(
+						score_possibilities,
+						note_distribution,
+						chart.note_count,
 					)
-				})?;
-				println!(
-					"Maybe fars {:?}, distribution {:?}",
-					maybe_fars, note_distribution
-				);
-				let play = CreatePlay::new(score, &chart, &user)
-					.with_attachment(file)
-					.with_fars(maybe_fars)
-					.save(&ctx.data())
-					.await?;
-				// }}}
-				// }}}
-				// {{{ Deliver embed
-				let (mut embed, attachment) = play
-					.to_embed(&ctx.data().db, &user, &song, &chart, i, None)
-					.await?;
-				if let Some(warning) = score_warning {
-					embed = embed.description(warning);
-				}
+					.map_err(|err| {
+						format!(
+							"Error occurred when disambiguating scores for '{}' [{:?}] by {}: {}",
+							song.title, difficulty, song.artist, err
+						)
+					})?;
 
-				embeds.push(embed);
-				attachments.extend(attachment);
-			// }}}
+					let play = CreatePlay::new(score, &chart, &user)
+						.with_attachment(file)
+						.with_fars(maybe_fars)
+						.with_max_recall(max_recall)
+						.save(&ctx.data())
+						.await?;
+					// }}}
+					// }}}
+					// {{{ Deliver embed
+					let (mut embed, attachment) = play
+						.to_embed(&ctx.data().db, &user, &song, &chart, i, None)
+						.await?;
+
+					if let Some(warning) = score_warning {
+						embed = embed.description(warning);
+					}
+
+					embeds.push(embed);
+					attachments.extend(attachment);
+					// }}}
+				};
+
+				if let Err(err) = result {
+					analyzer
+						.send_discord_error(ctx, &image, &file.filename, err)
+						.await?;
+				}
 			} else {
 				ctx.reply("One of the attached files is not an image!")
 					.await?;
 				continue;
 			}
 
-			let edited = CreateReply::default().reply(true).content(format!(
-				"Processed {}/{} scores",
-				i + 1,
-				files.len()
-			));
-
-			handle.edit(ctx, edited).await?;
+			edit_reply!(ctx, handle, "Processed {}/{} scores", i + 1, files.len()).await?;
 		}
 
 		handle.delete(ctx).await?;
@@ -330,14 +154,7 @@ pub async fn delete(
 	ctx: Context<'_>,
 	#[description = "Id of score to delete"] ids: Vec<u32>,
 ) -> Result<(), Error> {
-	let user = match User::from_context(&ctx).await {
-		Ok(user) => user,
-		Err(_) => {
-			ctx.say("You are not an user in my database, sorry!")
-				.await?;
-			return Ok(());
-		}
-	};
+	let user = get_user!(&ctx);
 
 	if ids.len() == 0 {
 		ctx.reply("Empty ID list provided").await?;
@@ -383,14 +200,14 @@ pub async fn show(
 	for (i, id) in ids.iter().enumerate() {
 		let res = query!(
 			"
-                SELECT 
-                    p.id,p.chart_id,p.user_id,p.score,p.zeta_score,
-                    p.max_recall,p.created_at,p.far_notes,
-                    u.discord_id
-                FROM plays p 
-                JOIN users u ON p.user_id = u.id
-                WHERE p.id=?
-            ",
+        SELECT 
+          p.id,p.chart_id,p.user_id,p.score,p.zeta_score,
+          p.max_recall,p.created_at,p.far_notes,
+          u.discord_id
+        FROM plays p 
+        JOIN users u ON p.user_id = u.id
+        WHERE p.id=?
+      ",
 			id
 		)
 		.fetch_one(&ctx.data().db)

src/commands/stats.rs

@@ -17,17 +17,20 @@ use poise::{
 use sqlx::query_as;
 
 use crate::{
+	arcaea::chart::{Chart, Song},
+	arcaea::jacket::BITMAP_IMAGE_SIZE,
+	arcaea::play::{DbPlay, Play},
+	arcaea::score::Score,
 	assets::{
 		get_b30_background, get_count_background, get_difficulty_background, get_grade_background,
 		get_name_backgound, get_ptt_emblem, get_score_background, get_status_background,
 		get_top_backgound, EXO_FONT,
 	},
 	bitmap::{Align, BitmapCanvas, Color, LayoutDrawer, LayoutManager, Rect},
-	chart::{Chart, Song},
 	context::{Context, Error},
-	jacket::BITMAP_IMAGE_SIZE,
-	score::{guess_song_and_chart, DbPlay, Play, Score},
-	user::{discord_it_to_discord_user, User},
+	get_user,
+	recognition::fuzzy_song_name::guess_song_and_chart,
+	user::discord_it_to_discord_user,
 };
 
 // {{{ Stats
@@ -63,14 +66,7 @@ pub async fn best(
 	#[description = "Name of chart to show (difficulty at the end)"]
 	name: String,
 ) -> Result<(), Error> {
-	let user = match User::from_context(&ctx).await {
-		Ok(user) => user,
-		Err(_) => {
-			ctx.say("You are not an user in my database, sorry!")
-				.await?;
-			return Ok(());
-		}
-	};
+	let user = get_user!(&ctx);
 
 	let (song, chart) = guess_song_and_chart(&ctx.data(), &name)?;
 	let play = query_as!(
@@ -121,14 +117,7 @@ pub async fn plot(
 	#[description = "Name of chart to show (difficulty at the end)"]
 	name: String,
 ) -> Result<(), Error> {
-	let user = match User::from_context(&ctx).await {
-		Ok(user) => user,
-		Err(_) => {
-			ctx.say("You are not an user in my database, sorry!")
-				.await?;
-			return Ok(());
-		}
-	};
+	let user = get_user!(&ctx);
 
 	let (song, chart) = guess_song_and_chart(&ctx.data(), &name)?;
 
@@ -240,14 +229,7 @@ pub async fn plot(
 /// Show the 30 best scores
 #[poise::command(prefix_command, slash_command)]
 pub async fn b30(ctx: Context<'_>) -> Result<(), Error> {
-	let user = match User::from_context(&ctx).await {
-		Ok(user) => user,
-		Err(_) => {
-			ctx.say("You are not an user in my database, sorry!")
-				.await?;
-			return Ok(());
-		}
-	};
+	let user = get_user!(&ctx);
 
 	let plays: Vec<DbPlay> = query_as(
 		"

src/commands/utils.rs

@@ -0,0 +1,24 @@
+#[macro_export]
+macro_rules! edit_reply {
+    ($ctx:expr, $handle:expr, $($arg:tt)*) => {{
+        let content = format!($($arg)*);
+        let edited = CreateReply::default()
+            .reply(true)
+            .content(content);
+        $handle.edit($ctx, edited)
+    }};
+}
+
+#[macro_export]
+macro_rules! get_user {
+	($ctx:expr) => {
+		match crate::user::User::from_context($ctx).await {
+			Ok(user) => user,
+			Err(_) => {
+				$ctx.say("You are not an user in my database, sorry!")
+					.await?;
+				return Ok(());
+			}
+		}
+	};
+}

src/context.rs

@@ -2,7 +2,9 @@ use std::{fs, path::PathBuf};
 
 use sqlx::SqlitePool;
 
-use crate::{chart::SongCache, jacket::JacketCache, ocr::ui::UIMeasurements};
+use crate::{
+	arcaea::chart::SongCache, arcaea::jacket::JacketCache, recognition::ui::UIMeasurements,
+};
 
 // Types used by all command functions
 pub type Error = Box<dyn std::error::Error + Send + Sync>;
@@ -12,6 +14,7 @@ pub type Context<'a> = poise::Context<'a, UserContext, Error>;
 pub struct UserContext {
 	#[allow(dead_code)]
 	pub data_dir: PathBuf,
+
 	pub db: SqlitePool,
 	pub song_cache: SongCache,
 	pub jacket_cache: JacketCache,

src/main.rs

@@ -3,17 +3,16 @@
 #![feature(let_chains)]
 #![feature(array_try_map)]
 #![feature(async_closure)]
+#![feature(try_blocks)]
 
+mod arcaea;
 mod assets;
 mod bitmap;
-mod chart;
 mod commands;
 mod context;
-mod image;
-mod jacket;
 mod levenshtein;
-mod ocr;
-mod score;
+mod recognition;
+mod transform;
 mod user;
 
 use assets::get_data_dir;

src/ocr/mod.rs

@@ -1 +0,0 @@
-pub mod ui;

src/recognition/fuzzy_song_name.rs

@@ -0,0 +1,127 @@
+use crate::arcaea::chart::{Chart, Difficulty, Song, SongCache};
+use crate::context::{Error, UserContext};
+use crate::levenshtein::edit_distance_with;
+
+/// Similar to `.strip_suffix`, but case insensitive
+#[inline]
+fn strip_case_insensitive_suffix<'a>(string: &'a str, suffix: &str) -> Option<&'a str> {
+	let suffix = suffix.to_lowercase();
+	if string.to_lowercase().ends_with(&suffix) {
+		Some(&string[0..string.len() - suffix.len()])
+	} else {
+		None
+	}
+}
+
+// {{{ Guess song and chart by name
+pub fn guess_song_and_chart<'a>(
+	ctx: &'a UserContext,
+	name: &'a str,
+) -> Result<(&'a Song, &'a Chart), Error> {
+	let name = name.trim();
+	let (name, difficulty) = name
+		.strip_suffix("PST")
+		.zip(Some(Difficulty::PST))
+		.or_else(|| strip_case_insensitive_suffix(name, "[PST]").zip(Some(Difficulty::PST)))
+		.or_else(|| strip_case_insensitive_suffix(name, "PRS").zip(Some(Difficulty::PRS)))
+		.or_else(|| strip_case_insensitive_suffix(name, "[PRS]").zip(Some(Difficulty::PRS)))
+		.or_else(|| strip_case_insensitive_suffix(name, "FTR").zip(Some(Difficulty::FTR)))
+		.or_else(|| strip_case_insensitive_suffix(name, "[FTR]").zip(Some(Difficulty::FTR)))
+		.or_else(|| strip_case_insensitive_suffix(name, "ETR").zip(Some(Difficulty::ETR)))
+		.or_else(|| strip_case_insensitive_suffix(name, "[ETR]").zip(Some(Difficulty::ETR)))
+		.or_else(|| strip_case_insensitive_suffix(name, "BYD").zip(Some(Difficulty::BYD)))
+		.or_else(|| strip_case_insensitive_suffix(name, "[BYD]").zip(Some(Difficulty::BYD)))
+		.unwrap_or((&name, Difficulty::FTR));
+
+	guess_chart_name(name, &ctx.song_cache, Some(difficulty), true)
+}
+// }}}
+// {{{ Guess chart by name
+/// Runs a specialized fuzzy-search through all charts in the game.
+///
+/// The `unsafe_heuristics` toggle increases the amount of resolvable queries, but might let in
+/// some false positives. We turn it on for simple user-search commands, but disallow it for things
+/// like OCR-generated text.
+pub fn guess_chart_name<'a>(
+	raw_text: &str,
+	cache: &'a SongCache,
+	difficulty: Option<Difficulty>,
+	unsafe_heuristics: bool,
+) -> Result<(&'a Song, &'a Chart), Error> {
+	let raw_text = raw_text.trim(); // not quite raw 🤔
+	let mut text: &str = &raw_text.to_lowercase();
+
+	// Cached vec used by the levenshtein distance function
+	let mut levenshtein_vec = Vec::with_capacity(20);
+	// Cached vec used to store distance calculations
+	let mut distance_vec = Vec::with_capacity(3);
+
+	let (song, chart) = loop {
+		let mut close_enough: Vec<_> = cache
+			.songs()
+			.filter_map(|item| {
+				let song = &item.song;
+				let chart = if let Some(difficulty) = difficulty {
+					item.lookup(difficulty).ok()?
+				} else {
+					item.charts().next()?
+				};
+
+				let song_title = &song.lowercase_title;
+				distance_vec.clear();
+
+				let base_distance = edit_distance_with(&text, &song_title, &mut levenshtein_vec);
+				if base_distance < 1.max(song.title.len() / 3) {
+					distance_vec.push(base_distance * 10 + 2);
+				}
+
+				let shortest_len = Ord::min(song_title.len(), text.len());
+				if let Some(sliced) = &song_title.get(..shortest_len)
+					&& (text.len() >= 6 || unsafe_heuristics)
+				{
+					let slice_distance = edit_distance_with(&text, sliced, &mut levenshtein_vec);
+					if slice_distance < 1 {
+						distance_vec.push(slice_distance * 10 + 3);
+					}
+				}
+
+				if let Some(shorthand) = &chart.shorthand
+					&& unsafe_heuristics
+				{
+					let short_distance = edit_distance_with(&text, shorthand, &mut levenshtein_vec);
+					if short_distance < 1.max(shorthand.len() / 3) {
+						distance_vec.push(short_distance * 10 + 1);
+					}
+				}
+
+				distance_vec
+					.iter()
+					.min()
+					.map(|distance| (song, chart, *distance))
+			})
+			.collect();
+
+		if close_enough.len() == 0 {
+			if text.len() <= 1 {
+				Err(format!(
+					"Could not find match for chart name '{}' [{:?}]",
+					raw_text, difficulty
+				))?;
+			} else {
+				text = &text[..text.len() - 1];
+			}
+		} else if close_enough.len() == 1 {
+			break (close_enough[0].0, close_enough[0].1);
+		} else {
+			if unsafe_heuristics {
+				close_enough.sort_by_key(|(_, _, distance)| *distance);
+				break (close_enough[0].0, close_enough[0].1);
+			} else {
+				return Err(format!("Name '{}' is too vague to choose a match", raw_text).into());
+			};
+		};
+	};
+
+	Ok((song, chart))
+}
+// }}}

src/recognition/mod.rs

@@ -0,0 +1,3 @@
+pub mod fuzzy_song_name;
+pub mod recognize;
+pub mod ui;

src/recognition/recognize.rs

@@ -0,0 +1,495 @@
+use std::fmt::Display;
+use std::io::Cursor;
+use std::str::FromStr;
+use std::{env, fs};
+
+use hypertesseract::{PageSegMode, Tesseract};
+use image::{DynamicImage, GenericImageView};
+use image::{ImageBuffer, Rgba};
+use num::integer::Roots;
+use poise::serenity_prelude::{CreateAttachment, CreateEmbed, CreateMessage, Timestamp};
+
+use crate::arcaea::chart::{Chart, Difficulty, Song, DIFFICULTY_MENU_PIXEL_COLORS};
+use crate::arcaea::jacket::IMAGE_VEC_DIM;
+use crate::arcaea::score::Score;
+use crate::bitmap::{Color, Rect};
+use crate::context::{Context, Error, UserContext};
+use crate::levenshtein::edit_distance;
+use crate::recognition::fuzzy_song_name::guess_chart_name;
+use crate::recognition::ui::{
+	ScoreScreenRect, SongSelectRect, UIMeasurementRect, UIMeasurementRect::*,
+};
+use crate::transform::rotate;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ScoreKind {
+	SongSelect,
+	ScoreScreen,
+}
+
+/// Caches a byte vector in order to prevent reallocation
+#[derive(Debug, Clone, Default)]
+pub struct ImageAnalyzer {
+	/// cached byte array
+	pub bytes: Vec<u8>,
+
+	/// Last rect used to crop something
+	last_rect: Option<(UIMeasurementRect, Rect)>,
+}
+
+impl ImageAnalyzer {
+	/// Similar to reinitializing this, but without deallocating memory
+	#[inline]
+	pub fn clear(&mut self) {
+		self.bytes.clear();
+		self.last_rect = None;
+	}
+
+	// {{{ Crop
+	pub fn crop_image_to_bytes(&mut self, image: &DynamicImage, rect: Rect) -> Result<(), Error> {
+		self.clear();
+		let image = image.crop_imm(rect.x as u32, rect.y as u32, rect.width, rect.height);
+		let mut cursor = Cursor::new(&mut self.bytes);
+		image.write_to(&mut cursor, image::ImageFormat::Png)?;
+
+		fs::write(format!("./logs/{}.png", Timestamp::now()), &self.bytes)?;
+
+		Ok(())
+	}
+
+	#[inline]
+	pub fn crop(&mut self, image: &DynamicImage, rect: Rect) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
+		if env::var("SHIMMERING_DEBUG_IMGS")
+			.map(|s| s == "1")
+			.unwrap_or(false)
+		{
+			self.crop_image_to_bytes(image, rect).unwrap();
+		}
+
+		image
+			.crop_imm(rect.x as u32, rect.y as u32, rect.width, rect.height)
+			.to_rgba8()
+	}
+
+	#[inline]
+	pub fn interp_crop(
+		&mut self,
+		ctx: &UserContext,
+		image: &DynamicImage,
+		ui_rect: UIMeasurementRect,
+	) -> Result<ImageBuffer<Rgba<u8>, Vec<u8>>, Error> {
+		let rect = ctx.ui_measurements.interpolate(ui_rect, image)?;
+		self.last_rect = Some((ui_rect, rect));
+		Ok(self.crop(image, rect))
+	}
+	// }}}
+	// {{{ Error handling
+	pub async fn send_discord_error(
+		&mut self,
+		ctx: Context<'_>,
+		image: &DynamicImage,
+		filename: &str,
+		err: impl Display,
+	) -> Result<(), Error> {
+		let mut embed = CreateEmbed::default().description(format!(
+			"Nerdy info
+```
+{}
+```",
+			err
+		));
+
+		if let Some((ui_rect, rect)) = self.last_rect {
+			self.crop_image_to_bytes(image, rect)?;
+
+			let bytes = std::mem::take(&mut self.bytes);
+			let error_attachement = CreateAttachment::bytes(bytes, filename);
+
+			embed = embed.attachment(filename).title(format!(
+				"An error occurred, around the time I was extracting data for {ui_rect:?}"
+			));
+
+			let msg = CreateMessage::default().embed(embed);
+			ctx.channel_id()
+				.send_files(ctx.http(), [error_attachement], msg)
+				.await?;
+		} else {
+			embed = embed.title("An error occurred");
+
+			let msg = CreateMessage::default().embed(embed);
+			ctx.channel_id().send_files(ctx.http(), [], msg).await?;
+		}
+
+		Ok(())
+	}
+	// }}}
+	// {{{ Read score
+	pub fn read_score(
+		&mut self,
+		ctx: &UserContext,
+		note_count: Option<u32>,
+		image: &DynamicImage,
+		kind: ScoreKind,
+	) -> Result<Vec<Score>, Error> {
+		let image = self.interp_crop(
+			ctx,
+			image,
+			if kind == ScoreKind::ScoreScreen {
+				ScoreScreen(ScoreScreenRect::Score)
+			} else {
+				SongSelect(SongSelectRect::Score)
+			},
+		)?;
+
+		let mut results = vec![];
+		for mode in [
+			PageSegMode::SingleWord,
+			PageSegMode::RawLine,
+			PageSegMode::SingleLine,
+			PageSegMode::SparseText,
+			PageSegMode::SingleBlock,
+		] {
+			let result: Result<_, Error> = try {
+				// {{{ Read score using tesseract
+				let text = Tesseract::builder()
+					.language(hypertesseract::Language::English)
+					.whitelist_str("0123456789'/")?
+					.page_seg_mode(mode)
+					.assume_numeric_input()
+					.build()?
+					.load_image(&image)?
+					.recognize()?
+					.get_text()?;
+
+				let text: String = text
+					.trim()
+					.chars()
+					.map(|char| if char == '/' { '7' } else { char })
+					.filter(|char| *char != ' ' && *char != '\'')
+					.collect();
+
+				let score = u32::from_str_radix(&text, 10)?;
+				Score(score)
+				// }}}
+			};
+
+			match result {
+				Ok(result) => {
+					results.push(result.0);
+				}
+				Err(err) => {
+					println!("OCR score result error: {}", err);
+				}
+			}
+		}
+
+		// {{{ Score correction
+		// The OCR sometimes fails to read "74" with the arcaea font,
+		// so we try to detect that and fix it
+		loop {
+			let old_stack_len = results.len();
+			println!("Results {:?}", results);
+			results = results
+				.iter()
+				.flat_map(|result| {
+					// If the length is correct, we are good to go!
+					if *result >= 8_000_000 {
+						vec![*result]
+					} else {
+						let mut results = vec![];
+						for i in [0, 1, 3, 4] {
+							let d = 10u32.pow(i);
+							if (*result / d) % 10 == 4 && (*result / d) % 100 != 74 {
+								let n = d * 10;
+								results.push((*result / n) * n * 10 + 7 * n + (*result % n));
+							}
+						}
+
+						results
+					}
+				})
+				.collect();
+
+			if old_stack_len == results.len() {
+				break;
+			}
+		}
+		// }}}
+		// {{{ Return score if consensus exists
+		// 1. Discard scores that are known to be impossible
+		let mut results: Vec<_> = results
+			.into_iter()
+			.filter(|result| {
+				8_000_000 <= *result
+					&& *result <= 10_010_000
+					&& note_count
+						.map(|note_count| {
+							let (zeta, shinies, score_units) = Score(*result).analyse(note_count);
+							8_000_000 <= zeta.0
+								&& zeta.0 <= 10_000_000 && shinies <= note_count
+								&& score_units <= 2 * note_count
+						})
+						.unwrap_or(true)
+			})
+			.map(|r| Score(r))
+			.collect();
+		println!("Results {:?}", results);
+
+		// 2. Look for consensus
+		for result in results.iter() {
+			if results.iter().filter(|e| **e == *result).count() > results.len() / 2 {
+				return Ok(vec![*result]);
+			}
+		}
+		// }}}
+
+		// If there's no consensus, we return everything
+		results.sort();
+		results.dedup();
+		println!("Results {:?}", results);
+
+		Ok(results)
+	}
+	// }}}
+	// {{{ Read difficulty
+	pub fn read_difficulty(
+		&mut self,
+		ctx: &UserContext,
+		image: &DynamicImage,
+		kind: ScoreKind,
+	) -> Result<Difficulty, Error> {
+		if kind == ScoreKind::SongSelect {
+			let min = DIFFICULTY_MENU_PIXEL_COLORS
+				.iter()
+				.zip(Difficulty::DIFFICULTIES)
+				.min_by_key(|(c, d)| {
+					let rect = ctx
+						.ui_measurements
+						.interpolate(
+							SongSelect(match d {
+								Difficulty::PST => SongSelectRect::Past,
+								Difficulty::PRS => SongSelectRect::Present,
+								Difficulty::FTR => SongSelectRect::Future,
+								_ => SongSelectRect::Beyond,
+							}),
+							image,
+						)
+						.unwrap();
+
+					// rect.width = 100;
+					// rect.height = 100;
+					// self.crop_image_to_bytes(image, rect).unwrap();
+
+					let image_color = image.get_pixel(rect.x as u32, rect.y as u32);
+					let image_color = Color::from_bytes(image_color.0);
+
+					let distance = c.distance(image_color);
+					(distance * 10000.0) as u32
+				})
+				.unwrap();
+
+			return Ok(min.1);
+		}
+
+		let mut ocr = Tesseract::builder()
+			.language(hypertesseract::Language::English)
+			.page_seg_mode(PageSegMode::RawLine)
+			.build()?;
+
+		ocr.load_image(&self.interp_crop(ctx, image, ScoreScreen(ScoreScreenRect::Difficulty))?)?
+			.recognize()?;
+
+		let text: &str = &ocr.get_text()?;
+		let text = text.trim().to_lowercase();
+
+		// let conf = t.mean_text_conf();
+		// if conf < 10 && conf != 0 {
+		// 	Err(format!(
+		// 		"Difficulty text is not readable (confidence = {}, text = {}).",
+		// 		conf, text
+		// 	))?;
+		// }
+
+		let difficulty = Difficulty::DIFFICULTIES
+			.iter()
+			.zip(Difficulty::DIFFICULTY_STRINGS)
+			.min_by_key(|(_, difficulty_string)| edit_distance(difficulty_string, &text))
+			.map(|(difficulty, _)| *difficulty)
+			.ok_or_else(|| format!("Unrecognised difficulty '{}'", text))?;
+
+		Ok(difficulty)
+	}
+	// }}}
+	// {{{ Read score kind
+	pub fn read_score_kind(
+		&mut self,
+		ctx: &UserContext,
+		image: &DynamicImage,
+	) -> Result<ScoreKind, Error> {
+		let text = Tesseract::builder()
+			.language(hypertesseract::Language::English)
+			.page_seg_mode(PageSegMode::RawLine)
+			.build()?
+			.load_image(&self.interp_crop(ctx, image, PlayKind)?)?
+			.recognize()?
+			.get_text()?
+			.trim()
+			.to_string();
+
+		// let conf = t.mean_text_conf();
+		// if conf < 10 && conf != 0 {
+		// 	Err(format!(
+		// 		"Score kind text is not readable (confidence = {}, text = {}).",
+		// 		conf, text
+		// 	))?;
+		// }
+
+		let result = if edit_distance(&text, "Result") < edit_distance(&text, "Select a song") {
+			ScoreKind::ScoreScreen
+		} else {
+			ScoreKind::SongSelect
+		};
+
+		Ok(result)
+	}
+	// }}}
+	// {{{ Read song
+	pub fn read_song<'a>(
+		&mut self,
+		ctx: &'a UserContext,
+		image: &DynamicImage,
+		difficulty: Difficulty,
+	) -> Result<(&'a Song, &'a Chart), Error> {
+		let text = Tesseract::builder()
+			.language(hypertesseract::Language::English)
+			.page_seg_mode(PageSegMode::SingleLine)
+			.whitelist_str("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789,.()- ")?
+			.build()?
+			.load_image(&self.interp_crop(ctx, image, ScoreScreen(ScoreScreenRect::Title))?)?
+			.recognize()?
+			.get_text()?;
+
+		// let conf = t.mean_text_conf();
+		// if conf < 20 && conf != 0 {
+		// 	Err(format!(
+		// 		"Title text is not readable (confidence = {}, text = {}).",
+		// 		conf,
+		// 		raw_text.trim()
+		// 	))?;
+		// }
+
+		guess_chart_name(&text, &ctx.song_cache, Some(difficulty), false)
+	}
+	// }}}
+	// {{{ Read jacket
+	pub async fn read_jacket<'a>(
+		&mut self,
+		ctx: &'a UserContext,
+		image: &mut DynamicImage,
+		kind: ScoreKind,
+		difficulty: Difficulty,
+	) -> Result<(&'a Song, &'a Chart), Error> {
+		let rect = ctx.ui_measurements.interpolate(
+			if kind == ScoreKind::ScoreScreen {
+				ScoreScreen(ScoreScreenRect::Jacket)
+			} else {
+				SongSelect(SongSelectRect::Jacket)
+			},
+			image,
+		)?;
+
+		let cropped = if kind == ScoreKind::ScoreScreen {
+			image.view(rect.x as u32, rect.y as u32, rect.width, rect.height)
+		} else {
+			let angle = f32::atan2(rect.height as f32, rect.width as f32);
+			let side = rect.height + rect.width;
+			rotate(
+				image,
+				Rect::new(rect.x, rect.y, side, side),
+				(rect.x, rect.y + rect.height as i32),
+				angle,
+			);
+
+			let len = (rect.width.pow(2) + rect.height.pow(2)).sqrt();
+
+			image.view(rect.x as u32, rect.y as u32 + rect.height, len, len)
+		};
+		let (distance, song_id) = ctx
+			.jacket_cache
+			.recognise(&*cropped)
+			.ok_or_else(|| "Could not recognise jacket")?;
+
+		if distance > (IMAGE_VEC_DIM * 3) as f32 {
+			Err("No known jacket looks like this")?;
+		}
+
+		let item = ctx.song_cache.lookup(*song_id)?;
+		let chart = item.lookup(difficulty)?;
+
+		// NOTE: this will reallocate a few strings, but it is what it is
+		Ok((&item.song, chart))
+	}
+	// }}}
+	// {{{ Read distribution
+	pub fn read_distribution(
+		&mut self,
+		ctx: &UserContext,
+		image: &DynamicImage,
+	) -> Result<(u32, u32, u32), Error> {
+		let mut ocr = Tesseract::builder()
+			.language(hypertesseract::Language::English)
+			.page_seg_mode(PageSegMode::SparseText)
+			.whitelist_str("0123456789")?
+			.assume_numeric_input()
+			.build()?;
+
+		let mut out = [0; 3];
+
+		use ScoreScreenRect::*;
+		static KINDS: [ScoreScreenRect; 3] = [Pure, Far, Lost];
+
+		for i in 0..3 {
+			let text = ocr
+				.load_image(&self.interp_crop(ctx, image, ScoreScreen(KINDS[i]))?)?
+				.recognize()?
+				.get_text()?;
+
+			println!("Raw '{}'", text.trim());
+			out[i] = u32::from_str(&text.trim()).unwrap_or(0);
+		}
+		println!("Ditribution {out:?}");
+
+		Ok((out[0], out[1], out[2]))
+	}
+	// }}}
+	// {{{ Read max recall
+	pub fn read_max_recall<'a>(
+		&mut self,
+		ctx: &'a UserContext,
+		image: &DynamicImage,
+	) -> Result<u32, Error> {
+		let text = Tesseract::builder()
+			.language(hypertesseract::Language::English)
+			.page_seg_mode(PageSegMode::SingleLine)
+			.whitelist_str("0123456789")?
+			.assume_numeric_input()
+			.build()?
+			.load_image(&self.interp_crop(ctx, image, ScoreScreen(ScoreScreenRect::MaxRecall))?)?
+			.recognize()?
+			.get_text()?;
+
+		let max_recall = u32::from_str_radix(text.trim(), 10)?;
+
+		// let conf = t.mean_text_conf();
+		// if conf < 20 && conf != 0 {
+		// 	Err(format!(
+		// 		"Title text is not readable (confidence = {}, text = {}).",
+		// 		conf,
+		// 		raw_text.trim()
+		// 	))?;
+		// }
+
+		Ok(max_recall)
+	}
+	// }}}
+}

src/recognition/ui.rs

@@ -1,5 +1,3 @@
-#![allow(dead_code)]
-
 use std::{fs, path::PathBuf};
 
 use image::GenericImage;