Lumotia/crates/ai-formatting/src/correction_learning.rs

use std::collections::HashSet;

const MAX_REWRITE_RATIO: f64 = 0.5;
const MIN_CORRECTION_LEN: usize = 3;
const MAX_DISTANCE_RATIO: f64 = 0.65;
const MAX_CORRECTIONS_PER_EDIT: usize = 8;

fn edit_distance(a: &str, b: &str) -> usize {
    let a_chars: Vec<char> = a.chars().collect();
    let b_chars: Vec<char> = b.chars().collect();
    let mut prev: Vec<usize> = (0..=b_chars.len()).collect();
    let mut curr = vec![0usize; b_chars.len() + 1];

    for (i, a_char) in a_chars.iter().enumerate() {
        curr[0] = i + 1;
        for (j, b_char) in b_chars.iter().enumerate() {
            curr[j + 1] = if a_char == b_char {
                prev[j]
            } else {
                1 + prev[j].min(prev[j + 1]).min(curr[j])
            };
        }
        prev.clone_from(&curr);
    }

    prev[b_chars.len()]
}

fn trim_non_word_edges(word: &str) -> &str {
    word.trim_matches(|c: char| !c.is_alphanumeric() && c != '_')
}

fn tokenize(text: &str) -> Vec<String> {
    text.split_whitespace()
        .filter_map(|word| {
            let trimmed = trim_non_word_edges(word);
            (!trimmed.is_empty()).then(|| trimmed.to_string())
        })
        .collect()
}

fn find_edited_region(original_text: &str, field_value: &str) -> String {
    if field_value.len() <= (original_text.len() * 3) / 2 {
        return field_value.to_string();
    }

    if field_value.contains(original_text) {
        return original_text.to_string();
    }

    let orig_words = tokenize(original_text);
    let field_words = tokenize(field_value);
    let window_size = orig_words.len();

    if field_words.len() <= window_size || window_size == 0 {
        return field_value.to_string();
    }

    let mut best_start = 0usize;
    let mut best_score = 0usize;
    for start in 0..=field_words.len() - window_size {
        let mut matches = 0usize;
        for offset in 0..window_size {
            if field_words[start + offset].eq_ignore_ascii_case(&orig_words[offset]) {
                matches += 1;
            }
        }
        if matches > best_score {
            best_score = matches;
            best_start = start;
        }
    }

    if (best_score as f64) < (window_size as f64 * 0.3) {
        return field_value.to_string();
    }

    field_words[best_start..best_start + window_size].join(" ")
}

fn find_substitutions(original_words: &[String], edited_words: &[String]) -> Vec<(String, String)> {
    let m = original_words.len();
    let n = edited_words.len();
    let mut dp = vec![vec![0usize; n + 1]; m + 1];

    for i in 1..=m {
        for j in 1..=n {
            dp[i][j] = if original_words[i - 1].eq_ignore_ascii_case(&edited_words[j - 1]) {
                dp[i - 1][j - 1] + 1
            } else {
                dp[i - 1][j].max(dp[i][j - 1])
            };
        }
    }

    let mut aligned: Vec<(Option<String>, Option<String>)> = Vec::new();
    let mut i = m;
    let mut j = n;
    while i > 0 || j > 0 {
        if i > 0 && j > 0 && original_words[i - 1].eq_ignore_ascii_case(&edited_words[j - 1]) {
            aligned.push((
                Some(original_words[i - 1].clone()),
                Some(edited_words[j - 1].clone()),
            ));
            i -= 1;
            j -= 1;
        } else if j > 0 && (i == 0 || dp[i][j - 1] >= dp[i - 1][j]) {
            aligned.push((None, Some(edited_words[j - 1].clone())));
            j -= 1;
        } else {
            aligned.push((Some(original_words[i - 1].clone()), None));
            i -= 1;
        }
    }
    aligned.reverse();

    let mut substitutions = Vec::new();
    for pair in aligned.windows(2) {
        let (orig_word, edited_word) = (&pair[0].0, &pair[0].1);
        let (next_orig_word, next_edited_word) = (&pair[1].0, &pair[1].1);
        if let (Some(orig_word), None, None, Some(corrected_word)) =
            (orig_word, edited_word, next_orig_word, next_edited_word)
        {
            substitutions.push((orig_word.clone(), corrected_word.clone()));
        }
    }

    substitutions
}

pub fn extract_corrections(
    original_text: &str,
    edited_text: &str,
    existing_terms: &[String],
) -> Vec<String> {
    if original_text.trim().is_empty()
        || edited_text.trim().is_empty()
        || original_text == edited_text
    {
        return Vec::new();
    }

    let edited_region = find_edited_region(original_text, edited_text);
    if edited_region == original_text {
        return Vec::new();
    }

    let original_words = tokenize(original_text);
    let edited_words = tokenize(&edited_region);
    if original_words.is_empty() || edited_words.is_empty() {
        return Vec::new();
    }

    let substitutions = find_substitutions(&original_words, &edited_words);
    if (substitutions.len() as f64) > (original_words.len() as f64 * MAX_REWRITE_RATIO) {
        return Vec::new();
    }

    let existing: HashSet<String> = existing_terms
        .iter()
        .map(|term| term.to_ascii_lowercase())
        .collect();
    let mut seen = HashSet::new();
    let mut results = Vec::new();

    for (original_word, corrected_word) in substitutions {
        let normalized_original = original_word.to_ascii_lowercase();
        let normalized_corrected = corrected_word.to_ascii_lowercase();
        if normalized_original == normalized_corrected
            || normalized_corrected.len() < MIN_CORRECTION_LEN
            || existing.contains(&normalized_corrected)
            || seen.contains(&normalized_corrected)
        {
            continue;
        }

        let max_len = original_word.len().max(corrected_word.len()).max(1);
        let distance = edit_distance(&normalized_original, &normalized_corrected);
        if distance as f64 / max_len as f64 > MAX_DISTANCE_RATIO {
            continue;
        }

        results.push(corrected_word);
        seen.insert(normalized_corrected);
        if results.len() >= MAX_CORRECTIONS_PER_EDIT {
            break;
        }
    }

    results
}

#[cfg(test)]
mod tests {
    use super::extract_corrections;

    #[test]
    fn extracts_phonetic_corrections_for_profile_learning() {
        let corrections = extract_corrections(
            "Email Shunade about the client deck tomorrow.",
            "Email Sinead about the client deck tomorrow.",
            &[],
        );

        assert_eq!(corrections, vec!["Sinead"]);
    }

    #[test]
    fn ignores_large_rewrites() {
        let corrections = extract_corrections(
            "This is a rough transcript of the meeting agenda.",
            "Let's throw this away and write something completely different instead.",
            &[],
        );

        assert!(corrections.is_empty());
    }

    #[test]
    fn skips_terms_already_in_profile_dictionary() {
        let corrections = extract_corrections(
            "Follow up with Corble tomorrow morning.",
            "Follow up with CORBEL tomorrow morning.",
            &[String::from("CORBEL")],
        );

        assert!(corrections.is_empty());
    }
}