Copies/copienator/commands/post_correction.py

from __future__ import annotations

import argparse
import re
from collections.abc import Sequence
from pathlib import Path
from typing import Any

from copienator import (
    EvaluationWorkspace,
    ExitCode,
    atomic_write_bytes,
    atomic_write_json,
    evaluation_parser,
    execute,
    read_json,
    workspace_from_args,
)
from copienator.latex_json import repair_decoded_latex_controls

WORD_LIST_FILE = Path(__file__).resolve().parents[1] / "data" / "liste_francais.txt"
ACCENT_PATTERN = re.compile(r"[éèêëàâäîïôöùûüçœÉÈÊËÀÂÄÎÏÔÖÙÛÜÇŒ]")
MATH_PATTERN = re.compile(
    r"(\$\$.*?\$\$|\$.*?\$|\\\(.*?\\\)|\\\[.*?\\\])",
    re.DOTALL,
)


def build_parser() -> argparse.ArgumentParser:
    return evaluation_parser("Clean encoding and LaTeX issues in correction.json.")


def escape_latex_underscores(text: str) -> str:
    r"""Escape underscores outside math without double-escaping existing ones."""

    def escape_plain(value: str) -> str:
        # Collapse any existing escape run as well, making cleanup idempotent.
        return re.sub(r"\\*_", lambda _match: r"\_", value)

    parts: list[str] = []
    last_end = 0
    for match in MATH_PATTERN.finditer(text):
        start, end = match.span()
        parts.append(escape_plain(text[last_end:start]))
        parts.append(match.group(0))
        last_end = end
    parts.append(escape_plain(text[last_end:]))
    return "".join(parts)


def escape_latex_percents(text: str) -> str:
    r"""Escape percent signs before annotation text is passed to LaTeX.

    Unlike underscores, percent signs are comments both inside and outside
    math mode.  Collapsing an existing escape run keeps cleanup idempotent.
    """
    return re.sub(r"\\*%", lambda _match: r"\%", text)


def normalize_overescaped_latex_commands(text: str) -> str:
    r"""Collapse doubled command escapes inside LaTeX math environments.

    Model responses occasionally contain ``\\mathbb`` after JSON decoding where
    LaTeX requires ``\mathbb``.  A doubled backslash followed by whitespace is a
    legitimate row break (for example in ``cases``), so it must be preserved.
    """

    def normalize_math(match: re.Match[str]) -> str:
        return re.sub(r"\\\\(?=[A-Za-z{}])", r"\\", match.group(0))

    return MATH_PATTERN.sub(normalize_math, text)


def build_lookup_map(word_list_path: Path = WORD_LIST_FILE) -> dict[str, str]:
    words = word_list_path.read_text(encoding="utf-8").splitlines()
    lookup: dict[str, str] = {}
    for word in words:
        broken_key = ACCENT_PATTERN.sub("\x00", word)
        if "\x00" in broken_key:
            lookup[broken_key.lower()] = word
    return lookup


def fast_fix(text: str, lookup: dict[str, str]) -> str:
    def replacer(match: re.Match[str]) -> str:
        broken_word = match.group(0)
        return lookup.get(broken_word.lower(), broken_word)

    return re.sub(r"[a-zA-Z\x00]+", replacer, text)


def fix_hex_corruption_safe(text: str) -> str:
    return re.sub(
        r"\x00([eEfF][0-9a-fA-F])",
        lambda match: chr(int(match.group(1), 16)),
        text,
    )


def repair_json_escape_corruption(text: str) -> str:
    r"""Restore observed LaTeX commands consumed as JSON control escapes."""
    text = repair_decoded_latex_controls(text)
    replacements = (
        ("\x0crac", r"\frac"),
        ("\x0ceuille", r"\equiv"),
        ("\theta", r"\theta"),
        ("\times", r"\times"),
        ("\textbackslash ", "\\"),
        ("\negthinspace", r"\negthinspace"),
        ("\neq", r"\neq"),
        ("\not", r"\not"),
        ("∈", r"\ensuremath{\in}"),
        ("⊂", r"\ensuremath{\subset}"),
    )
    for broken, repaired in replacements:
        text = text.replace(broken, repaired)
    return text


def clean_string(text: str, lookup: dict[str, str]) -> str:
    text = fix_hex_corruption_safe(text)
    text = text.replace("\x19", "\x00")
    text = text.replace("\x18", "\x00")
    text = text.replace("\x00\x00", "\x00")
    text = re.sub(r" \x00{1,2} ", " à ", text)
    if "\x00" in text:
        text = fast_fix(text, lookup).replace("\x00", "")
    text = repair_json_escape_corruption(text)
    text = normalize_overescaped_latex_commands(text)
    return escape_latex_percents(escape_latex_underscores(text))


def clean_obj(value: Any, lookup: dict[str, str]) -> Any:
    if isinstance(value, str):
        return clean_string(value, lookup)
    if isinstance(value, list):
        return [clean_obj(item, lookup) for item in value]
    if isinstance(value, dict):
        return {
            key: item if key == "suffix" else clean_obj(item, lookup)
            for key, item in value.items()
        }
    return value


def unresolved_math_control_paths(value: Any, path: str = "$") -> list[str]:
    """Locate destructive JSON control escapes left in mathematical text."""
    if isinstance(value, str):
        return (
            [path]
            if any(
                control in {"\b", "\f"}
                for match in MATH_PATTERN.finditer(value)
                for control in match.group(0)
            )
            else []
        )
    if isinstance(value, list):
        return [
            issue
            for index, item in enumerate(value)
            for issue in unresolved_math_control_paths(item, f"{path}[{index}]")
        ]
    if isinstance(value, dict):
        return [
            issue
            for key, item in value.items()
            for issue in unresolved_math_control_paths(item, f"{path}[{key!r}]")
        ]
    return []


def run(
    workspace: EvaluationWorkspace,
    *,
    word_list_path: Path = WORD_LIST_FILE,
) -> ExitCode:
    workspace.require_files("correction.json")
    lookup = build_lookup_map(word_list_path)
    data = read_json(workspace.correction_file)
    cleaned = clean_obj(data, lookup)
    unresolved = unresolved_math_control_paths(cleaned)
    backup = workspace.root / "correction_precleanup.json"
    atomic_write_bytes(backup, workspace.correction_file.read_bytes())
    print(f"Original JSON backed up to {backup}")
    atomic_write_json(workspace.correction_file, cleaned)
    print(f"Fixed JSON saved to {workspace.correction_file}")
    if unresolved:
        print(
            "Warning: unresolved JSON control characters remain in mathematical "
            "text at: " + ", ".join(unresolved)
        )
        return ExitCode.PARTIAL
    return ExitCode.SUCCESS


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    return execute(parser, argv, lambda args: run(workspace_from_args(args)))


if __name__ == "__main__":
    raise SystemExit(main())