path: root/src/gap_pipeline/offline.py

"""No-API validation, released-data alignment, and run statistics."""

from __future__ import annotations

import json
import re
from collections import Counter
from pathlib import Path
from typing import Any

from .store import sha256_payload
from .surface import IDENTIFIER_RE, SURFACE_FAMILIES


REQUIRED_TOP_LEVEL = {
    "index",
    "question",
    "solution",
    "vars",
    "params",
    "sci_consts",
    "variants",
    "problem_type",
}


def load_dataset(dataset_dir: Path) -> dict[str, dict[str, Any]]:
    records: dict[str, dict[str, Any]] = {}
    for path in sorted(dataset_dir.glob("*.json")):
        record = json.loads(path.read_text(encoding="utf-8"))
        item_id = str(record.get("index", ""))
        if not item_id:
            raise ValueError(f"{path}: missing index")
        if item_id in records:
            raise ValueError(f"duplicate dataset ID {item_id}")
        records[item_id] = record
    return records


def normalize_latex_symbol(value: str) -> str:
    """Normalize notation aliases used in the released rename maps.

    The public records sometimes contain both ``x_n`` and ``x_{n}``, or use a
    different number of escaped backslashes between the symbol inventory and a
    map key.  These are representation aliases, not two mathematical symbols.
    This normalization is deliberately used only for diagnostics; it never
    rewrites released text.
    """

    normalized = re.sub(r"\\\\+", r"\\", str(value))
    normalized = re.sub(r"\s+", "", normalized)
    previous = None
    while previous != normalized:
        previous = normalized
        normalized = re.sub(r"_\{([^{}]+)\}", r"_\1", normalized)
    return normalized


def validate_public_dataset(dataset_dir: Path) -> dict[str, Any]:
    records = load_dataset(dataset_dir)
    errors: list[str] = []
    warnings: list[str] = []
    family_counts: Counter[str] = Counter()
    kernel_metadata_counts: Counter[str] = Counter()
    potential_target_collision_count = 0
    unchanged_mapping_count = 0
    tex_target_count = 0
    alias_source_count = 0
    source_outside_inventory_count = 0

    for item_id, record in records.items():
        missing = REQUIRED_TOP_LEVEL - set(record)
        if missing:
            errors.append(f"{item_id}: missing top-level fields {sorted(missing)}")
            continue
        if not record["question"].strip() or not record["solution"].strip():
            errors.append(f"{item_id}: empty canonical question or solution")
        expected_symbols = set(record.get("vars", [])) | set(record.get("params", []))
        normalized_expected = {
            normalize_latex_symbol(value) for value in expected_symbols
        }
        variants = record["variants"]

        for family in SURFACE_FAMILIES:
            family_counts[family] += int(family in variants)
            if family not in variants:
                errors.append(f"{item_id}: missing surface family {family}")
                continue
            variant = variants[family]
            missing_variant = {"map", "question", "solution"} - set(variant)
            if missing_variant:
                errors.append(
                    f"{item_id}/{family}: missing fields {sorted(missing_variant)}"
                )
                continue
            rename_map = variant["map"]
            canonical_sources: dict[str, str] = {}
            missing_sources: list[str] = []
            for source in rename_map:
                normalized_source = normalize_latex_symbol(source)
                canonical_sources[source] = normalized_source
                if source not in expected_symbols and normalized_source in normalized_expected:
                    alias_source_count += 1
                elif normalized_source not in normalized_expected:
                    missing_sources.append(source)
                    source_outside_inventory_count += 1
            if missing_sources:
                warnings.append(
                    f"{item_id}/{family}: map sources absent from normalized "
                    f"symbol inventory {sorted(missing_sources)}"
                )

            # Alias keys such as x_n and x_{n} are one source.  A collision is
            # counted only if two distinct normalized sources share a target.
            target_to_sources: dict[str, set[str]] = {}
            for source, replacement in rename_map.items():
                target_to_sources.setdefault(
                    normalize_latex_symbol(replacement), set()
                ).add(canonical_sources[source])
                if normalize_latex_symbol(source) == normalize_latex_symbol(
                    replacement
                ):
                    unchanged_mapping_count += 1
                if not IDENTIFIER_RE.fullmatch(str(replacement)):
                    tex_target_count += 1
                if not str(replacement).strip():
                    errors.append(f"{item_id}/{family}: empty replacement target")
            potential_target_collision_count += sum(
                len(sources) - 1
                for sources in target_to_sources.values()
                if len(sources) > 1
            )

        if "kernel_variant" not in variants:
            errors.append(f"{item_id}: missing kernel_variant")
        else:
            family_counts["kernel_variant"] += 1
            kernel = variants["kernel_variant"]
            if not str(kernel.get("question", "")).strip():
                errors.append(f"{item_id}/kernel_variant: empty question")
            if not str(kernel.get("solution", "")).strip():
                errors.append(f"{item_id}/kernel_variant: empty solution")
            if "_meta" in kernel:
                kernel_metadata_counts["_meta"] += 1
            if "metadata" in kernel:
                kernel_metadata_counts["metadata"] += 1
            if "_replacement_note" in kernel:
                kernel_metadata_counts["_replacement_note"] += 1
    return {
        "dataset_dir": str(dataset_dir.resolve()),
        "record_count": len(records),
        "expected_record_count": 1051,
        "record_count_matches": len(records) == 1051,
        "family_counts": dict(family_counts),
        "kernel_metadata_shapes": dict(kernel_metadata_counts),
        "surface_map_checks": {
            "potential_target_collision_count": potential_target_collision_count,
            "unchanged_mapping_count": unchanged_mapping_count,
            "tex_or_structured_target_count": tex_target_count,
            "notation_alias_source_count": alias_source_count,
            "source_outside_normalized_inventory_count": (
                source_outside_inventory_count
            ),
            "interpretation": (
                "TeX/structured targets, notation aliases, and potential "
                "target collisions are descriptive diagnostics, not validation "
                "errors. Strict rules apply only to newly generated maps."
            ),
        },
        "errors": errors,
        "warnings": warnings,
        "valid": not errors,
    }


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


