dataforge/bench/methods.py

"""Benchmark method implementations for DataForge."""

from __future__ import annotations

import json
import random
import time
from math import ceil
from statistics import median
from typing import TYPE_CHECKING, cast

from dataforge.bench.core import BenchmarkRepair, SeedBenchmarkResult, quota_units, score_repairs
from dataforge.bench.error_classes import (
    expected_calibration_error,
    precision_at_auto_apply,
    score_repairs_by_class,
)
from dataforge.bench.groq_client import BenchLLMClient, CostCapExceededError, ProviderRequestError
from dataforge.datasets.real_world import RealWorldDataset
from dataforge.detectors import run_all_detectors
from dataforge.repairers import propose_fixes
from dataforge.repairers.llm_corrector import LLMCorrectorRepairer
from dataforge.schema_inference import infer_schema

if TYPE_CHECKING:
    from dataforge.agent.providers import Message

# The corrector auto-apply diagnostic uses a fixed, pre-committed self-consistency
# agreement threshold (not an in-sample fit): precision_at_auto_apply reports the
# correction precision among proposals whose agreement clears this bar, which is
# the honest "if we auto-applied only high-agreement outputs" signal.
_CORRECTOR_AUTO_APPLY_CONFIDENCE = 0.95
_CORRECTOR_SAMPLES = 3


def _reproduction_command(method: str, dataset: str, seeds: int) -> str:
    """Build the canonical reproduction command for one method/dataset selection."""
    return f"dataforge bench --methods {method} --datasets {dataset} --seeds {seeds}"


def _repairs_from_proposed_fixes(
    dataset: RealWorldDataset,
) -> tuple[list[BenchmarkRepair], set[tuple[int, str]]]:
    """Run the shipped deterministic detector/repair stack on one dataset.

    Returns both the proposed repairs and the set of cells the detectors flagged
    (detection coverage), so detection and correction can be scored separately.
    """
    inferred_schema = infer_schema(dataset.dirty_df.copy(deep=True)).to_schema(
        include_inferred_constraints=True
    )
    issues = run_all_detectors(dataset.dirty_df.copy(deep=True), schema=inferred_schema)
    detected_cells = {(issue.row, issue.column) for issue in issues}
    proposals = propose_fixes(
        issues,
        dataset.dirty_df.copy(deep=True),
        inferred_schema,
        cache_dir=None,
        allow_llm=False,
    )
    repairs = [
        BenchmarkRepair(
            row=proposal.fix.row,
            column=proposal.fix.column,
            new_value=proposal.fix.new_value,
            reason=proposal.reason,
        )
        for proposal in proposals
    ]
    return repairs, detected_cells


def run_heuristic_episode(dataset: RealWorldDataset, *, seed: int) -> SeedBenchmarkResult:
    """Run the current deterministic DataForge stack as the heuristic baseline."""
    start = time.perf_counter()
    repairs, detected_cells = _repairs_from_proposed_fixes(dataset)
    metrics = score_repairs(dataset.ground_truth, repairs)
    by_class = score_repairs_by_class(dataset.ground_truth, repairs, detected_cells)
    runtime_s = round(time.perf_counter() - start, 4)
    return SeedBenchmarkResult(
        method="heuristic",
        by_class=by_class,
        dataset=dataset.metadata.name,
        seed=seed,
        status="ok",
        precision=metrics.precision,
        recall=metrics.recall,
        f1=metrics.f1,
        tp=metrics.tp,
        fp=metrics.fp,
        fn=metrics.fn,
        avg_steps=float(1 + len(repairs)),
        llm_calls=0,
        prompt_tokens=0,
        completion_tokens=0,
        quota_units=0.0,
        runtime_s=runtime_s,
        provider="local",
        model="deterministic",
        reproduction_command=_reproduction_command("heuristic", dataset.metadata.name, 1),
    )


