tests.py

import json
import re
import numpy as np
from bs4 import BeautifulSoup

from dataclasses import asdict, dataclass
from enum import Enum
from typing import List, Optional, Tuple, Dict, Any

from fuzzysearch import find_near_matches
from rapidfuzz import fuzz
from tqdm import tqdm

from olmocr.repeatdetect import RepeatDetector
from .katex.render import render_equation, compare_rendered_equations

class TestType(str, Enum):
    BASELINE = "baseline"
    PRESENT = "present"
    ABSENT = "absent"
    ORDER = "order"
    TABLE = "table"
    MATH = "math"


class TestChecked(str, Enum):
    VERIFIED = "verified"
    REJECTED = "rejected"


class ValidationError(Exception):
    """Exception raised for validation errors."""

    pass


def normalize_text(md_content: str) -> str:
    # Normalize whitespace in the md_content
    md_content = re.sub(r'\s+', " ", md_content)

    # Dictionary of characters to replace: keys are fancy characters, values are ASCII equivalents
    replacements = {
        "‘": "'", "’": "'", "‚": "'",
        "“": "\"", "”": "\"", "„": "\"",
        "＿": "_",
        "–": "-", "—": "-", "‑": "-", "‒": "-"
    }

     # Apply all replacements from the dictionary
    for fancy_char, ascii_char in replacements.items():
        md_content = md_content.replace(fancy_char, ascii_char)

    return md_content


@dataclass(kw_only=True)
class BasePDFTest:
    """
    Base class for all PDF test types.

    Attributes:
        pdf: The PDF filename.
        page: The page number for the test.
        id: Unique identifier for the test.
        type: The type of test.
        threshold: A float between 0 and 1 representing the threshold for fuzzy matching.
    """

    pdf: str
    page: int
    id: str
    type: str
    max_diffs: int = 0
    checked: Optional[TestChecked] = None

    def __post_init__(self):
        if not self.pdf:
            raise ValidationError("PDF filename cannot be empty")
        if not self.id:
            raise ValidationError("Test ID cannot be empty")
        if not isinstance(self.max_diffs, int) or self.max_diffs < 0:
            raise ValidationError(f"Max diffs must be positive number or 0")
        if self.type not in {t.value for t in TestType}:
            raise ValidationError(f"Invalid test type: {self.type}")

    def run(self, md_content: str) -> Tuple[bool, str]:
        """
        Run the test on the provided markdown content.

        Args:
            md_content: The content of the .md file.

        Returns:
            A tuple (passed, explanation) where 'passed' is True if the test passes,
            and 'explanation' provides details when the test fails.
        """
        raise NotImplementedError("Subclasses must implement the run method")


@dataclass
class TextPresenceTest(BasePDFTest):
    """
    Test to verify the presence or absence of specific text in a PDF.

    Attributes:
        text: The text string to search for.
    """

    text: str
    case_sensitive: bool=True

    def __post_init__(self):
        super().__post_init__()
        if self.type not in {TestType.PRESENT.value, TestType.ABSENT.value}:
            raise ValidationError(f"Invalid type for TextPresenceTest: {self.type}")
        if not self.text.strip():
            raise ValidationError("Text field cannot be empty")

    def run(self, md_content: str) -> Tuple[bool, str]:
        reference_query = self.text

        # Normalize whitespace in the md_content
        md_content = normalize_text(md_content)

        if not self.case_sensitive:
            reference_query = reference_query.lower()
            md_content = md_content.lower()

        # Threshold for fuzzy matching derived from max_diffs
        threshold = 1.0 - (self.max_diffs / (len(reference_query) if len(reference_query) > 0 else 1))
        best_ratio = fuzz.partial_ratio(reference_query, md_content) / 100.0

        if self.type == TestType.PRESENT.value:
            if best_ratio >= threshold:
                return True, ""
            else:
                msg = f"Expected '{reference_query[:40]}...' with threshold {threshold} " f"but best match ratio was {best_ratio:.3f}"
                return False, msg
        else:  # ABSENT
            if best_ratio < threshold:
                return True, ""
            else:
                msg = f"Expected absence of '{reference_query[:40]}...' with threshold {threshold} " f"but best match ratio was {best_ratio:.3f}"
                return False, msg


