clue_trainer.py

"""Load the single Chroniclues Parseword dataset and build few-shot examples.

The project uses one dataset only: cryptic_parseword_combined.csv

The loader detects the file format from its contents and requires exactly two
columns: input-text and target-text. Operations are inferred from the known
seven blocks of 100 rows.
"""
from __future__ import annotations

import csv
import importlib
import os
import random
import re
import zipfile
from io import BytesIO
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Iterable

BASE_DIR = Path(__file__).resolve().parent
DATASET_PATH = Path(
    os.getenv("CLUE_DATASET", str(BASE_DIR / "cryptic_parseword_combined.csv"))
).expanduser().resolve()

OPERATION_ORDER = [
    "hidden_word",
    "deletion",
    "reverse",
    "anagram",
    "container",
    "homophone",
    "selection",
]

BLOCK_SIZE = 100
EXPECTED_ROWS = BLOCK_SIZE * len(OPERATION_ORDER)
EXPECTED_COLUMNS = ["input-text", "target-text"]
MAX_PER_OP_BLOCK = int(os.getenv("MAX_PER_OP_BLOCK", "3"))
MAX_PER_OP_FORMULA = int(os.getenv("MAX_PER_OP_FORMULA", "5"))


def _clean_answer(value: Any) -> str:
    return re.sub(r"[^A-Z]", "", str(value or "").upper())


def _clean_clue(value: Any) -> str:
    return re.sub(r"\s+", " ", str(value or "")).strip()


def _is_xlsx(path: Path) -> bool:
    """Detect an XLSX workbook by file contents, not filename extension."""
    try:
        return zipfile.is_zipfile(path)
    except OSError:
        return False


def _read_xlsx(path: Path) -> tuple[list[str], list[dict[str, Any]]]:
    try:
        openpyxl = importlib.import_module("openpyxl")
        load_workbook = openpyxl.load_workbook
    except ImportError as exc:
        raise RuntimeError(
            "The clue dataset is an XLSX workbook. Install openpyxl: "
            "python -m pip install openpyxl"
        ) from exc

    workbook = load_workbook(
        BytesIO(path.read_bytes()),
        read_only=True,
        data_only=True,
    )

    if not workbook.worksheets:
        raise ValueError(f"Dataset workbook contains no worksheets: {path}")

    sheet = workbook.worksheets[0]
    iterator = sheet.iter_rows(values_only=True)

    try:
        headers = [str(value or "").strip() for value in next(iterator)]
    except StopIteration as exc:
        raise ValueError(f"Dataset is empty: {path}") from exc

    rows = [dict(zip(headers, values)) for values in iterator]
    return headers, rows


def _read_text_csv(path: Path) -> tuple[list[str], list[dict[str, Any]]]:
    with path.open("r", newline="", encoding="utf-8-sig") as handle:
        reader = csv.DictReader(handle)
        headers = [(value or "").strip() for value in (reader.fieldnames or [])]
        return headers, [dict(row) for row in reader]


def _read_source_rows(path: Path) -> tuple[list[str], list[dict[str, Any]]]:
    if _is_xlsx(path):
        return _read_xlsx(path)
    return _read_text_csv(path)


