scripts/score_external_ocr.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

from scripts.benchmark_ocr import text_metrics


def read_candidate_text(path: Path) -> str:
    if not path.exists():
        raise FileNotFoundError(f"OCR text path not found: {path}")
    if path.is_dir():
        pieces = [
            child.read_text(encoding="utf-8", errors="replace")
            for child in sorted(path.glob("*.txt"))
            if child.is_file()
        ]
        if not pieces:
            raise ValueError(f"No .txt files found in OCR output directory: {path}")
        return "\n\n".join(pieces)
    return path.read_text(encoding="utf-8", errors="replace")


def parse_candidate(value: str) -> tuple[str, Path]:
    if "=" not in value:
        path = Path(value)
        return path.stem or "external", path
    label, path_text = value.split("=", 1)
    label = label.strip()
    if not label:
        raise ValueError("Candidate label cannot be empty.")
    return label, Path(path_text)


def score_candidate(label: str, path: Path) -> dict[str, Any]:
    text = read_candidate_text(path)
    metrics = text_metrics(text)
    return {
        "label": label,
        "path": str(path),
        "ok": bool(text.strip()),
        **metrics,
    }


def choose_best(results: list[dict[str, Any]]) -> dict[str, Any] | None:
    successful = [item for item in results if item.get("ok")]
    if not successful:
        return None
    return max(successful, key=lambda item: (item.get("qualityScore", 0), item.get("arabicWords", 0)))


def load_baseline(path: Path | None) -> dict[str, Any] | None:
    if path is None:
        return None
    if not path.exists():
        raise FileNotFoundError(f"Baseline JSON not found: {path}")
    payload = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(payload, list):
        baseline = choose_best([item for item in payload if isinstance(item, dict)])
    elif isinstance(payload, dict):
        if isinstance(payload.get("best"), dict):
            baseline = payload["best"]
        elif isinstance(payload.get("selected"), dict):
            baseline = payload["selected"]
        elif isinstance(payload.get("results"), list):
            baseline = choose_best([item for item in payload["results"] if isinstance(item, dict)])
        elif isinstance(payload.get("benchmark"), list):
            baseline = choose_best([item for item in payload["benchmark"] if isinstance(item, dict)])
        else:
            baseline = payload
    else:
        baseline = None
    if not baseline:
        raise ValueError(f"Could not find a usable baseline result in {path}")
    baseline = dict(baseline)
    baseline.setdefault("label", baseline.get("engine") or baseline.get("extraction") or "wired-baseline")
    baseline["path"] = str(path)
    return baseline


def compare_to_baseline(best: dict[str, Any] | None, baseline: dict[str, Any] | None) -> dict[str, Any] | None:
    if not best or not baseline:
        return None
    best_score = float(best.get("qualityScore") or 0)
    baseline_score = float(baseline.get("qualityScore") or 0)
    best_words = int(best.get("arabicWords") or 0)
    baseline_words = int(baseline.get("arabicWords") or 0)
    score_delta = round(best_score - baseline_score, 2)
    word_delta = best_words - baseline_words
    beats = score_delta > 0 or (score_delta == 0 and word_delta > 0)
    return {
        "baselineLabel": baseline.get("label") or baseline.get("engine") or baseline.get("extraction") or "wired-baseline",
        "baselineScore": baseline_score,
        "baselineArabicWords": baseline_words,
        "bestLabel": best.get("label"),
        "bestScore": best_score,
        "bestArabicWords": best_words,
        "scoreDelta": score_delta,
        "arabicWordDelta": word_delta,
        "beatsBaseline": beats,
        "promotionReady": bool(beats and best.get("quality") in {"good", "warning"}),
    }


def score_external_ocr(
    candidates: list[str],
    report_path: Path | None = None,
    baseline_json: Path | None = None,
    json_path: Path | None = None,
) -> dict[str, Any]:
    if not candidates:
        raise ValueError("At least one --candidate is required.")
    results = [score_candidate(*parse_candidate(candidate)) for candidate in candidates]
    best = choose_best(results)
    baseline = load_baseline(baseline_json)
    comparison = compare_to_baseline(best, baseline)
    payload = {
        "ready": bool(best and best.get("quality") in {"good", "warning"}),
        "promotionReady": bool(comparison.get("promotionReady")) if comparison else False,
        "best": best,
        "baseline": baseline,
        "comparison": comparison,
        "results": results,
    }
    if report_path:
        write_score_report(report_path, payload)
        payload["reportPath"] = str(report_path)
    if json_path:
        json_path.parent.mkdir(parents=True, exist_ok=True)
        json_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        payload["jsonPath"] = str(json_path)
    return payload


def markdown_value(value: Any) -> str:
    if value is None or value == "":
        return "-"
    return str(value)