def run_llm_corrector_episode(
    dataset: RealWorldDataset,
    *,
    seed: int,
    client: BenchLLMClient,
    samples: int = _CORRECTOR_SAMPLES,
    max_issues: int | None = None,
) -> SeedBenchmarkResult:
    """Run the grounded, contract-bound LLM corrector as a benchmark method.

    Dirty-only: detection uses the inferred schema (same as the heuristic
    baseline), then the corrector proposes an exact value per detected issue.
    Every LLM call is routed through the quota-tracking ``client`` so calls,
    tokens, and latency are measured. Beyond correction F1 and per-class
    coverage, this reports calibration quality (ECE) and the precision the tool
    would achieve if it auto-applied only high-agreement outputs
    (``precision_at_auto_apply`` at a fixed agreement threshold).

    When ``max_issues`` is set and the detector finds more issues than the cap,
    a deterministic ``seed``-derived random subset of that size is scored. This
    bounds LLM spend for real-provider runs; the returned warnings record the
    sampling so the report stays honest about coverage.
    """
    start = time.perf_counter()
    counters = {"llm_calls": 0, "prompt_tokens": 0, "completion_tokens": 0}
    warnings: list[str] = []

    def _adapter(messages: list[Message]) -> str:
        completion = client.complete(cast("list[dict[str, str]]", messages))
        counters["llm_calls"] += 1
        counters["prompt_tokens"] += completion.prompt_tokens
        counters["completion_tokens"] += completion.completion_tokens
        warnings.extend(completion.warnings)
        return completion.text

    inferred_schema = infer_schema(dataset.dirty_df.copy(deep=True)).to_schema(
        include_inferred_constraints=True
    )
    working = dataset.dirty_df.copy(deep=True)
    issues = run_all_detectors(working, schema=inferred_schema)
    if max_issues is not None and len(issues) > max_issues:
        warnings.append(f"corrector_sampled_{max_issues}_of_{len(issues)}")
        sampler = random.Random(seed)
        issues = sorted(
            sampler.sample(issues, max_issues),
            key=lambda issue: (issue.row, issue.column),
        )
    detected_cells = {(issue.row, issue.column) for issue in issues}

    corrector = LLMCorrectorRepairer(
        cache_dir=None,
        allow_llm=True,
        model=client.model,
        samples=samples,
        completion_fn=_adapter,
    )
    ground_truth = {(cell.row, cell.column): cell.clean_value for cell in dataset.ground_truth}

    repairs: list[BenchmarkRepair] = []
    calibration_samples: list[tuple[float, bool]] = []
    auto_apply_samples: list[tuple[bool, bool]] = []
    call_failures = 0
    # The corrector runs against the schema-less product path (schema=None) so it
    # infers its own high-confidence verification constraints, mirroring how a
    # user with an undeclared CSV would run it.
    for issue in issues:
        try:
            fix = corrector.propose(issue, working, None)
        except CostCapExceededError:
            # A tripped spend guard is a hard stop; never swallow it.
            raise
        except (ProviderRequestError, TimeoutError):
            # A single throttled/failed provider call must not abort the whole
            # run: skip this issue (counts as no repair) and keep going so the
            # benchmark still produces a report on flaky-quota accounts.
            call_failures += 1
            continue
        if fix is None:
            continue
        new_value = fix.fix.new_value
        repairs.append(
            BenchmarkRepair(
                row=fix.fix.row,
                column=fix.fix.column,
                new_value=new_value,
                reason=fix.reason,
            )
        )
        was_correct = ground_truth.get((fix.fix.row, fix.fix.column)) == new_value
        calibration_samples.append((fix.confidence, was_correct))
        auto_applied = fix.confidence >= _CORRECTOR_AUTO_APPLY_CONFIDENCE
        auto_apply_samples.append((auto_applied, was_correct))

    if call_failures:
        warnings.append(f"corrector_call_failures_{call_failures}")

    metrics = score_repairs(dataset.ground_truth, repairs)
    by_class = score_repairs_by_class(dataset.ground_truth, repairs, detected_cells)
    runtime_s = round(time.perf_counter() - start, 4)
    return SeedBenchmarkResult(
        method="llm_corrector",
        by_class=by_class,
        dataset=dataset.metadata.name,
        seed=seed,
        status="ok",
        precision=metrics.precision,
        recall=metrics.recall,
        f1=metrics.f1,
        tp=metrics.tp,
        fp=metrics.fp,
        fn=metrics.fn,
        avg_steps=float(counters["llm_calls"]),
        llm_calls=counters["llm_calls"],
        prompt_tokens=counters["prompt_tokens"],
        completion_tokens=counters["completion_tokens"],
        quota_units=quota_units(
            llm_calls=counters["llm_calls"],
            prompt_tokens=counters["prompt_tokens"],
            completion_tokens=counters["completion_tokens"],
        ),
        runtime_s=runtime_s,
        provider=client.provider,
        model=client.model,
        warnings=warnings,
        ece=expected_calibration_error(calibration_samples),
        precision_at_auto_apply=precision_at_auto_apply(auto_apply_samples),
        auto_apply_count=sum(1 for auto_applied, _ in auto_apply_samples if auto_applied),
        reproduction_command=_reproduction_command("llm_corrector", dataset.metadata.name, 1),
    )