def load_dataset(path: Path = DATASET_PATH) -> list[dict[str, Any]]:
    """Load the single two-column dataset and assign operations by row block."""
    if not path.exists():
        raise FileNotFoundError(
            f"Clue dataset not found at {path}. Place "
            "cryptic_parseword_combined.csv beside clue_trainer.py or set "
            "CLUE_DATASET to its absolute path."
        )

    headers, source_rows = _read_source_rows(path)
    if headers != EXPECTED_COLUMNS:
        raise ValueError(
            f"Dataset must contain exactly {EXPECTED_COLUMNS}; found {headers}."
        )
    if len(source_rows) != EXPECTED_ROWS:
        raise ValueError(
            f"Dataset must contain exactly {EXPECTED_ROWS} data rows; "
            f"found {len(source_rows)}. Operations are inferred from seven "
            "100-row blocks."
        )

    rows: list[dict[str, Any]] = []
    for zero_index, source in enumerate(source_rows):
        operation = OPERATION_ORDER[zero_index // BLOCK_SIZE]
        clue_text = _clean_clue(source.get("input-text"))
        answer = _clean_answer(source.get("target-text"))
        rows.append(
            {
                "operation": operation,
                "clue_text": clue_text,
                "answer": answer,
                "source_row": zero_index + 2,
            }
        )
    return rows


def validate_dataset(rows: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
    """Run structural checks supported by this two-column dataset."""
    results: list[dict[str, Any]] = []
    supported = set(OPERATION_ORDER)

    for index, row in enumerate(rows, start=1):
        operation = str(row.get("operation", ""))
        clue_text = _clean_clue(row.get("clue_text"))
        answer = _clean_answer(row.get("answer"))
        errors: list[str] = []

        if operation not in supported:
            errors.append(f"unsupported operation '{operation}'")
        if not clue_text:
            errors.append("input-text is empty")
        if not answer:
            errors.append("target-text is empty")
        elif not 1 <= len(answer) <= 12:
            errors.append(f"answer length {len(answer)} is outside 1-12")

        results.append(
            {
                "index": index,
                "row": row,
                "operation": operation,
                "answer": answer,
                "valid": not errors,
                "level": "structural",
                "message": "; ".join(errors) if errors else "Required fields are present",
            }
        )
    return results


def _select_examples(
    rows: list[dict[str, Any]], max_per_op: int, seed: int = 42
) -> dict[str, list[dict[str, Any]]]:
    rng = random.Random(seed)
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        grouped[row["operation"]].append(row)

    selected: dict[str, list[dict[str, Any]]] = {}
    for operation in OPERATION_ORDER:
        pool = list(grouped.get(operation, []))
        rng.shuffle(pool)
        selected[operation] = pool[:max_per_op]
    return selected


def build_fewshot_block(examples: dict[str, list[dict[str, Any]]]) -> str:
    lines = ["PARSWORDS DATASET STYLE REFERENCE", ""]
    for operation in OPERATION_ORDER:
        lines.append(f"[{operation.upper()}]")
        for row in examples.get(operation, []):
            lines.append(f"- clue: {row['clue_text']}")
            lines.append(f"  answer: {row['answer']}")
        lines.append("")
    return "\n".join(lines).strip()


def build_clue_formula(examples: dict[str, list[dict[str, Any]]]) -> str:
    lines = ["PARSWORDS CLUE EXAMPLES", ""]
    for operation in OPERATION_ORDER:
        lines.append(f"[{operation.upper()}]")
        for row in examples.get(operation, []):
            lines.append(f"  {row['answer']:12s} -> {row['clue_text']}")
        lines.append("")
    return "\n".join(lines).strip()


def get_schema_examples(
    examples: dict[str, list[dict[str, Any]]]
) -> dict[str, str]:
    return {
        operation: rows[0]["clue_text"]
        for operation, rows in examples.items()
        if rows
    }


# Load exactly once at import time.
ROWS = load_dataset()
_VALIDATION = validate_dataset(ROWS)
VALID_ROWS = [item["row"] for item in _VALIDATION if item["valid"]]

VALIDATION_REPORT = {
    "total": len(_VALIDATION),
    "passed": sum(1 for item in _VALIDATION if item["valid"]),
    "failed": [item for item in _VALIDATION if not item["valid"]],
    "level": "structural",
}

CLUE_EXAMPLES = _select_examples(VALID_ROWS, MAX_PER_OP_BLOCK)
_FORMULA_EXAMPLES = _select_examples(VALID_ROWS, MAX_PER_OP_FORMULA)
FEWSHOT_BLOCK = build_fewshot_block(CLUE_EXAMPLES)
CLUE_FORMULA = build_clue_formula(_FORMULA_EXAMPLES)
SCHEMA_CLUE_EXAMPLES = get_schema_examples(CLUE_EXAMPLES)

DATASET_SUMMARY = {
    "dataset_path": str(DATASET_PATH),
    "dataset_format": "xlsx" if _is_xlsx(DATASET_PATH) else "csv",
    "total_rows": len(ROWS),
    "operation_counts": dict(Counter(row["operation"] for row in ROWS)),
    "validation_passed": VALIDATION_REPORT["passed"],
    "validation_total": VALIDATION_REPORT["total"],
}


if __name__ == "__main__":
    print(f"Dataset path: {DATASET_PATH}")
    print(f"Detected format: {DATASET_SUMMARY['dataset_format']}")
    print(f"Rows loaded: {len(ROWS)}")
    print(
        f"Validation: {VALIDATION_REPORT['passed']}/"
        f"{VALIDATION_REPORT['total']} structural checks passed"
    )
    for operation in OPERATION_ORDER:
        print(f"  {operation:14s} {DATASET_SUMMARY['operation_counts'].get(operation, 0)}")
    if VALIDATION_REPORT["failed"]:
        print("First 10 failures:")
        for failure in VALIDATION_REPORT["failed"][:10]:
            print(
                f"  source_row={failure['row'].get('source_row')} "
                f"operation={failure['operation']} "
                f"answer={failure['answer']} reason={failure['message']}"
            )