@dataclass
class TextOrderTest(BasePDFTest):
    """
    Test to verify that one text appears before another in a PDF.

    Attributes:
        before: The text expected to appear first.
        after: The text expected to appear after the 'before' text.
    """

    before: str
    after: str

    def __post_init__(self):
        super().__post_init__()
        if self.type != TestType.ORDER.value:
            raise ValidationError(f"Invalid type for TextOrderTest: {self.type}")
        if not self.before.strip():
            raise ValidationError("Before field cannot be empty")
        if not self.after.strip():
            raise ValidationError("After field cannot be empty")

    def run(self, md_content: str) -> Tuple[bool, str]:
        md_content = normalize_text(md_content)

        before_matches = find_near_matches(self.before, md_content, max_l_dist=self.max_diffs)
        after_matches = find_near_matches(self.after, md_content, max_l_dist=self.max_diffs)

        if not before_matches:
            return False, f"'before' text '{self.before[:40]}...' not found with max_l_dist {self.max_diffs}"
        if not after_matches:
            return False, f"'after' text '{self.after[:40]}...' not found with max_l_dist {self.max_diffs}"

        for before_match in before_matches:
            for after_match in after_matches:
                if before_match.start < after_match.start:
                    return True, ""
        return False, (f"Could not find a location where '{self.before[:40]}...' appears before " f"'{self.after[:40]}...'.")



