spectre_pdf 1.0.0

//! Whitespace-column table detection for borderless PDF tables.
//!
//! Extract per-page text, identify regions of consistent multi-column
//! structure by gap patterns, find column break positions where most
//! lines have whitespace, split each line into cells. Tuned for the
//! borderless tables that dominate financial filings (10-Ks,
//! prospectuses, municipal Official Statements).

use crate::ExtractError;

/// Returns one table per detection: rows × cells. `page_num` is
/// 0-indexed; `None` scans every page.
pub fn extract_tables_impl(
    pdf_bytes: &[u8],
    page_num: Option<usize>,
) -> Result<Vec<Vec<Vec<String>>>, ExtractError> {
    let page_texts = extract_page_texts(pdf_bytes)?;

    let pages_to_process: Vec<&str> = match page_num {
        Some(n) if n < page_texts.len() => vec![page_texts[n].as_str()],
        Some(_) => return Ok(vec![]),
        None => page_texts.iter().map(String::as_str).collect(),
    };

    let mut all_tables: Vec<Vec<Vec<String>>> = Vec::new();
    for page_text in pages_to_process {
        all_tables.extend(detect_tables_in_page(page_text));
    }
    Ok(all_tables)
}

fn extract_page_texts(pdf_bytes: &[u8]) -> Result<Vec<String>, ExtractError> {
    let doc = crate::document::open_sp_with_password(pdf_bytes, b"")?;
    let mut page_nums: Vec<u32> = doc.get_pages().into_keys().collect();
    page_nums.sort();
    Ok(page_nums
        .iter()
        .map(|&n| doc.extract_text(&[n]).unwrap_or_default())
        .collect())
}

/// Detect whitespace-column tables in page text.
///
/// Key insight: borderless financial tables (10-Ks, prospectuses,
/// municipal Official Statements) use large whitespace gaps to
/// separate columns. Lines within a table region have consistent
/// gap patterns.
fn detect_tables_in_page(page_text: &str) -> Vec<Vec<Vec<String>>> {
    let lines: Vec<&str> = page_text.lines().collect();
    if lines.is_empty() {
        return vec![];
    }

    // Group lines into candidate table regions
    let regions = identify_table_regions(&lines);
    let mut tables = Vec::new();

    for region in regions {
        if region.len() < 2 {
            continue;
        }
        let table = parse_table_region(&region);
        if !table.is_empty() && table.len() >= 2 {
            // At least 2 rows with at least 2 columns
            let has_multi_col = table.iter().any(|row| row.len() >= 2);
            if has_multi_col {
                tables.push(table);
            }
        }
    }

    tables
}

/// Identify groups of lines that look like table rows.
///
/// Heuristic: a table region is a consecutive block of lines where
/// multiple lines share similar column structure (detected by whitespace gaps ≥ 2 spaces).
fn identify_table_regions(lines: &[&str]) -> Vec<Vec<String>> {
    let mut regions: Vec<Vec<String>> = Vec::new();
    let mut current_region: Vec<String> = Vec::new();
    let mut consecutive_structured = 0usize;

    for &line in lines {
        let stripped = line.trim();
        if stripped.is_empty() {
            if consecutive_structured >= 2 && !current_region.is_empty() {
                // End of a structured block — save if it has table characteristics
                regions.push(current_region.clone());
                current_region.clear();
            }
            consecutive_structured = 0;
            continue;
        }

 // A "structured" line has multiple whitespace-separated columns
        // (detected by 2+ consecutive spaces suggesting column gaps)
        let has_gaps = has_column_gaps(stripped);
        if has_gaps {
            consecutive_structured += 1;
            current_region.push(line.to_string());
        } else {
 // Single-column line: could be header or paragraph
            // Include it if we're inside a table region
            if consecutive_structured >= 1 {
                current_region.push(line.to_string());
            } else {
                if !current_region.is_empty() {
                    if consecutive_structured >= 2 {
                        regions.push(current_region.clone());
                    }
                    current_region.clear();
                }
                consecutive_structured = 0;
            }
        }
    }

    if consecutive_structured >= 2 && !current_region.is_empty() {
        regions.push(current_region);
    }

    regions
}

/// Check if a line has significant whitespace gaps (≥2 spaces) suggesting columns.
fn has_column_gaps(line: &str) -> bool {
    // Count transitions from non-space to double-space
    let mut gap_count = 0usize;
    let chars: Vec<char> = line.chars().collect();
    let n = chars.len();
    let mut i = 0;

    while i < n {
        if chars[i] == ' ' {
            let gap_start = i;
            while i < n && chars[i] == ' ' {
                i += 1;
            }
            let gap_len = i - gap_start;
            if gap_len >= 2 {
                gap_count += 1;
            }
        } else {
            i += 1;
        }
    }

    gap_count >= 1
}