def corrector_promotion_verdict(
    record: SeedBenchmarkResult,
    *,
    min_precision_at_auto_apply: float = 0.95,
    max_ece: float = 0.1,
    min_auto_apply_count: int = 1,
) -> tuple[bool, list[str]]:
    """Decide whether corrector outputs may be promoted to auto-apply.

    Promotion requires measured evidence: enough high-agreement proposals to
    judge, a precision at that agreement clearing the floor, and a calibrated
    (low-ECE) confidence signal. Anything short of that stays propose-not-apply.
    """
    reasons: list[str] = []
    if record.precision_at_auto_apply is None or record.auto_apply_count is None:
        return False, ["record does not carry corrector auto-apply metrics"]
    if record.auto_apply_count < min_auto_apply_count:
        reasons.append(
            f"only {record.auto_apply_count} high-agreement proposal(s); "
            f"need >= {min_auto_apply_count}"
        )
    if record.precision_at_auto_apply < min_precision_at_auto_apply:
        reasons.append(
            f"precision_at_auto_apply {record.precision_at_auto_apply} < "
            f"{min_precision_at_auto_apply}"
        )
    if record.ece is not None and record.ece > max_ece:
        reasons.append(f"ECE {record.ece} > {max_ece} (confidence poorly calibrated)")
    return (not reasons), reasons


def run_random_episode(dataset: RealWorldDataset, *, seed: int) -> SeedBenchmarkResult:
    """Run the bounded random baseline on one dataset."""
    rng = random.Random(seed)
    start = time.perf_counter()
    budget = min(200, max(25, ceil(len(dataset.ground_truth) / 10)))
    column_values = {
        column: [str(value) for value in dataset.dirty_df[column].tolist()]
        for column in dataset.canonical_columns
    }
    repairs: list[BenchmarkRepair] = []
    for _ in range(budget):
        row_index = rng.randrange(len(dataset.dirty_df.index))
        column = rng.choice(dataset.canonical_columns)
        new_value = rng.choice(column_values[column])
        repairs.append(
            BenchmarkRepair(
                row=row_index,
                column=column,
                new_value=new_value,
                reason="random baseline",
            )
        )
    metrics = score_repairs(dataset.ground_truth, repairs)
    by_class = score_repairs_by_class(dataset.ground_truth, repairs)
    runtime_s = round(time.perf_counter() - start, 4)
    return SeedBenchmarkResult(
        method="random",
        by_class=by_class,
        dataset=dataset.metadata.name,
        seed=seed,
        status="ok",
        precision=metrics.precision,
        recall=metrics.recall,
        f1=metrics.f1,
        tp=metrics.tp,
        fp=metrics.fp,
        fn=metrics.fn,
        avg_steps=float(budget),
        llm_calls=0,
        prompt_tokens=0,
        completion_tokens=0,
        quota_units=0.0,
        runtime_s=runtime_s,
        provider="local",
        model="random",
        reproduction_command=_reproduction_command("random", dataset.metadata.name, 1),
    )


def _chunk_records(dataset: RealWorldDataset, row_indices: tuple[int, ...]) -> list[dict[str, str]]:
    """Serialize one row chunk for prompting."""
    records: list[dict[str, str]] = []
    for row_index in row_indices:
        row_payload: dict[str, str] = {"_row": str(row_index)}
        for column in dataset.canonical_columns:
            row_payload[column] = str(dataset.dirty_df.iloc[row_index][column])
        records.append(row_payload)
    return records


def _column_stats(
    dataset: RealWorldDataset, columns: list[str]
) -> dict[str, dict[str, str | float | int]]:
    """Return simple benchmark-local column statistics for ReAct prompting."""
    stats: dict[str, dict[str, str | float | int]] = {}
    for column in columns:
        series = dataset.dirty_df[column].astype(str)
        non_empty = [value for value in series.tolist() if value != ""]
        numeric_values: list[float] = []
        for value in non_empty:
            try:
                numeric_values.append(float(value))
            except ValueError:
                continue
        stats[column] = {
            "non_empty_count": len(non_empty),
            "unique_count": len(set(non_empty)),
        }
        if numeric_values:
            stats[column]["median"] = round(float(median(numeric_values)), 4)
    return stats