def write_score_report(path: Path, payload: dict[str, Any]) -> None:
    best = payload.get("best") or {}
    baseline = payload.get("baseline") or {}
    comparison = payload.get("comparison") or {}
    lines = [
        "# External Arabic OCR Score Report",
        "",
        "Scores use the same Arabic speech-readiness metrics as the app's OCR benchmark.",
        "",
        f"Best candidate: {markdown_value(best.get('label'))}",
        f"Best quality: {markdown_value(best.get('quality'))}",
        f"Best score: {markdown_value(best.get('qualityScore'))}",
        f"Baseline: {markdown_value(baseline.get('label'))}",
        f"Beats baseline: {markdown_value(comparison.get('beatsBaseline'))}",
        f"Promotion ready: {markdown_value(comparison.get('promotionReady'))}",
        "",
        "| Candidate | Quality | Score | Arabic Words | Speech Chars | Placeholder Ratio | Fragment Ratio | Source | Notes |",
        "| --- | --- | ---: | ---: | ---: | ---: | ---: | --- | --- |",
    ]
    for item in payload.get("results", []):
        notes = "; ".join(item.get("qualityReasons") or [])
        lines.append(
            "| "
            + " | ".join(
                [
                    markdown_value(item.get("label")),
                    markdown_value(item.get("quality")),
                    markdown_value(item.get("qualityScore")),
                    markdown_value(item.get("arabicWords")),
                    markdown_value(item.get("speechCharacters")),
                    markdown_value(item.get("placeholderRatio")),
                    markdown_value(item.get("fragmentLineRatio")),
                    markdown_value(item.get("path")),
                    markdown_value(notes),
                ]
            )
            + " |"
        )
    lines.extend(
        [
            "",
            "## Baseline Comparison",
            "",
            f"- Baseline score: {markdown_value(comparison.get('baselineScore'))}",
            f"- Best external score: {markdown_value(comparison.get('bestScore'))}",
            f"- Score delta: {markdown_value(comparison.get('scoreDelta'))}",
            f"- Baseline Arabic words: {markdown_value(comparison.get('baselineArabicWords'))}",
            f"- Best external Arabic words: {markdown_value(comparison.get('bestArabicWords'))}",
            f"- Arabic word delta: {markdown_value(comparison.get('arabicWordDelta'))}",
            "",
            "## Promotion Rule",
            "",
            "Promote an external OCR model only if its text is ready for TTS, its score beats the wired Arabic OCR benchmark on the same page images, and the worker can handle the model runtime.",
        ]
    )
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")


def print_table(payload: dict[str, Any]) -> None:
    print("candidate       quality  score    words   speech  fragments  source")
    print("--------------  -------  -------  ------  ------  ---------  ------")
    for item in payload["results"]:
        print(
            f"{item['label']:<14}  "
            f"{item.get('quality', '-'):<7}  "
            f"{item.get('qualityScore', 0):>7}  "
            f"{item.get('arabicWords', 0):>6}  "
            f"{item.get('speechCharacters', 0):>6}  "
            f"{item.get('fragmentLineRatio', 0):>9}  "
            f"{item.get('path', '-')}"
        )
    best = payload.get("best") or {}
    if best:
        print()
        print(f"Best external OCR candidate: {best.get('label')} quality={best.get('quality')} score={best.get('qualityScore')}")
    comparison = payload.get("comparison") or {}
    if comparison:
        print(
            f"Baseline comparison: beatsBaseline={comparison.get('beatsBaseline')} "
            f"scoreDelta={comparison.get('scoreDelta')} promotionReady={comparison.get('promotionReady')}"
        )


def main_cli() -> None:
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    parser = argparse.ArgumentParser(description="Score external Arabic OCR text outputs with the app's TTS-readiness metrics.")
    parser.add_argument(
        "--candidate",
        action="append",
        default=[],
        help="OCR text output as label=path or plain path. Directories concatenate sorted *.txt files.",
    )
    parser.add_argument("--write-report", type=Path, help="Optional Markdown report destination.")
    parser.add_argument("--write-json", type=Path, help="Optional JSON report destination for model_promotion_gate.py.")
    parser.add_argument(
        "--baseline-json",
        type=Path,
        help="Optional benchmark_ocr.py --json or prepare_book_workflow.py --json file for wired OCR baseline comparison.",
    )
    parser.add_argument("--json", action="store_true", help="Print JSON instead of a compact table.")
    args = parser.parse_args()

    payload = score_external_ocr(
        args.candidate,
        report_path=args.write_report,
        baseline_json=args.baseline_json,
        json_path=args.write_json,
    )
    if args.json:
        print(json.dumps(payload, ensure_ascii=False, indent=2))
    else:
        print_table(payload)


if __name__ == "__main__":
    main_cli()