def summarize_run(run_dir: Path) -> dict[str, Any]:
    finals = sorted(run_dir.glob("items/*/final.json"))
    counts: Counter[str] = Counter()
    iterations: list[int] = []
    first_round_failed = 0
    repaired_accepts = 0
    cap_reached = 0
    accepted_ids: list[str] = []
    rejected_ids: list[str] = []

    for path in finals:
        payload = json.loads(path.read_text(encoding="utf-8"))
        status = str(payload["status"])
        counts[status] += 1
        item_id = str(payload["item_id"])
        history = payload.get("iterations", [])
        iterations.append(len(history))
        if history and not history[0].get("unanimous", False):
            first_round_failed += 1
        if status == "accepted":
            accepted_ids.append(item_id)
            if any(row.get("repaired", False) for row in history):
                repaired_accepts += 1
        else:
            rejected_ids.append(item_id)
            if len(history) >= 15:
                cap_reached += 1

    return {
        "run_dir": str(run_dir.resolve()),
        "candidate_count": len(finals),
        "status_counts": dict(counts),
        "failed_first_unanimous_vote": first_round_failed,
        "accepted_after_repair": repaired_accepts,
        "finally_rejected": counts["rejected"],
        "reached_iteration_cap": cap_reached,
        "iteration_min": min(iterations) if iterations else None,
        "iteration_max": max(iterations) if iterations else None,
        "iteration_mean": sum(iterations) / len(iterations) if iterations else None,
        "accepted_ids": accepted_ids,
        "rejected_ids": rejected_ids,
    }


def align_run_to_release(run_dir: Path, dataset_dir: Path) -> dict[str, Any]:
    release = load_dataset(dataset_dir)
    rows = []
    for path in sorted(run_dir.glob("items/*/final.json")):
        payload = json.loads(path.read_text(encoding="utf-8"))
        if payload.get("status") != "accepted":
            continue
        item_id = str(payload["item_id"])
        candidate = payload["accepted_candidate"]
        if item_id not in release:
            rows.append({"item_id": item_id, "status": "missing_from_release"})
            continue
        released_kernel = release[item_id]["variants"]["kernel_variant"]
        question_exact = candidate["problem"] == released_kernel["question"]
        solution_exact = candidate["proof"] == released_kernel["solution"]
        question_normalized = normalize_text(candidate["problem"]) == normalize_text(
            released_kernel["question"]
        )
        solution_normalized = normalize_text(candidate["proof"]) == normalize_text(
            released_kernel["solution"]
        )
        rows.append(
            {
                "item_id": item_id,
                "status": (
                    "exact_match"
                    if question_exact and solution_exact
                    else "normalized_match"
                    if question_normalized and solution_normalized
                    else "mismatch"
                ),
                "question_exact": question_exact,
                "solution_exact": solution_exact,
                "question_normalized": question_normalized,
                "solution_normalized": solution_normalized,
                "candidate_sha256": sha256_payload(
                    {
                        "question": candidate["problem"],
                        "solution": candidate["proof"],
                    }
                ),
                "release_sha256": sha256_payload(
                    {
                        "question": released_kernel["question"],
                        "solution": released_kernel["solution"],
                    }
                ),
            }
        )
    return {
        "run_dir": str(run_dir.resolve()),
        "dataset_dir": str(dataset_dir.resolve()),
        "checked": len(rows),
        "status_counts": dict(Counter(row["status"] for row in rows)),
        "rows": rows,
    }