@dataclass
class TableTest(BasePDFTest):
    """
    Test to verify certain properties of a table are held, namely that some cells appear relative to other cells correctly
    """
    # This is the target cell, which must exist in at least one place in the table
    cell: str

    # These properties say that the cell immediately up/down/left/right of the target cell has the string specified
    up: str = ""
    down: str = ""
    left: str = ""
    right: str = ""

    # These properties say that the cell all the way up, or all the way left of the target cell (ex. headings) has the string value specified
    top_heading: str = ""
    left_heading: str = ""


    def __post_init__(self):
        super().__post_init__()
        if self.type != TestType.TABLE.value:
            raise ValidationError(f"Invalid type for TableTest: {self.type}")

    def parse_markdown_tables(self, md_content: str) -> List[np.ndarray]:
        """
        Extract and parse all markdown tables from the provided content.

        Args:
            md_content: The markdown content containing tables

        Returns:
            A list of numpy arrays, each representing a parsed table
        """
        import re
        import numpy as np

        # Updated regex to allow optional leading and trailing pipes
        table_pattern = (
            r'(\|?(?:[^|\n]*\|)+[^|\n]*\|?)\s*\n'
            r'\|?(?:[ :-]+\|)+[ :-]+\|?\s*\n'
            r'((?:\|?(?:[^|\n]*\|)+[^|\n]*\|?\s*\n)+)'
        )
        table_matches = re.finditer(table_pattern, md_content)

        parsed_tables = []

        for table_match in table_matches:
            # Extract header and body from the table match
            header_row = table_match.group(1).strip()
            body_rows = table_match.group(2).strip().split('\n')

            # Process header and rows to remove leading/trailing pipes
            header_cells = [cell.strip() for cell in header_row.split('|')]
            if header_cells and header_cells[0] == '':
                header_cells = header_cells[1:]
            if header_cells and header_cells[-1] == '':
                header_cells = header_cells[:-1]

            # Process table body rows
            table_data = []
            for row in [header_row] + body_rows:
                if '|' not in row:  # Skip separator row
                    continue

                cells = [cell.strip() for cell in row.split('|')]
                if cells and cells[0] == '':
                    cells = cells[1:]
                if cells and cells[-1] == '':
                    cells = cells[:-1]

                table_data.append(cells)

            # Skip separator row (second row with dashes)
            if len(table_data) > 1 and all('-' in cell for cell in table_data[1]):
                table_data = [table_data[0]] + table_data[2:]

            # Convert to numpy array for easier manipulation
            # Ensure all rows have the same number of columns by padding if necessary
            max_cols = max(len(row) for row in table_data)
            padded_data = [row + [''] * (max_cols - len(row)) for row in table_data]
            table_array = np.array(padded_data)

            parsed_tables.append(table_array)

        return parsed_tables

    def parse_html_tables(self, html_content: str) -> List[np.ndarray]:
        """
        Extract and parse all HTML tables from the provided content.
        
        Args:
            html_content: The HTML content containing tables
            
        Returns:
            A list of numpy arrays, each representing a parsed table
        """
        soup = BeautifulSoup(html_content, 'html.parser')
        tables = soup.find_all('table')
        
        parsed_tables = []
        
        for table in tables:
            rows = table.find_all(['tr'])
            table_data = []
            
            for row in rows:
                cells = row.find_all(['th', 'td'])
                row_data = [cell.get_text().strip() for cell in cells]
                table_data.append(row_data)
            
            # Ensure all rows have the same number of columns
            if table_data:
                max_cols = max(len(row) for row in table_data)
                padded_data = [row + [''] * (max_cols - len(row)) for row in table_data]
                table_array = np.array(padded_data)
                parsed_tables.append(table_array)
        
        return parsed_tables

    def run(self, content: str) -> Tuple[bool, str]:
        """
        Run the table test on provided content.
        
        Finds all tables (markdown and/or HTML based on content_type) and checks if any cell 
        matches the target cell and satisfies the specified relationships.
        
        Args:
            content: The content containing tables (markdown or HTML)
            
        Returns:
            A tuple (passed, explanation) where 'passed' is True if the test passes,
            and 'explanation' provides details when the test fails.
        """
        # Initialize variables to track tables and results
        tables_to_check = []
        failed_reasons = []
        
        # Threshold for fuzzy matching derived from max_diffs
        threshold = 1.0 - (self.max_diffs / (len(self.cell) if len(self.cell) > 0 else 1))
        
        # Parse tables based on content_type
        md_tables = self.parse_markdown_tables(content)
        tables_to_check.extend(md_tables)
        
        html_tables = self.parse_html_tables(content)
        tables_to_check.extend(html_tables)
        
        # If no tables found, return failure
        if not tables_to_check:
            return False, f"No tables found in the content"
            
        # Check each table
        for table_array in tables_to_check:
            # Find all cells that match the target cell using fuzzy matching
            matches = []
            for i in range(table_array.shape[0]):
                for j in range(table_array.shape[1]):
                    cell_content = table_array[i, j]
                    similarity = fuzz.ratio(self.cell, cell_content) / 100.0
                    
                    if similarity >= threshold:
                        matches.append((i, j))
            
            # If no matches found in this table, continue to the next table
            if not matches:
                continue
                
            # Check the relationships for each matching cell
            for row_idx, col_idx in matches:
                all_relationships_satisfied = True
                current_failed_reasons = []
                
                # Check up relationship
                if self.up and row_idx > 0:
                    up_cell = table_array[row_idx - 1, col_idx]
                    up_similarity = fuzz.ratio(self.up, up_cell) / 100.0
                    if up_similarity < threshold:
                        all_relationships_satisfied = False
                        current_failed_reasons.append(f"Cell above '{up_cell}' doesn't match expected '{self.up}' (similarity: {up_similarity:.2f})")
                
                # Check down relationship
                if self.down and row_idx < table_array.shape[0] - 1:
                    down_cell = table_array[row_idx + 1, col_idx]
                    down_similarity = fuzz.ratio(self.down, down_cell) / 100.0
                    if down_similarity < threshold:
                        all_relationships_satisfied = False
                        current_failed_reasons.append(f"Cell below '{down_cell}' doesn't match expected '{self.down}' (similarity: {down_similarity:.2f})")
                
                # Check left relationship
                if self.left and col_idx > 0:
                    left_cell = table_array[row_idx, col_idx - 1]
                    left_similarity = fuzz.ratio(self.left, left_cell) / 100.0
                    if left_similarity < threshold:
                        all_relationships_satisfied = False
                        current_failed_reasons.append(f"Cell to the left '{left_cell}' doesn't match expected '{self.left}' (similarity: {left_similarity:.2f})")
                
                # Check right relationship
                if self.right and col_idx < table_array.shape[1] - 1:
                    right_cell = table_array[row_idx, col_idx + 1]
                    right_similarity = fuzz.ratio(self.right, right_cell) / 100.0
                    if right_similarity < threshold:
                        all_relationships_satisfied = False
                        current_failed_reasons.append(f"Cell to the right '{right_cell}' doesn't match expected '{self.right}' (similarity: {right_similarity:.2f})")
                
                # Check top heading relationship
                if self.top_heading and row_idx > 0:
                    # Find the first non-empty cell in the same column (starting from the top)
                    top_heading_cell = ""
                    for i in range(row_idx):
                        if table_array[i, col_idx].strip():
                            top_heading_cell = table_array[i, col_idx]
                            break
                    
                    if not top_heading_cell:
                        all_relationships_satisfied = False
                        current_failed_reasons.append(f"No non-empty top heading found in column {col_idx}")
                    else:
                        top_similarity = fuzz.ratio(self.top_heading, top_heading_cell) / 100.0
                        if top_similarity < threshold:
                            all_relationships_satisfied = False
                            current_failed_reasons.append(f"Top heading '{top_heading_cell}' doesn't match expected '{self.top_heading}' (similarity: {top_similarity:.2f})")
                
                # Check left heading relationship
                if self.left_heading and col_idx > 0:
                    # Find the first non-empty cell in the same row (starting from the left)
                    left_heading_cell = ""
                    for j in range(col_idx):
                        if table_array[row_idx, j].strip():
                            left_heading_cell = table_array[row_idx, j]
                            break
                    
                    if not left_heading_cell:
                        all_relationships_satisfied = False
                        current_failed_reasons.append(f"No non-empty left heading found in row {row_idx}")
                    else:
                        left_heading_similarity = fuzz.ratio(self.left_heading, left_heading_cell) / 100.0
                        if left_heading_similarity < threshold:
                            all_relationships_satisfied = False
                            current_failed_reasons.append(f"Left heading '{left_heading_cell}' doesn't match expected '{self.left_heading}' (similarity: {left_heading_similarity:.2f})")
            
                # If all relationships are satisfied for this cell, the test passes
                if all_relationships_satisfied:
                    return True, ""
                else:
                    failed_reasons.extend(current_failed_reasons)
        
        # If we've gone through all tables and all matching cells and none satisfied all relationships
        if not failed_reasons:
            return False, f"No cell matching '{self.cell}' found in any table with threshold {threshold}"
        else:
            return False, f"Found cells matching '{self.cell}' but relationships were not satisfied: {'; '.join(failed_reasons)}"