def _strip_json_line_comments(text: str) -> str:
    """Remove JavaScript-style line comments outside JSON strings."""
    result: list[str] = []
    in_string = False
    escaped = False
    index = 0
    while index < len(text):
        char = text[index]
        next_char = text[index + 1] if index + 1 < len(text) else ""
        if in_string:
            result.append(char)
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
            index += 1
            continue
        if char == '"':
            in_string = True
            result.append(char)
            index += 1
            continue
        if char == "/" and next_char == "/":
            index += 2
            while index < len(text) and text[index] not in "\r\n":
                index += 1
            continue
        result.append(char)
        index += 1
    return "".join(result)


def _extract_json_object(text: str) -> dict[str, object] | None:
    """Parse the first JSON object found in an LLM response string."""
    stripped = text.strip()
    if stripped.startswith("```"):
        stripped = stripped.strip("`")
        if stripped.lower().startswith("json"):
            stripped = stripped[4:].strip()
    stripped = _strip_json_line_comments(stripped)
    decoder = json.JSONDecoder()
    for offset, char in enumerate(stripped):
        if char != "{":
            continue
        try:
            payload, _ = decoder.raw_decode(stripped[offset:])
        except json.JSONDecodeError:
            continue
        if isinstance(payload, dict):
            return payload
    return None


def _repairs_from_payload(payload: dict[str, object]) -> list[BenchmarkRepair]:
    """Convert a parsed JSON payload into benchmark repairs."""
    raw_repairs = payload.get("repairs", [])
    if not isinstance(raw_repairs, list):
        return []
    repairs: list[BenchmarkRepair] = []
    for raw_repair in raw_repairs:
        if not isinstance(raw_repair, dict):
            continue
        row = raw_repair.get("row")
        column = raw_repair.get("column")
        new_value = raw_repair.get("new_value")
        reason = raw_repair.get("reason", "LLM repair")
        if (
            not isinstance(row, int)
            or not isinstance(column, str)
            or not isinstance(new_value, str)
        ):
            continue
        repairs.append(
            BenchmarkRepair(
                row=row,
                column=column,
                new_value=new_value,
                reason=str(reason),
            )
        )
    return repairs


def run_llm_zeroshot_episode(
    dataset: RealWorldDataset,
    *,
    seed: int,
    client: BenchLLMClient,
) -> SeedBenchmarkResult:
    """Run the zero-shot Groq baseline across fixed contiguous row chunks."""
    start = time.perf_counter()
    llm_calls = 0
    prompt_tokens = 0
    completion_tokens = 0
    warnings: list[str] = []
    repairs: list[BenchmarkRepair] = []

    for row_indices in chunk_row_indices(len(dataset.dirty_df.index)):
        chunk_payload = _chunk_records(dataset, row_indices)
        messages = [
            {
                "role": "system",
                "content": (
                    "You are benchmarking tabular data cleaning. Reply with strict JSON: "
                    '{"repairs":[{"row":0,"column":"Column","new_value":"value","reason":"why"}]}.'
                ),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "dataset": dataset.metadata.name,
                        "columns": list(dataset.canonical_columns),
                        "rows": chunk_payload,
                    },
                    sort_keys=True,
                ),
            },
        ]
        completion = client.complete(messages)
        llm_calls += 1
        prompt_tokens += completion.prompt_tokens
        completion_tokens += completion.completion_tokens
        warnings.extend(list(completion.warnings))
        parsed = _extract_json_object(completion.text)
        if parsed is not None:
            repairs.extend(_repairs_from_payload(parsed))

    metrics = score_repairs(dataset.ground_truth, repairs)
    by_class = score_repairs_by_class(dataset.ground_truth, repairs)
    runtime_s = round(time.perf_counter() - start, 4)
    return SeedBenchmarkResult(
        method="llm_zeroshot",
        by_class=by_class,
        dataset=dataset.metadata.name,
        seed=seed,
        status="ok",
        precision=metrics.precision,
        recall=metrics.recall,
        f1=metrics.f1,
        tp=metrics.tp,
        fp=metrics.fp,
        fn=metrics.fn,
        avg_steps=float(llm_calls),
        llm_calls=llm_calls,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        quota_units=quota_units(
            llm_calls=llm_calls,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
        ),
        runtime_s=runtime_s,
        provider="groq",
        model=client.model,
        warnings=warnings,
        reproduction_command=_reproduction_command("llm_zeroshot", dataset.metadata.name, 1),
    )


