1use crate::labels::{self, Copy, Label};
14use crate::spamcore::{N, NAMES};
15use crate::{Dirs, LearnError, corpus, train};
16use serde::Serialize;
17use std::collections::{BTreeMap, HashMap};
18
19#[derive(Debug, Clone, PartialEq, Serialize)]
22pub struct Group {
23 pub account: String,
24 pub label: Label,
25 pub messages: u64,
26 pub means: Vec<Option<f64>>,
28 pub known: Vec<u64>,
29}
30
31pub fn names() -> &'static [&'static str] {
33 &NAMES
34}
35
36pub fn feature_means(dirs: &Dirs, trusted: &BTreeMap<String, Vec<String>>, domain: Option<&str>) -> Result<Vec<Group>, LearnError> {
40 let mut copies = Vec::new();
41 corpus::read_all(dirs, |record| copies.push(Copy::of(&record)))?;
42 let (labeled, _) = labels::decide(copies, &labels::read_log(dirs), &labels::read_moved(dirs), &labels::read_flagged(dirs));
43 let chosen: HashMap<corpus::Place, Label> = labeled.into_iter().map(|l| (l.place, l.label)).collect();
44 let domain = domain.map(|d| d.trim().trim_start_matches('@').to_ascii_lowercase());
45 let mut sums: BTreeMap<(String, bool), (u64, [f64; N], [u64; N])> = BTreeMap::new();
47 corpus::read_all(dirs, |record| {
48 let Some(&label) = chosen.get(&record.place()) else { return };
49 if let Some(domain) = &domain {
50 let from = labels::from_address(&record.header_bytes()).and_then(|a| a.rsplit_once('@').map(|(_, d)| d.to_string())).unwrap_or_default();
51 if !(from == *domain || from.ends_with(&format!(".{domain}"))) {
52 return;
53 }
54 }
55 let (_, features) = train::read_message(&record, trusted, &crate::spamcore::Lexicon::builtin());
57 let group = sums.entry((record.account.clone(), label == Label::Spam)).or_insert((0, [0.0; N], [0; N]));
58 group.0 += 1;
59 for (h, x) in features.iter().enumerate() {
60 if !x.is_nan() {
61 group.1[h] += f64::from(*x);
62 group.2[h] += 1;
63 }
64 }
65 })?;
66 Ok(sums
67 .into_iter()
68 .map(|((account, spam), (messages, sum, known))| Group {
69 account,
70 label: if spam { Label::Spam } else { Label::Ham },
71 messages,
72 means: sum.iter().zip(&known).map(|(s, &k)| (k > 0).then(|| s / k as f64)).collect(),
73 known: known.to_vec(),
74 })
75 .collect())
76}
77
78#[cfg(test)]
79mod tests {
80 use super::*;
81 use crate::synthetic;
82
83 #[test]
84 fn means_by_account_and_label() {
85 let root = std::env::temp_dir().join(format!("sioul-learn-diagnose-{}", std::process::id()));
86 let _ = std::fs::remove_dir_all(&root);
87 let dirs = Dirs::under(&root);
88 let mail = synthetic::mailbox(2, 60, 40);
89 let records: Vec<_> = mail.iter().enumerate().map(|(i, m)| synthetic::record_of(m, 1, i as u32 + 1)).collect();
90 corpus::store(&dirs, &records).unwrap();
91 let groups = feature_means(&dirs, &BTreeMap::new(), None).unwrap();
92 assert_eq!(groups.iter().map(|g| g.messages).sum::<u64>(), 100);
93 assert!(groups.iter().all(|g| g.means.len() == N && g.known.len() == N));
94 let list = NAMES.iter().position(|n| *n == "list_unsubscribe").unwrap();
95 assert!(groups.iter().all(|g| g.means[list].is_some_and(|m| (0.0..=1.0).contains(&m))), "{groups:?}");
96 let lottery = feature_means(&dirs, &BTreeMap::new(), Some("lottery.test")).unwrap();
98 assert!(!lottery.is_empty() && lottery.iter().all(|g| g.label == Label::Spam), "{lottery:?}");
99 let _ = std::fs::remove_dir_all(root);
100 }
101}