@dataclass
class BaselineTest(BasePDFTest):
    """
    This test makes sure that several baseline quality checks pass for the output generation.

    Namely, the output is not blank, not endlessly repeating, and contains characters of the proper
    character sets.

    """
    max_repeats: int=30

    def run(self, content: str) -> Tuple[bool, str]:
        if len("".join(c for c in content if c.isalnum()).strip()) == 0:
            return False, "The text contains no alpha numeric characters"

        # Makes sure that the content has no egregious repeated ngrams at the end, which indicate a degradation of quality
        # Honestly, this test doesn't seem to catch anything at the moment, maybe it can be refactored to a "text-quality"
        # test or something, that measures repetition, non-blanks, charsets, etc
        d = RepeatDetector(max_ngram_size=5)
        d.add_letters(content)
        repeats = d.ngram_repeats()

        for index, count in enumerate(repeats):
            if count > self.max_repeats:
                return False, f"Text ends with {count} repeating {index+1}-grams, invalid"

        pattern = re.compile(
            r'['
            r'\u4e00-\u9FFF'   # CJK Unified Ideographs (Chinese characters)
            r'\u3040-\u309F'   # Hiragana (Japanese)
            r'\u30A0-\u30FF'   # Katakana (Japanese)
            r'\U0001F600-\U0001F64F'  # Emoticons (Emoji)
            r'\U0001F300-\U0001F5FF'  # Miscellaneous Symbols and Pictographs (Emoji)
            r'\U0001F680-\U0001F6FF'  # Transport and Map Symbols (Emoji)
            r'\U0001F1E0-\U0001F1FF'  # Regional Indicator Symbols (flags, Emoji)
            r']',
            flags=re.UNICODE)
        matches = pattern.findall(content)
        if matches:
            return False, f"Text contains disallowed characters {matches}"

        return True, ""