def run_llm_react_episode(
    dataset: RealWorldDataset,
    *,
    seed: int,
    client: BenchLLMClient,
) -> SeedBenchmarkResult:
    """Run the constrained ReAct-style Groq baseline with one optional tool step."""
    start = time.perf_counter()
    llm_calls = 0
    tool_calls = 0
    prompt_tokens = 0
    completion_tokens = 0
    warnings: list[str] = []
    repairs: list[BenchmarkRepair] = []

    for row_indices in chunk_row_indices(len(dataset.dirty_df.index)):
        chunk_payload = _chunk_records(dataset, row_indices)
        schema_summary = {
            "dataset": dataset.metadata.name,
            "columns": list(dataset.canonical_columns),
            "chunk_rows": len(row_indices),
        }
        messages = [
            {
                "role": "system",
                "content": (
                    "You are benchmarking tabular data cleaning with a constrained tool loop. "
                    "Respond with one JSON action object. Allowed actions: "
                    "inspect_rows, column_stats, submit_repairs, finish."
                ),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "schema_summary": schema_summary,
                        "rows": chunk_payload,
                    },
                    sort_keys=True,
                ),
            },
        ]
        first = client.complete(messages)
        llm_calls += 1
        prompt_tokens += first.prompt_tokens
        completion_tokens += first.completion_tokens
        warnings.extend(list(first.warnings))
        first_payload = _extract_json_object(first.text)
        if first_payload is None:
            continue

        action = first_payload.get("action")
        if action == "submit_repairs":
            repairs.extend(_repairs_from_payload(first_payload))
            continue
        if action == "finish":
            continue

        tool_result: dict[str, object]
        if action == "inspect_rows":
            requested_rows = first_payload.get("row_indices", [])
            if not isinstance(requested_rows, list):
                requested_rows = []
            safe_rows = [
                row for row in requested_rows if isinstance(row, int) and row in row_indices
            ]
            tool_result = {"rows": _chunk_records(dataset, tuple(safe_rows))}
        elif action == "column_stats":
            requested_columns = first_payload.get("columns", [])
            if not isinstance(requested_columns, list):
                requested_columns = []
            safe_columns = [
                column
                for column in requested_columns
                if isinstance(column, str) and column in dataset.canonical_columns
            ]
            tool_result = {"column_stats": _column_stats(dataset, safe_columns)}
        else:
            continue
        tool_calls += 1
        messages.append({"role": "assistant", "content": first.text})
        messages.append({"role": "user", "content": json.dumps(tool_result, sort_keys=True)})
        second = client.complete(messages)
        llm_calls += 1
        prompt_tokens += second.prompt_tokens
        completion_tokens += second.completion_tokens
        warnings.extend(list(second.warnings))
        second_payload = _extract_json_object(second.text)
        if second_payload is not None and second_payload.get("action") == "submit_repairs":
            repairs.extend(_repairs_from_payload(second_payload))

    metrics = score_repairs(dataset.ground_truth, repairs)
    by_class = score_repairs_by_class(dataset.ground_truth, repairs)
    runtime_s = round(time.perf_counter() - start, 4)
    return SeedBenchmarkResult(
        method="llm_react",
        by_class=by_class,
        dataset=dataset.metadata.name,
        seed=seed,
        status="ok",
        precision=metrics.precision,
        recall=metrics.recall,
        f1=metrics.f1,
        tp=metrics.tp,
        fp=metrics.fp,
        fn=metrics.fn,
        avg_steps=float(llm_calls + tool_calls),
        llm_calls=llm_calls,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        quota_units=quota_units(
            llm_calls=llm_calls,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
        ),
        runtime_s=runtime_s,
        provider="groq",
        model=client.model,
        warnings=warnings,
        reproduction_command=_reproduction_command("llm_react", dataset.metadata.name, 1),
    )


def chunk_row_indices(n_rows: int) -> tuple[tuple[int, ...], ...]:
    """Local import wrapper that avoids circular imports in the LLM helpers."""
    from dataforge.bench.core import chunk_row_indices as _chunk_row_indices

    return _chunk_row_indices(n_rows)