/// Parse a table region into rows and cells using column gap detection.
fn parse_table_region(lines: &[String]) -> Vec<Vec<String>> {
    if lines.is_empty() {
        return vec![];
    }

    // Find column break positions across all lines
    let col_breaks = find_column_breaks(lines);

    // Split each line at column breaks
    lines
        .iter()
        .map(|line| split_line_at_breaks(line, &col_breaks))
        .filter(|row| !row.is_empty() && row.iter().any(|c| !c.is_empty()))
        .collect()
}

/// Find column break positions that appear consistently across lines.
///
/// A break position is a **character** index (NOT a byte index) where most
/// lines have whitespace. Multi-byte UTF-8 characters (e.g. CJK, accented
/// Latin, emoji) advance the position by one, matching how a human reader
/// counts columns in a fixed-pitch table. The downstream
/// [`split_line_at_breaks`] call also operates on `chars`, so the indices
/// produced here are consumed consistently.
///
/// Vote-based: at each char position, count lines that have ` ` there.
/// Lines shorter than `max_chars` are treated as if they had spaces in the
/// missing trailing positions.
fn find_column_breaks(lines: &[String]) -> Vec<usize> {
    let lines_chars: Vec<Vec<char>> = lines.iter().map(|l| l.chars().collect()).collect();
    let max_chars = lines_chars.iter().map(|c| c.len()).max().unwrap_or(0);
    if max_chars == 0 {
        return vec![];
    }

    let mut space_votes: Vec<usize> = vec![0; max_chars + 1];
    let n_lines = lines_chars.len();

    for chars in &lines_chars {
        for (i, &ch) in chars.iter().enumerate() {
            if ch == ' ' {
                space_votes[i] += 1;
            }
        }
 // Lines shorter than max_chars vote "space" for the missing trailing
 // positions — that's how a column at position N still gets credited
        // when some rows have empty trailing cells.
        for vote in &mut space_votes[chars.len()..max_chars] {
            *vote += 1;
        }
    }

    // Threshold: a position is a "gap" if ≥60% of lines have space there.
    let threshold = (n_lines as f64 * 0.6) as usize;
    let gap_mask: Vec<bool> = space_votes.iter().map(|&v| v >= threshold).collect();

    let mut breaks: Vec<usize> = Vec::new();
    let mut in_gap = false;
    let mut gap_start = 0;

    for (i, &is_gap) in gap_mask.iter().enumerate() {
        if is_gap && !in_gap {
            in_gap = true;
            gap_start = i;
        } else if !is_gap && in_gap {
            in_gap = false;
            let gap_len = i - gap_start;
            if gap_len >= 2 {
                // Use the end of the gap as the break (start of next column).
                breaks.push(i);
            }
        }
    }
    if in_gap {
        let gap_len = max_chars - gap_start;
        if gap_len >= 2 {
            breaks.push(max_chars);
        }
    }

    breaks
}

/// Split a line into cells at the given character break positions.
fn split_line_at_breaks(line: &str, breaks: &[usize]) -> Vec<String> {
    if breaks.is_empty() {
        let t = line.trim().to_string();
        return if t.is_empty() { vec![] } else { vec![t] };
    }

    let chars: Vec<char> = line.chars().collect();
    let mut cells: Vec<String> = Vec::new();
    let mut prev = 0;

    for &brk in breaks {
        let end = brk.min(chars.len());
        let cell: String = chars[prev..end].iter().collect();
        cells.push(cell.trim().to_string());
        prev = end;
    }

    // Last segment
    if prev < chars.len() {
        let cell: String = chars[prev..].iter().collect();
        cells.push(cell.trim().to_string());
    }

    cells
}

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

    #[test]
    fn test_column_gap_detection() {
        // Simulate a rating table line
        let line = "Moody's             Aa2";
        assert!(has_column_gaps(line));
    }

    #[test]
    fn test_split_at_breaks() {
        let line = "Moody's             Aa2".to_string();
        let breaks = find_column_breaks(&[
            "Agency              Rating".to_string(),
            "Moody's             Aa2".to_string(),
            "S&P                 AA+".to_string(),
        ]);
        assert!(!breaks.is_empty(), "Should detect column breaks");
        let cells = split_line_at_breaks(&line, &breaks);
        assert!(cells.len() >= 2, "Should have at least 2 cells");
    }

    #[test]
    fn test_simple_table_region() {
        let lines = vec![
            "Agency              Rating    Outlook".to_string(),
            "Moody's             Aa2       Stable".to_string(),
            "S&P                 AA+       Stable".to_string(),
            "Fitch               AA        Stable".to_string(),
        ];
        let table = parse_table_region(&lines);
        assert!(!table.is_empty(), "Should parse table");
        assert!(table[0].len() >= 2, "Should have multiple columns");
    }
}