@dataclass
class MathTest(BasePDFTest):
    math: str

    def __post_init__(self):
        super().__post_init__()
        if self.type != TestType.MATH.value:
            raise ValidationError(f"Invalid type for MathTest: {self.type}")
        if len(self.math.strip()) == 0:
            raise ValidationError(f"Math test must have non-empty math expression")

        self.reference_render = render_equation(self.math)

        if self.reference_render is None:
            raise ValidationError(f"Math equation {self.math} was not able to render")


    def run(self, content: str) -> Tuple[bool, str]:
        # Store both the search pattern and the full pattern to replace
        patterns = [
            (r'\$\$(.+?)\$\$', r'\$\$(.+?)\$\$'),        # $$...$$
            (r'\\\((.+?)\\\)', r'\\\((.+?)\\\)'),        # \(...\)
            (r'\\\[(.+?)\\\]', r'\\\[(.+?)\\\]'),        # \[...\]
            (r'\$(.+?)\$', r'\$(.+?)\$')                 # $...$
        ]

        equations = []
        modified_content = content
        
        for search_pattern, replace_pattern in patterns:
            # Find all matches for the current pattern
            matches = re.findall(search_pattern, modified_content, re.DOTALL)
            equations.extend([e.strip() for e in matches])
            
            # Replace all instances of this pattern with empty strings
            modified_content = re.sub(replace_pattern, '', modified_content, flags=re.DOTALL)

        # If an equation in the markdown exactly matches our math string, then that's good enough
        # we don't have to do a more expensive comparison
        if any(hyp == self.math for hyp in equations):
            return True, ""

        # If not, then let's render the math equation itself and now compare to each hypothesis
        best_match_score = 0.0
        best_match_render = None

        for hypothesis in equations:
            hypothesis_render = render_equation(hypothesis)

            if not hypothesis_render:
                continue

            if compare_rendered_equations(self.reference_render, hypothesis_render):
                return True, ""

        # self.reference_render.save(f"maths/{self.id}_ref.png", format="PNG")
        # best_match_render.save(f"maths/{self.id}_hyp.png", format="PNG")

        return False, f"No match found for {self.math} anywhere in content, best match threshold was {best_match_score:.3f}"    



def load_tests(jsonl_file: str) -> List[BasePDFTest]:
    """
    Load tests from a JSONL file.

    Args:
        jsonl_file: Path to the JSONL file containing test definitions.

    Returns:
        A list of test objects.
    """
    tests: List[BasePDFTest] = []
    unique_ids = set()
    with open(jsonl_file, "r") as file:
        for line_number, line in tqdm(enumerate(file, start=1), desc="Loading tests"):
            line = line.strip()
            if not line:
                continue

            try:
                data = json.loads(line)
                test_type = data.get("type")
                if test_type in {TestType.PRESENT.value, TestType.ABSENT.value}:
                    test = TextPresenceTest(**data)
                elif test_type == TestType.ORDER.value:
                    test = TextOrderTest(**data)
                elif test_type == TestType.TABLE.value:
                    test = TableTest(**data)
                elif test_type == TestType.MATH.value:
                    test = MathTest(**data)
                else:
                    raise ValidationError(f"Unknown test type: {test_type}")

                if test.id in unique_ids:
                    raise ValidationError(f"Test with duplicate id {test.id} found, error loading tests.")
                else:
                    unique_ids.add(test.id)

                tests.append(test)
            except json.JSONDecodeError as e:
                print(f"Error parsing JSON on line {line_number}: {e}")
                raise
            except (ValidationError, KeyError) as e:
                print(f"Error on line {line_number}: {e}")
                raise
            except Exception as e:
                print(f"Unexpected error on line {line_number}: {e}")
                raise

    return tests


def save_tests(tests: List[BasePDFTest], jsonl_file: str) -> None:
    """
    Save tests to a JSONL file using asdict for conversion.

    Args:
        tests: A list of test objects.
        jsonl_file: Path to the output JSONL file.
    """
    with open(jsonl_file, "w") as file:
        for test in tests:
            file.write(json.dumps(asdict(test)) + "\n")