use crate::config::extraction_profiles::DocumentType;
#[derive(Debug, Clone)]
pub struct DocumentStats {
pub line_count: usize,
pub avg_chars_per_line: f32,
pub char_variance: f32,
pub justified_lines_percentage: f32,
pub avg_word_gap: f32,
pub word_gap_variance: f32,
pub form_field_percentage: f32,
pub citation_percentage: f32,
pub math_symbols_count: usize,
pub avg_line_spacing: f32,
pub spacing_tightness: f32,
}
impl Default for DocumentStats {
fn default() -> Self {
Self {
line_count: 0,
avg_chars_per_line: 0.0,
char_variance: 0.0,
justified_lines_percentage: 0.0,
avg_word_gap: 0.0,
word_gap_variance: 0.0,
form_field_percentage: 0.0,
citation_percentage: 0.0,
math_symbols_count: 0,
avg_line_spacing: 0.0,
spacing_tightness: 0.0,
}
}
}
pub struct DocumentClassifier;
impl DocumentClassifier {
pub fn classify_lines<'a, I>(lines: I) -> (DocumentType, DocumentStats)
where
I: Iterator<Item = &'a str>,
{
let mut stats = DocumentStats::default();
let mut line_lengths = Vec::new();
let mut justified_count = 0;
let mut form_field_count = 0;
let mut citation_count = 0;
let mut math_symbol_count = 0;
let mut line_spacing_values = Vec::new();
let mut word_gaps = Vec::new();
for (idx, line) in lines.enumerate() {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
stats.line_count += 1;
line_lengths.push(trimmed.len());
if Self::looks_justified(trimmed) {
justified_count += 1;
}
if Self::contains_form_field_markers(trimmed) {
form_field_count += 1;
}
if Self::looks_like_citation(trimmed) {
citation_count += 1;
}
if Self::contains_math_symbols(trimmed) {
math_symbol_count += 1;
}
let gaps = Self::extract_word_gaps(trimmed);
word_gaps.extend(gaps);
if idx > 0 {
line_spacing_values.push(1.0);
}
}
if !line_lengths.is_empty() {
let total_chars: usize = line_lengths.iter().sum();
stats.avg_chars_per_line = total_chars as f32 / line_lengths.len() as f32;
let mean = stats.avg_chars_per_line;
let variance: f32 = line_lengths
.iter()
.map(|&len| {
let diff = len as f32 - mean;
diff * diff
})
.sum::<f32>()
/ line_lengths.len() as f32;
stats.char_variance = variance.sqrt();
}
if stats.line_count > 0 {
stats.justified_lines_percentage = (justified_count as f32 / stats.line_count as f32) * 100.0;
stats.form_field_percentage = (form_field_count as f32 / stats.line_count as f32) * 100.0;
stats.citation_percentage = (citation_count as f32 / stats.line_count as f32) * 100.0;
}
stats.math_symbols_count = math_symbol_count;
if !word_gaps.is_empty() {
let total_gap: f32 = word_gaps.iter().sum();
stats.avg_word_gap = total_gap / word_gaps.len() as f32;
let mean = stats.avg_word_gap;
let variance: f32 = word_gaps
.iter()
.map(|&gap| {
let diff = gap - mean;
diff * diff
})
.sum::<f32>()
/ word_gaps.len() as f32;
stats.word_gap_variance = variance.sqrt();
stats.spacing_tightness = word_gaps.iter().filter(|&&g| g < 0.15).count() as f32 / word_gaps.len() as f32;
}
if !line_spacing_values.is_empty() {
let total_spacing: f32 = line_spacing_values.iter().sum();
stats.avg_line_spacing = total_spacing / line_spacing_values.len() as f32;
}
let doc_type = Self::classify_from_stats(&stats);
(doc_type, stats)
}
pub fn classify(_data: &str) -> (DocumentType, DocumentStats) {
(DocumentType::Mixed, DocumentStats::default())
}
fn classify_from_stats(stats: &DocumentStats) -> DocumentType {
if stats.form_field_percentage > 15.0 {
return DocumentType::Form;
}
if stats.math_symbols_count > 5 && stats.citation_percentage > 5.0 {
return DocumentType::Academic;
}
if stats.justified_lines_percentage > 30.0 && stats.word_gap_variance > 1.0 && stats.form_field_percentage < 5.0
{
return DocumentType::Policy;
}
if stats.spacing_tightness > 0.7 && stats.math_symbols_count > 2 {
return DocumentType::Academic;
}
if stats.avg_word_gap > 0.3 && stats.word_gap_variance > 2.0 {
return DocumentType::ScannedOCR;
}
if stats.char_variance < 10.0 && stats.justified_lines_percentage > 20.0 {
return DocumentType::Government;
}
DocumentType::Mixed
}
fn looks_justified(line: &str) -> bool {
if line.is_empty() {
return false;
}
let double_space_count = line.matches(" ").count();
let word_count = line.split_whitespace().count();
if word_count > 5 && double_space_count > 0 {
return true;
}
false
}
fn contains_form_field_markers(line: &str) -> bool {
let has_underscores = line.matches('_').count() >= 3;
let has_brackets = line.contains('[') || line.contains(']');
let has_boxes = line.contains('☐') || line.contains('☒') || line.contains('□');
has_underscores || has_brackets || has_boxes
}
fn extract_word_gaps(line: &str) -> Vec<f32> {
let mut gaps = Vec::new();
let mut in_gap = false;
let mut gap_size = 0;
for ch in line.chars() {
if ch == ' ' {
gap_size += 1;
in_gap = true;
} else if in_gap {
gaps.push(gap_size as f32 * 0.1);
gap_size = 0;
in_gap = false;
}
}
if in_gap {
gaps.push(gap_size as f32 * 0.1);
}
gaps
}
fn looks_like_citation(text: &str) -> bool {
if text.contains("et al") || text.contains("et. al") {
return true;
}
if text.contains('[') && text.contains(']') {
return true;
}
if text.len() >= 4 {
for chunk in text.chars().collect::<Vec<_>>().windows(4) {
let s: String = chunk.iter().collect();
if let Ok(year) = s.parse::<u32>()
&& (1900..=2100).contains(&year)
{
return true;
}
}
}
false
}
fn contains_math_symbols(text: &str) -> bool {
text.chars().any(Self::is_math_symbol)
}
fn is_math_symbol(c: char) -> bool {
matches!(
c,
'∑' | '∫'
| '∂'
| '∇'
| '√'
| '∞'
| '≈'
| '≠'
| '≤'
| '≥'
| '±'
| '×'
| '÷'
| 'α'
| 'β'
| 'γ'
| 'δ'
| 'ε'
| 'θ'
| 'λ'
| 'μ'
| 'π'
| 'σ'
| 'ω'
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_citation_detection() {
assert!(DocumentClassifier::looks_like_citation("et al"));
assert!(DocumentClassifier::looks_like_citation("et. al"));
assert!(DocumentClassifier::looks_like_citation("[1]"));
assert!(DocumentClassifier::looks_like_citation("[42]"));
assert!(!DocumentClassifier::looks_like_citation("hello world"));
}
#[test]
fn test_math_symbol_detection() {
assert!(DocumentClassifier::contains_math_symbols("π is pi"));
assert!(DocumentClassifier::contains_math_symbols("∫ is integral"));
assert!(!DocumentClassifier::contains_math_symbols("hello world"));
}
#[test]
fn test_math_symbol_recognition() {
assert!(DocumentClassifier::is_math_symbol('π'));
assert!(DocumentClassifier::is_math_symbol('∫'));
assert!(DocumentClassifier::is_math_symbol('≈'));
assert!(!DocumentClassifier::is_math_symbol('a'));
}
#[test]
fn test_form_field_detection() {
assert!(DocumentClassifier::contains_form_field_markers("Name: ___________"));
assert!(DocumentClassifier::contains_form_field_markers("[X] Check here"));
assert!(DocumentClassifier::contains_form_field_markers("Address: _______"));
assert!(!DocumentClassifier::contains_form_field_markers("This is normal text"));
}
#[test]
fn test_justified_text_detection() {
assert!(DocumentClassifier::looks_justified(
"This is justified text with variable spacing"
));
assert!(!DocumentClassifier::looks_justified("Short text"));
assert!(!DocumentClassifier::looks_justified(""));
}
#[test]
fn test_word_gap_extraction() {
let gaps = DocumentClassifier::extract_word_gaps("word1 word2 word3");
assert!(!gaps.is_empty());
assert!(gaps.len() >= 2);
}
#[test]
fn test_classify_form_document() {
let lines = vec![
"Name: ___________",
"Address: _______",
"Phone: ___________",
"[X] Check here",
];
let (doc_type, stats) = DocumentClassifier::classify_lines(lines.into_iter());
assert_eq!(doc_type, DocumentType::Form);
assert!(stats.form_field_percentage > 15.0);
}
#[test]
fn test_academic_characteristics_detection() {
let lines = vec![
"Abstract: We prove that π ≈ ∑ contribution",
"Smith et al. (2020) showed π in ∫ dx",
"From [1] we know ∂ exists with ∞ solutions",
"Therefore λ > 0 and α ∈ (0,1)",
];
let (_doc_type, stats) = DocumentClassifier::classify_lines(lines.into_iter());
assert!(stats.math_symbols_count > 2, "Should detect multiple math symbols");
assert!(stats.citation_percentage > 0.0, "Should detect academic citations");
}
#[test]
fn test_classify_justified_document() {
let lines = vec![
"This is justified text with variable spacing throughout",
"The document maintains consistent margins on both sides",
"Justified alignment is common in policy and legal texts",
"Each line extends to the right margin with adjustments",
];
let (doc_type, stats) = DocumentClassifier::classify_lines(lines.into_iter());
assert!(stats.justified_lines_percentage > 30.0);
assert!(matches!(
doc_type,
DocumentType::Policy | DocumentType::Government | DocumentType::Academic | DocumentType::Mixed
));
}
#[test]
fn test_statistics_calculation() {
let lines = vec!["short", "medium length", "very long line here"];
let (_doc_type, stats) = DocumentClassifier::classify_lines(lines.into_iter());
assert!(stats.line_count > 0);
assert!(stats.avg_chars_per_line > 0.0);
assert!(stats.char_variance >= 0.0);
}
#[test]
fn test_empty_document() {
let lines: Vec<&str> = vec![];
let (_doc_type, stats) = DocumentClassifier::classify_lines(lines.into_iter());
assert_eq!(stats.line_count, 0);
assert_eq!(stats.avg_chars_per_line, 0.0);
}
}