experiments/bioagent_bench/evaluate_bioagent_bench.py

#!/usr/bin/env python3
"""Evaluate Hypo_Bio_OS BioAgent Bench runs.

This evaluator follows the BioAgent Bench paper's grader design:
evaluate each trial as a pipeline execution, prioritize demonstrable
pipeline completion over exact numeric agreement, and return an
EvaluationResults-style JSON object.
"""

from __future__ import annotations

import argparse
import csv
import gzip
import json
import os
import re
from bisect import bisect_right
from datetime import timezone, datetime
from difflib import SequenceMatcher
from pathlib import Path
from statistics import mean
from typing import Any


DATASET_ROOT = Path("/225040511/project/bioagent-bench/dataset")
METADATA_PATH = Path("/225040511/project/bioagent-bench/src/task_metadata.json")
DEFAULT_RUNS_ROOT = Path("/225040511/project/Hypo_Bio_OS/bioagent-bench-runs")


FLEXIBLE_TABLE_MATCH_CONFIGS: dict[tuple[str, str], dict[str, Any]] = {
    (
        "alzheimer-mouse",
        "pathway_comparison.csv",
    ): {
        "soft_key_columns": ["Pathway"],
        "threshold": 0.62,
        "text_weight": 0.9,
        "numeric_weight": 0.1,
    },
    (
        "comparative-genomics",
        "cluster_annotation_mapping.csv",
    ): {
        "soft_key_columns": ["consensus_annotation"],
        "threshold": 0.68,
    },
    (
        "metagenomics",
        "phylum_relative_abundances.csv",
    ): {
        "soft_key_columns": ["Phylum"],
        "threshold": 0.9,
        "text_weight": 1.0,
        "numeric_weight": 0.0,
    },
    (
        "single-cell",
        "all_clusters_de_genes.csv",
    ): {
        "soft_key_columns": ["gene_name", "cluster_id"],
        "threshold": 0.78,
        "text_weight": 0.85,
        "numeric_weight": 0.15,
    },
    (
        "viral-metagenomics",
        "taxonomy.csv",
    ): {
        "soft_key_columns": ["domain", "species"],
        "threshold": 0.72,
        "text_weight": 0.85,
        "numeric_weight": 0.15,
    },
}


TASK_CONFIGS: dict[str, dict[str, Any]] = {
    "alzheimer-mouse": {
        "truth_files": ["pathway_comparison.csv"],
        "result_files": ["pathway_comparison.csv"],
        "key_columns": ["Pathway"],
        "numeric_columns": ["5xFAD_pvalue", "3xTG_AD_pvalue", "PS3O1S_pvalue"],
        "pipeline_steps": [
            "inspect mouse count/DEA inputs",
            "prepare metadata and count matrices for 5xFAD and 3xTG-AD",
            "perform differential expression for 5xFAD",
            "perform differential expression for 3xTG-AD",
            "use provided PS3O1S differential expression results",
            "run pathway enrichment per model",
            "merge shared/comparative pathway p-values into final CSV",
        ],
        "results_match_guidance": (
            "Treat mouse mmu/Mus musculus pathway labels as semantically compatible with hsa/Homo sapiens labels "
            "when the pathway identity is the same. Do not require exact p-value equality if the pipeline is plausible."
        ),
    },
    "comparative-genomics": {
        "truth_files": ["cluster_annotation_mapping.csv"],
        "result_files": ["cluster_annotation_mapping.csv"],
        "key_columns": ["cluster_number", "consensus_annotation"],
        "numeric_columns": [],
        "pipeline_steps": [
            "inspect Micrococcus FASTA/GFF/reference inputs",
            "predict or extract protein-coding genes",
            "identify orthologous/co-evolving clusters across genomes",
            "filter clusters present across intended genomes and coding-only",
            "assign high-confidence consensus annotations",
            "write cluster_number,consensus_annotation CSV",
        ],
        "results_match_guidance": (
            "Prioritize whether the output represents conserved annotated coding clusters. Exact cluster numbering "
            "can vary by method, but placeholder annotations or unrelated organisms should not match."
        ),
    },
    "cystic-fibrosis": {
        "truth_files": ["cf_variants.csv"],
        "result_files": ["cf_variants.csv"],
        "key_columns": ["chromosome", "position", "reference", "alternate"],
        "numeric_columns": [],
        "verifiable": True,
        "pipeline_steps": [
            "inspect family description and variant VCF",
            "filter variants by recessive inheritance in affected siblings",
            "exclude variants inconsistent with unaffected relatives/parents",
            "annotate candidate variant with ClinVar/reference metadata",
            "write the requested causal-variant CSV schema",
        ],
        "results_match_guidance": (
            "The result should identify the CFTR pathogenic recessive variant. Exact textual disease lists may differ, "
            "but chromosome, position, ref/alt, gene, and clinical interpretation must be consistent."
        ),
    },
    "deseq": {
        "truth_files": ["up_regulated_genes.csv"],
        "result_files": ["up_regulated_genes.csv"],
        "key_columns": ["gene_id"],
        "numeric_columns": ["log2FoldChange", "pvalue", "padj"],
        "pipeline_steps": [
            "inspect RNA-seq reads and Candida reference files",
            "prepare genome annotation/index",
            "align reads or otherwise quantify genes",
            "count reads per gene",
            "construct biofilm/planktonic sample metadata",
            "run differential expression",
            "filter up-regulated significant genes and write final CSV",
        ],
        "results_match_guidance": (
            "Prioritize evidence of a complete RNA-seq DE pipeline and a plausible up-regulated gene table. "
            "Do not require exact equality for all p-values/log fold changes."
        ),
    },
    "evolution": {
        "truth_files": ["variants_shared.csv", "gene_annotations.csv"],
        "result_files": ["variants_shared.csv", "gene_annotations.csv"],
        "key_columns": {
            "variants_shared.csv": ["CHROM", "POS", "REF", "ALT"],
            "gene_annotations.csv": ["Gene_Name"],
        },
        "numeric_columns": {"variants_shared.csv": [], "gene_annotations.csv": []},
        "pipeline_steps": [
            "inspect ancestor/evolved-line reads",
            "prepare or identify valid E. coli reference/assembly",
            "align ancestor and evolved reads",
            "call variants for all samples",
            "identify variants shared by evolved lines and absent from ancestor",
            "annotate variant/gene effects",
            "write variants_shared.csv and gene_annotations.csv",
        ],
        "results_match_guidance": (
            "Reward a biologically coherent shared-variant workflow. Penalize using forbidden external/sibling references "
            "or hallucinated annotations even if the final schema exists."
        ),
    },
    "giab": {
        "truth_files": ["HG001_GRCh38_1_22_v4.2.1_benchmark.vcf.gz"],
        "result_files": ["predicted.vcf.gz"],
        "vcf": True,
        "verifiable": True,
        "pipeline_steps": [
            "inspect paired reads, BED targets, and GRCh38 reference",
            "align reads to reference",
            "sort/index BAM",
            "mark duplicates or prepare analysis-ready BAM",
            "call variants",
            "compress/index final VCF if needed",
            "write predicted.vcf.gz",
        ],
        "results_match_guidance": (
            "Use GIAB variant concordance as the correctness signal. The f1_score field should reflect variant-level "
            "overlap where available."
        ),
    },
    "metagenomics": {
        "truth_files": ["phylum_relative_abundances.csv"],
        "result_files": ["phylum_relative_abundances.csv"],
        "key_columns": ["OTU"],
        "numeric_columns": ["JP4D", "JC1A"],
        "pipeline_steps": [
            "inspect paired metagenomic reads and reference database",
            "classify reads taxonomically",
            "aggregate classifications at bacterial phylum level",
            "normalize relative abundances for JP4D and JC1A",
            "write OTU,Kingdom,Phylum,JP4D,JC1A CSV",
        ],
        "results_match_guidance": (
            "Prioritize correct use of the bacterial/metagenomic reference and plausible phylum-level abundance table. "
            "Small abundance differences are acceptable."
        ),
    },
    "single-cell": {
        "truth_files": ["all_clusters_de_genes.csv"],
        "result_files": ["all_clusters_de_genes.csv"],
        "key_columns": ["cluster_id", "gene_name"],
        "numeric_columns": ["logfoldchanges", "pvals", "pvals_adj", "abs_logfc"],
        "pipeline_steps": [
            "load 10X matrices and metadata",
            "perform QC/normalization",
            "cluster cells and preserve cluster IDs",
            "annotate cell types using marker evidence",
            "compare pre/post exercise within cell types or clusters",
            "write all significant DE genes in requested schema",
        ],
        "results_match_guidance": (
            "Cell-type labels and cluster identities may vary, so prioritize marker-supported annotation, DE evidence, "
            "and schema correctness over exact row equality."
        ),
    },
    "transcript-quant": {
        "truth_files": ["truth.tsv"],
        "result_files": ["truth.tsv"],
        "tsv_no_header": True,
        "key_columns": ["transcript_id"],
        "numeric_columns": ["count"],
        "verifiable": True,
        "pipeline_steps": [
            "inspect paired RNA-seq reads and transcriptome reference",
            "build transcriptome index",
            "quantify transcript abundance/counts",
            "extract transcript_id and count columns",
            "write no-header two-column TSV",
        ],
        "results_match_guidance": (
            "Because the data is simulated, counts should closely reproduce the truth. Headerless two-column TSV format "
            "is required."
        ),
    },
    "viral-metagenomics": {
        "truth_files": ["taxonomy.csv"],
        "result_files": ["taxonomy.csv"],
        "key_columns": ["domain", "species"],
        "numeric_columns": ["contig_count"],
        "verifiable": True,
        "pipeline_steps": [
            "inspect dolphin metagenomic reads and viral reference resources",
            "assemble reads into contigs",
            "classify contigs against compatible viral reference resources",
            "aggregate contig counts by domain/species",
            "write contig_count,domain,species CSV",
        ],
        "results_match_guidance": (
            "Prioritize using this task's viral reference resources and correctly identifying viral species. "
            "Unclassified contigs are allowed when supported by classification output."
        ),
    },
}


def utc_timestamp() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")


def normalize_text(value: Any) -> str:
    return " ".join(str(value).strip().split()).lower()


def normalize_for_soft_match(value: Any) -> str:
    text = normalize_text(value)
    text = re.sub(r"\b(homo sapiens|mus musculus|human|mouse)\b", " ", text)
    text = re.sub(r"\b(hsa|mmu)(?=\d)", " ", text)
    text = re.sub(r"[^a-z0-9]+", " ", text)
    return " ".join(text.split())


def normalize_chrom(value: Any) -> str:
    chrom = normalize_text(value)
    if chrom.startswith("chr"):
        chrom = chrom[3:]
    return {"m": "mt", "mitochondria": "mt"}.get(chrom, chrom)


def maybe_float(value: Any) -> float | None:
    try:
        return float(str(value).strip())
    except Exception:
        return None


def get_row_value(row: dict, column: str, default: str = ""):
    if column in row:
        return row.get(column, default)
    target = normalize_text(column)
    for key, value in row.items():
        if normalize_text(key) == target:
            return value
    return default


def token_similarity(left: str, right: str) -> float:
    left_tokens = set(left.split())
    right_tokens = set(right.split())
    if not left_tokens and not right_tokens:
        return 1.0
    if not left_tokens or not right_tokens:
        return 0.0
    return len(left_tokens & right_tokens) / len(left_tokens | right_tokens)


def text_similarity(left: Any, right: Any) -> float:
    left_norm = normalize_for_soft_match(left)
    right_norm = normalize_for_soft_match(right)
    if left_norm == right_norm:
        return 1.0
    if not left_norm or not right_norm:
        return 0.0
    sequence_score = SequenceMatcher(None, left_norm, right_norm).ratio()
    token_score = token_similarity(left_norm, right_norm)
    return max(sequence_score, token_score)


def numeric_similarity(left: Any, right: Any, rel_scale: float = 0.1) -> float | None:
    left_value = maybe_float(left)
    right_value = maybe_float(right)
    if left_value is None or right_value is None:
        return None
    if left_value == right_value:
        return 1.0
    scale = max(abs(left_value), abs(right_value), 1e-12)
    relative_error = abs(left_value - right_value) / scale
    return max(0.0, 1.0 - (relative_error / rel_scale))


def load_task_metadata(metadata_path: Path) -> dict[str, dict]:
    if not metadata_path.exists():
        return {}
    payload = json.loads(metadata_path.read_text(encoding="utf-8"))
    return {item["task_id"]: item for item in payload}


def load_run_metadata(run_dir: Path) -> dict:
    path = run_dir / "run_metadata.json"
    if not path.exists():
        return {}
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return {}


def load_csv_rows(path: Path) -> list[dict]:
    with path.open("r", encoding="utf-8", errors="ignore", newline="") as handle:
        return list(csv.DictReader(handle))


def load_tsv_no_header(path: Path) -> list[dict]:
    rows = []
    with path.open("r", encoding="utf-8", errors="ignore") as handle:
        for line in handle:
            line = line.rstrip("\n")
            if not line:
                continue
            parts = line.split("\t")
            if len(parts) >= 2:
                rows.append({"transcript_id": parts[0], "count": parts[1]})
    return rows


def read_text_preview(path: Path, max_lines: int = 80, max_chars: int = 24000) -> str:
    opener = gzip.open if path.suffix == ".gz" else open
    preview_lines = []
    total_chars = 0
    with opener(path, "rt", encoding="utf-8", errors="ignore") as handle:
        for i, line in enumerate(handle):
            if i >= max_lines:
                preview_lines.append("... [truncated]")
                break
            preview_lines.append(line.rstrip("\n"))
            total_chars += len(line)
            if total_chars >= max_chars:
                preview_lines.append("... [truncated]")
                break
    return "\n".join(preview_lines)


def load_prediction_rows(task_id: str, file_name: str, prediction_path: Path) -> list[dict]:
    if TASK_CONFIGS[task_id].get("tsv_no_header"):
        return load_tsv_no_header(prediction_path)
    return load_csv_rows(prediction_path)


def open_text_auto(path: Path):
    return gzip.open(path, "rt", encoding="utf-8", errors="ignore") if path.suffix == ".gz" else path.open(
        "r", encoding="utf-8", errors="ignore"
    )


def parse_bed_intervals(path: Path) -> dict[str, list[tuple[int, int]]]:
    intervals: dict[str, list[tuple[int, int]]] = {}
    if not path.exists():
        return intervals
    with open_text_auto(path) as handle:
        for line in handle:
            if not line.strip() or line.startswith(("#", "track", "browser")):
                continue
            fields = line.rstrip("\n").split("\t")
            if len(fields) < 3:
                continue
            try:
                start = int(fields[1])
                end = int(fields[2])
            except ValueError:
                continue
            chrom = normalize_chrom(fields[0])
            intervals.setdefault(chrom, []).append((start, end))
    for chrom in intervals:
        intervals[chrom].sort()
    return intervals


def position_in_intervals(chrom: str, pos: str, intervals: dict[str, list[tuple[int, int]]] | None) -> bool:
    if not intervals:
        return True
    try:
        pos_1_based = int(pos)
    except ValueError:
        return False
    pos_0_based = pos_1_based - 1
    chrom_intervals = intervals.get(normalize_chrom(chrom), [])
    interval_index = bisect_right(chrom_intervals, (pos_0_based, 10**18)) - 1
    if interval_index < 0:
        return False
    start, end = chrom_intervals[interval_index]
    return start <= pos_0_based < end


def load_vcf_keys(path: Path, intervals: dict[str, list[tuple[int, int]]] | None = None) -> set[tuple[str, str, str, str]]:
    keys = set()
    with open_text_auto(path) as handle:
        for line in handle:
            if not line or line.startswith("#"):
                continue
            fields = line.rstrip("\n").split("\t")
            if len(fields) < 5:
                continue
            chrom, pos, _vid, ref, alt = fields[:5]
            if not position_in_intervals(chrom, pos, intervals):
                continue
            for alt_item in alt.split(","):
                keys.add((normalize_chrom(chrom), normalize_text(pos), normalize_text(ref), normalize_text(alt_item)))
    return keys


def f1_from_counts(tp: int, pred_count: int, truth_count: int) -> dict[str, float]:
    precision = tp / pred_count if pred_count else 0.0
    recall = tp / truth_count if truth_count else 0.0
    f1 = 0.0 if precision + recall == 0 else (2 * precision * recall) / (precision + recall)
    return {"precision": precision, "recall": recall, "f1": f1}


def compare_variant_sets(pred_keys: set[tuple[str, str, str, str]], truth_keys: set[tuple[str, str, str, str]]) -> dict:
    tp = len(pred_keys & truth_keys)
    metrics = f1_from_counts(tp, len(pred_keys), len(truth_keys))
    return {
        "pred_variant_count": len(pred_keys),
        "truth_variant_count": len(truth_keys),
        "shared_variant_count": tp,
        "precision": metrics["precision"],
        "recall": metrics["recall"],
        "f1": metrics["f1"],
    }


def compare_vcf(prediction_path: Path, truth_path: Path, task_id: str, dataset_root: Path) -> dict:
    bed_candidates = [
        dataset_root / task_id / "data" / "Agilent_v7.chr.bed",
        dataset_root / task_id / "results" / "HG001_GRCh38_1_22_v4.2.1_benchmark.bed",
    ]
    bed_path = next((path for path in bed_candidates if path.exists()), None)
    if not bed_path:
        unfiltered_pred_keys = load_vcf_keys(prediction_path)
        unfiltered_truth_keys = load_vcf_keys(truth_path)
        unfiltered = compare_variant_sets(unfiltered_pred_keys, unfiltered_truth_keys)
        return {
            **unfiltered,
            "truth_scope": "all_truth_variants",
            "unfiltered_f1": unfiltered["f1"],
            "note": "Approximate VCF comparison by normalized CHROM,POS,REF,ALT; no BED scope was available.",
        }

    intervals = parse_bed_intervals(bed_path)
    scoped_pred_keys = load_vcf_keys(prediction_path, intervals=intervals)
    scoped_truth_keys = load_vcf_keys(truth_path, intervals=intervals)
    scoped = compare_variant_sets(scoped_pred_keys, scoped_truth_keys)
    return {
        **scoped,
        "truth_scope": "target_bed",
        "target_bed": str(bed_path),
        "unfiltered_f1": None,
        "note": "Primary F1 is restricted to the task target BED and uses normalized CHROM,POS,REF,ALT.",
    }


def row_similarity(
    pred: dict,
    truth: dict,
    soft_key_columns: list[str],
    numeric_columns: list[str],
    text_weight: float,
    numeric_weight: float,
) -> float:
    text_scores = [
        text_similarity(get_row_value(pred, column, ""), get_row_value(truth, column, ""))
        for column in soft_key_columns
    ]
    text_score = sum(text_scores) / len(text_scores) if text_scores else 0.0

    numeric_scores = [
        score
        for column in numeric_columns
        if (score := numeric_similarity(get_row_value(pred, column, ""), get_row_value(truth, column, ""))) is not None
    ]
    if not numeric_scores or numeric_weight <= 0:
        return text_score
    numeric_score = sum(numeric_scores) / len(numeric_scores)
    total_weight = text_weight + numeric_weight
    return ((text_score * text_weight) + (numeric_score * numeric_weight)) / total_weight


def greedy_soft_row_match(
    pred_rows: list[dict],
    truth_rows: list[dict],
    soft_key_columns: list[str],
    numeric_columns: list[str],
    threshold: float,
    text_weight: float,
    numeric_weight: float,
) -> tuple[list[tuple[int, int, float]], list[int], list[int]]:
    candidates = []
    first_soft_column = soft_key_columns[0] if soft_key_columns else None
    truth_buckets: dict[str, list[tuple[int, dict]]] = {}
    use_bucketed_candidates = first_soft_column is not None and len(pred_rows) * len(truth_rows) > 200_000
    if use_bucketed_candidates:
        for truth_index, truth in enumerate(truth_rows):
            bucket_key = normalize_for_soft_match(get_row_value(truth, first_soft_column, ""))
            truth_buckets.setdefault(bucket_key, []).append((truth_index, truth))

    for pred_index, pred in enumerate(pred_rows):
        if use_bucketed_candidates:
            bucket_key = normalize_for_soft_match(get_row_value(pred, first_soft_column, ""))
            candidate_truth_rows = truth_buckets.get(bucket_key, [])
        else:
            candidate_truth_rows = list(enumerate(truth_rows))
        for truth_index, truth in candidate_truth_rows:
            score = row_similarity(pred, truth, soft_key_columns, numeric_columns, text_weight, numeric_weight)
            if score >= threshold:
                candidates.append((score, pred_index, truth_index))

    candidates.sort(reverse=True)
    used_pred = set()
    used_truth = set()
    matches = []
    for score, pred_index, truth_index in candidates:
        if pred_index in used_pred or truth_index in used_truth:
            continue
        used_pred.add(pred_index)
        used_truth.add(truth_index)
        matches.append((pred_index, truth_index, score))
    unmatched_pred = [index for index in range(len(pred_rows)) if index not in used_pred]
    unmatched_truth = [index for index in range(len(truth_rows)) if index not in used_truth]
    return matches, unmatched_pred, unmatched_truth


def compare_table_rows(task_id: str, file_name: str, pred_rows: list[dict], truth_rows: list[dict]) -> dict:
    config = TASK_CONFIGS[task_id]
    key_columns = config["key_columns"][file_name] if isinstance(config["key_columns"], dict) else config["key_columns"]
    numeric_columns = (
        config["numeric_columns"][file_name] if isinstance(config["numeric_columns"], dict) else config["numeric_columns"]
    )

    pred_map = {
        tuple(normalize_text(get_row_value(row, col, "")) for col in key_columns): row
        for row in pred_rows
    }
    truth_map = {
        tuple(normalize_text(get_row_value(row, col, "")) for col in key_columns): row
        for row in truth_rows
    }

    pred_keys = set(pred_map)
    truth_keys = set(truth_map)
    shared_keys = pred_keys & truth_keys
    metrics = f1_from_counts(len(shared_keys), len(pred_keys), len(truth_keys))

    numeric_diffs = {col: [] for col in numeric_columns}
    for key in shared_keys:
        pred = pred_map[key]
        truth = truth_map[key]
        for col in numeric_columns:
            pred_value = maybe_float(get_row_value(pred, col, ""))
            truth_value = maybe_float(get_row_value(truth, col, ""))
            if pred_value is not None and truth_value is not None:
                numeric_diffs[col].append(abs(pred_value - truth_value))

    summary = {
        "pred_row_count": len(pred_rows),
        "truth_row_count": len(truth_rows),
        "shared_key_count": len(shared_keys),
        "exact_key_precision": metrics["precision"],
        "exact_key_recall": metrics["recall"],
        "exact_key_f1": metrics["f1"],
        "key_precision": metrics["precision"],
        "key_recall": metrics["recall"],
        "key_f1": metrics["f1"],
        "match_strategy": "exact_key",
        "match_precision": metrics["precision"],
        "match_recall": metrics["recall"],
        "match_f1": metrics["f1"],
    }

    flexible_config = FLEXIBLE_TABLE_MATCH_CONFIGS.get((task_id, file_name))
    if flexible_config:
        soft_key_columns = flexible_config.get("soft_key_columns", key_columns)
        threshold = flexible_config.get("threshold", 0.75)
        text_weight = flexible_config.get("text_weight", 1.0)
        numeric_weight = flexible_config.get("numeric_weight", 0.0)
        matches, unmatched_pred, unmatched_truth = greedy_soft_row_match(
            pred_rows,
            truth_rows,
            soft_key_columns,
            numeric_columns,
            threshold,
            text_weight,
            numeric_weight,
        )
        soft_metrics = f1_from_counts(len(matches), len(pred_rows), len(truth_rows))
        summary.update(
            {
                "match_strategy": "soft_row_similarity",
                "soft_key_columns": soft_key_columns,
                "match_threshold": threshold,
                "soft_match_count": len(matches),
                "match_precision": soft_metrics["precision"],
                "match_recall": soft_metrics["recall"],
                "match_f1": soft_metrics["f1"],
                "mean_match_score": mean([score for *_unused, score in matches]) if matches else 0.0,
                "unmatched_prediction_examples": [
                    {column: get_row_value(pred_rows[index], column, "") for column in soft_key_columns}
                    for index in unmatched_pred[:5]
                ],
                "unmatched_truth_examples": [
                    {column: get_row_value(truth_rows[index], column, "") for column in soft_key_columns}
                    for index in unmatched_truth[:5]
                ],
            }
        )
    if numeric_columns:
        summary["numeric_mae"] = {
            col: (sum(values) / len(values) if values else None) for col, values in numeric_diffs.items()
        }
    return summary


def summarize_file_for_judge(task_id: str, path: Path) -> dict:
    config = TASK_CONFIGS[task_id]
    if not path.exists():
        return {"path": str(path), "exists": False}
    if config.get("vcf"):
        return {
            "path": str(path),
            "exists": True,
            "type": "vcf.gz",
            "size_bytes": path.stat().st_size,
            "preview": read_text_preview(path, max_lines=60, max_chars=20000),
        }
    if path.suffix.lower() == ".tsv" and config.get("tsv_no_header"):
        rows = load_tsv_no_header(path)
        return {
            "path": str(path),
            "exists": True,
            "type": "tsv",
            "row_count": len(rows),
            "columns": ["transcript_id", "count"],
            "preview": read_text_preview(path, max_lines=80, max_chars=22000),
        }
    rows = load_csv_rows(path)
    return {
        "path": str(path),
        "exists": True,
        "type": path.suffix.lower().lstrip(".") or "text",
        "row_count": len(rows),
        "columns": list(rows[0].keys()) if rows else [],
        "preview": read_text_preview(path, max_lines=80, max_chars=22000),
    }


def locate_prediction_file(task_id: str, run_dir: Path, file_name: str) -> Path | None:
    candidate = run_dir / file_name
    if candidate.exists():
        return candidate
    matches = sorted(run_dir.rglob(file_name))
    return matches[0] if matches else None


def infer_latest_run_dir(task_id: str, runs_root: Path) -> Path | None:
    candidates = sorted(path for path in runs_root.glob(f"{task_id}_*") if path.is_dir())
    return candidates[-1] if candidates else None


def build_processing_tree(run_dir: Path, max_entries: int = 600) -> list[str]:
    entries = []
    for path in sorted(run_dir.rglob("*")):
        rel = path.relative_to(run_dir)
        if len(entries) >= max_entries:
            entries.append("... [truncated]")
            break
        if path.is_dir():
            entries.append(f"{rel}/")
        else:
            entries.append(f"{rel}\t{path.stat().st_size} bytes")
    return entries


def collect_trace_path_evidence(run_dir: Path) -> dict:
    """Collect only folders/file paths from the run, matching the paper's trace input."""
    execution_log = run_dir / "execution_log.txt"
    path_mentions = []
    if execution_log.exists():
        text = execution_log.read_text(encoding="utf-8", errors="ignore")
        for token in text.replace('"', " ").replace("'", " ").split():
            if token.startswith("/") and ("/bioagent-bench/" in token or "/bioagent-bench-runs/" in token):
                cleaned = token.rstrip("),.;:<>")
                if cleaned not in path_mentions:
                    path_mentions.append(cleaned)
            if len(path_mentions) >= 250:
                break
    return {
        "processing_tree": build_processing_tree(run_dir),
        "path_mentions_from_trace": path_mentions,
    }


def build_artifact_metrics(task_id: str, run_dir: Path, dataset_root: Path) -> list[dict]:
    config = TASK_CONFIGS[task_id]
    truth_dir = dataset_root / task_id / "results"
    artifacts = []
    for truth_name, result_name in zip(config["truth_files"], config["result_files"], strict=False):
        truth_path = truth_dir / truth_name
        prediction_path = locate_prediction_file(task_id, run_dir, result_name)
        entry = {
            "truth_file": str(truth_path),
            "prediction_file": str(prediction_path) if prediction_path else None,
            "prediction_exists": bool(prediction_path and prediction_path.exists()),
        }
        if not prediction_path or not prediction_path.exists() or not truth_path.exists():
            entry["metrics"] = {"error": "missing prediction or truth file"}
            artifacts.append(entry)
            continue
        if config.get("vcf"):
            entry["metrics"] = compare_vcf(prediction_path, truth_path, task_id, dataset_root)
        else:
            pred_rows = load_prediction_rows(task_id, truth_name, prediction_path)
            truth_rows = load_prediction_rows(task_id, truth_name, truth_path)
            entry["metrics"] = compare_table_rows(task_id, truth_name, pred_rows, truth_rows)
        artifacts.append(entry)
    return artifacts


def result_artifact_summaries(task_id: str, run_dir: Path, dataset_root: Path) -> tuple[list[dict], list[dict]]:
    config = TASK_CONFIGS[task_id]
    truth_dir = dataset_root / task_id / "results"
    result_summaries = []
    truth_summaries = []
    for truth_name, result_name in zip(config["truth_files"], config["result_files"], strict=False):
        prediction_path = locate_prediction_file(task_id, run_dir, result_name)
        truth_path = truth_dir / truth_name
        result_summaries.append(
            summarize_file_for_judge(task_id, prediction_path) if prediction_path else {"path": result_name, "exists": False}
        )
        truth_summaries.append(summarize_file_for_judge(task_id, truth_path))
    return result_summaries, truth_summaries


def infer_rule_steps(task_id: str, artifacts: list[dict], trace_evidence: dict) -> tuple[int, int, list[str]]:
    config = TASK_CONFIGS[task_id]
    expected_steps = config["pipeline_steps"]
    total = len(expected_steps)
    haystack = "\n".join(trace_evidence["processing_tree"] + trace_evidence["path_mentions_from_trace"]).lower()
    completed = 0
    evidence = []
    keyword_sets = {
        "inspect": ["task_query", "run_metadata", "data"],
        "metadata": ["metadata", "counts"],
        "differential": ["deseq", "de_results", "differential", "log2fold"],
        "enrichment": ["kegg", "enrichment", "pathway"],
        "protein": ["prodigal", ".faa", "protein"],
        "ortholog": ["orthofinder", "mmseqs", "diamond", "cluster"],
        "variant": [".vcf", "bcftools", "variant"],
        "align": [".bam", ".sam", "bwa", "star", "aligned"],
        "count": ["feature_counts", "counts", "abundance"],
        "classify": ["kraken", "kaiju", "classification", "report"],
        "assemble": ["spades", "megahit", "contigs", "assembly"],
        "single-cell": ["matrix", "scanpy", "cluster", "umap", "markers"],
        "index": ["index", ".idx", "star_index"],
        "final": [artifact["prediction_file"] or "" for artifact in artifacts],
    }
    for step in expected_steps:
        step_l = step.lower()
        if (
            ("write" in step_l or "final" in step_l)
            and any(name in step_l for name in ["csv", "tsv", "vcf", "predicted"])
            and not any(artifact.get("prediction_exists") for artifact in artifacts)
        ):
            continue
        if (
            ("write" in step_l or "final" in step_l)
            and any(name in step_l for name in ["csv", "tsv", "vcf", "predicted"])
            and all(artifact.get("prediction_exists") for artifact in artifacts)
        ):
            completed += 1
            evidence.append(step)
            continue
        keys = []
        for concept, words in keyword_sets.items():
            if concept in step_l or any(word in step_l for word in words[:2]):
                keys.extend(words)
        if not keys:
            generic = {"and", "or", "the", "to", "a", "an", "write", "final", "requested"}
            keys = [token for token in step_l.replace(",", " ").split() if len(token) > 3 and token not in generic][:4]
        if any(key and key.lower() in haystack for key in keys):
            completed += 1
            evidence.append(step)

    # Final artifacts are the strongest evidence for final-result step.
    final_files = [artifact for artifact in artifacts if artifact.get("prediction_exists")]
    if final_files and completed < total:
        completed = max(completed, total - 1)
        evidence.append("final artifact(s) exist; inferred most upstream steps completed")
    return min(completed, total), total, evidence


def rule_trial_judge(task_id: str, artifacts: list[dict], trace_evidence: dict) -> dict:
    steps_completed, steps_to_completion, evidence = infer_rule_steps(task_id, artifacts, trace_evidence)
    final_result_reached = all(artifact.get("prediction_exists") for artifact in artifacts)
    metric_scores = []
    giab_f1 = None
    for artifact in artifacts:
        metrics = artifact.get("metrics", {})
        if "match_f1" in metrics:
            metric_scores.append(metrics["match_f1"])
        elif "f1" in metrics:
            metric_scores.append(metrics["f1"])
            if task_id == "giab":
                giab_f1 = metrics["f1"]
        elif "key_f1" in metrics:
            metric_scores.append(metrics["key_f1"])
    mean_metric_score = mean(metric_scores) if metric_scores else 0.0

    # Paper-style correctness flag is task/rubric specific. In rule mode, use a conservative helper.
    if TASK_CONFIGS[task_id].get("verifiable"):
        threshold = 0.5 if task_id == "giab" else 0.8
        results_match = final_result_reached and mean_metric_score >= threshold
    else:
        results_match = final_result_reached and steps_completed == steps_to_completion

    notes = (
        "Rule-mode approximation of the paper's LLM grader. "
        f"Evidence-backed steps: {', '.join(evidence[:8]) if evidence else 'none detected'}."
    )
    return {
        "steps_completed": steps_completed,
        "steps_to_completion": steps_to_completion,
        "final_result_reached": final_result_reached,
        "notes": notes,
        "results_match": results_match,
        "f1_score": giab_f1,
    }


def resolve_judge_client(provider: str, model: str | None, base_url: str | None, api_key: str | None):
    try:
        from openai import OpenAI
    except ImportError as exc:
        raise ImportError("openai package is required for LLM judging.") from exc

    provider = provider.lower().strip()
    if provider == "deepseek":
        resolved_model = model or os.getenv("DEEPSEEK_MODEL_NAME", "deepseek-chat")
        resolved_base_url = base_url or os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1")
        resolved_api_key = api_key or os.getenv("DEEPSEEK_API_KEY")
        if not resolved_api_key:
            raise RuntimeError("Missing DEEPSEEK_API_KEY for DeepSeek LLM judge.")
        return OpenAI(api_key=resolved_api_key, base_url=resolved_base_url), resolved_model, provider
    if provider == "openai":
        resolved_model = model or os.getenv("OPENAI_MODEL_NAME", "gpt-5.1")
        resolved_api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not resolved_api_key:
            raise RuntimeError("Missing OPENAI_API_KEY for OpenAI LLM judge.")
        return OpenAI(api_key=resolved_api_key), resolved_model, provider
    raise ValueError(f"Unsupported LLM judge provider: {provider}")


def extract_json_object(text: str) -> dict:
    text = text.strip()
    if not text:
        raise ValueError("Empty LLM judge response.")
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        start = text.find("{")
        end = text.rfind("}")
        if start != -1 and end != -1 and end > start:
            return json.loads(text[start : end + 1])
        raise


def llm_trial_judge(
    task_id: str,
    task_prompt: str,
    input_data_path: Path,
    reference_data_path: Path,
    result_summaries: list[dict],
    truth_summaries: list[dict],
    trace_evidence: dict,
    artifacts: list[dict],
    provider: str,
    model: str | None,
    base_url: str | None,
    api_key: str | None,
) -> dict:
    client, resolved_model, resolved_provider = resolve_judge_client(provider, model, base_url, api_key)
    config = TASK_CONFIGS[task_id]
    compact_tree = "\n".join(trace_evidence["processing_tree"][:450])
    path_mentions = "\n".join(trace_evidence["path_mentions_from_trace"][:180])

    user_prompt = f"""
You are a strict, impartial Bioinformatics Pipeline Judge. Evaluate an LLM agent's work for executing a bioinformatics pipeline instructed by the prompt.

Inputs:
1. Input data: {input_data_path}
2. Reference data: {reference_data_path if reference_data_path.exists() else "<none>"}
3. Processing tree:
{compact_tree}

Trace path mentions:
{path_mentions or "<none>"}

4. Results:
{json.dumps(result_summaries, ensure_ascii=False, indent=2)}

5. Truth:
{json.dumps(truth_summaries, ensure_ascii=False, indent=2)}

6. Prompt:
{task_prompt}

Task expected pipeline steps:
{json.dumps(config["pipeline_steps"], ensure_ascii=False, indent=2)}

Auxiliary rule-based artifact metrics:
{json.dumps(artifacts, ensure_ascii=False, indent=2)}

Evaluation rules:
- Prioritize evaluation of pipeline completion over exact numerical correctness.
- If gene names are of different naming conventions, the result can still be valid.
- Estimate steps_to_completion from bioinformatics-relevant steps required for this task.
- Count upstream steps only if their expected artifacts are present.
- Do not count placeholders or mock completion as completed steps.
- For p-values, log-fold values, abundances, or other statistics, check that plausible prior steps generated them.
- Task-specific results-match guidance: {config.get("results_match_guidance", "")}

Return JSON only with exactly these fields:
- steps_completed: integer
- steps_to_completion: integer
- final_result_reached: boolean
- notes: string
- results_match: boolean
- f1_score: number or null; only use a real F1 for GIAB/variant concordance, otherwise null
""".strip()

    response = client.chat.completions.create(
        model=resolved_model,
        temperature=0,
        response_format={"type": "json_object"},
        messages=[
            {"role": "system", "content": "Return strict JSON only. Follow the BioAgent Bench EvaluationResults schema."},
            {"role": "user", "content": user_prompt},
        ],
    )
    parsed = extract_json_object(response.choices[0].message.content or "{}")
    parsed["provider"] = resolved_provider
    parsed["model"] = resolved_model
    return parsed


def normalize_evaluation_results(raw: dict) -> dict:
    steps_completed = int(raw.get("steps_completed", 0) or 0)
    steps_to_completion = int(raw.get("steps_to_completion", 0) or 0)
    final_result_reached = bool(raw.get("final_result_reached", False))
    results_match = bool(raw.get("results_match", False))
    f1_score = raw.get("f1_score")
    if f1_score is not None:
        try:
            f1_score = float(f1_score)
        except Exception:
            f1_score = None
    completion_rate = steps_completed / steps_to_completion if steps_to_completion else 0.0
    return {
        "steps_completed": steps_completed,
        "steps_to_completion": steps_to_completion,
        "completion_rate": completion_rate,
        "final_result_reached": final_result_reached,
        "results_match": results_match,
        "f1_score": f1_score,
        "notes": str(raw.get("notes", "")),
    }


def evaluate_task(
    task_id: str,
    run_dir: Path,
    dataset_root: Path,
    task_metadata: dict[str, dict],
    judge_mode: str = "rule",
    llm_provider: str = "deepseek",
    llm_model: str | None = None,
    llm_base_url: str | None = None,
    llm_api_key: str | None = None,
) -> dict:
    task_dir = dataset_root / task_id
    run_metadata = load_run_metadata(run_dir)
    task_prompt = (
        run_metadata.get("benchmark_task_context", {}).get("task_prompt")
        or task_metadata.get(task_id, {}).get("task_prompt")
        or run_metadata.get("query")
        or ""
    )
    trace_evidence = collect_trace_path_evidence(run_dir)
    artifacts = build_artifact_metrics(task_id, run_dir, dataset_root)
    result_summaries, truth_summaries = result_artifact_summaries(task_id, run_dir, dataset_root)

    rule_raw = rule_trial_judge(task_id, artifacts, trace_evidence)
    rule_result = normalize_evaluation_results(rule_raw)

    selected_raw = rule_raw
    llm_result = None
    if judge_mode in {"llm", "both"}:
        selected_raw = llm_trial_judge(
            task_id=task_id,
            task_prompt=task_prompt,
            input_data_path=task_dir / "data",
            reference_data_path=task_dir / "reference",
            result_summaries=result_summaries,
            truth_summaries=truth_summaries,
            trace_evidence=trace_evidence,
            artifacts=artifacts,
            provider=llm_provider,
            model=llm_model,
            base_url=llm_base_url,
            api_key=llm_api_key,
        )
        llm_result = normalize_evaluation_results(selected_raw)

    selected = rule_result if judge_mode == "rule" else llm_result
    assert selected is not None
    return {
        "task_id": task_id,
        "run_dir": str(run_dir),
        "evaluated_at_utc": utc_timestamp(),
        "judge_mode": judge_mode,
        "evaluation_results": selected,
        "overall_score": selected["completion_rate"],
        "score_definition": (
            "BioAgent Bench-style completion rate: steps_completed / steps_to_completion. "
            "results_match and f1_score are reported separately."
        ),
        "rule_evaluation_results": rule_result,
        "llm_evaluation_results": llm_result,
        "artifacts": artifacts,
        "result_summaries": result_summaries,
        "truth_summaries": truth_summaries,
        "trace_evidence": trace_evidence,
        "paper_alignment": {
            "grader_inputs": [
                "input data path",
                "reference data path",
                "expected outcome/truth as text summary",
                "agent outcome as text summary",
                "agent trace represented as folders/file paths",
                "task prompt and grading logic",
            ],
            "grader_outputs": [
                "steps_completed",
                "steps_to_completion",
                "final_result_reached",
                "notes",
                "results_match",
                "f1_score",
            ],
        },
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Evaluate Hypo_Bio_OS outputs on bioagent-bench.")
    parser.add_argument("--task", action="append", help="Task ID to evaluate. Can be provided multiple times.")
    parser.add_argument("--all", action="store_true", help="Evaluate all tasks that have run directories.")
    parser.add_argument("--run-dir", help="Specific run directory for single-task evaluation.")
    parser.add_argument("--runs-root", default=str(DEFAULT_RUNS_ROOT))
    parser.add_argument("--dataset-root", default=str(DATASET_ROOT))
    parser.add_argument("--metadata", default=str(METADATA_PATH))
    parser.add_argument("--output", default=None, help="Where to save the evaluation JSON.")
    parser.add_argument(
        "--judge-mode",
        choices=["rule", "llm", "both"],
        default="rule",
        help="Use local paper-shaped heuristic grading, LLM grading, or both.",
    )
    parser.add_argument(
        "--llm-provider",
        choices=["openai", "deepseek"],
        default=os.getenv("BIOAGENT_BENCH_JUDGE_PROVIDER", "deepseek"),
        help="OpenAI-compatible backend provider for the LLM judge.",
    )
    parser.add_argument("--llm-model", default=None, help="Override model name for the LLM judge.")
    parser.add_argument("--llm-base-url", default=None, help="Override base URL for the LLM judge.")
    parser.add_argument("--llm-api-key", default=None, help="Override API key for the LLM judge.")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    dataset_root = Path(args.dataset_root)
    runs_root = Path(args.runs_root)
    task_metadata = load_task_metadata(Path(args.metadata))

    if args.all:
        task_ids = list(TASK_CONFIGS)
    else:
        task_ids = args.task or []
    if not task_ids:
        raise SystemExit("Provide --task <task_id> or use --all.")

    results = []
    for task_id in task_ids:
        if task_id not in TASK_CONFIGS:
            raise SystemExit(f"Unsupported task ID: {task_id}")
        if args.run_dir and len(task_ids) == 1:
            run_dir = Path(args.run_dir)
        else:
            run_dir = infer_latest_run_dir(task_id, runs_root)
            if run_dir is None:
                print(f"Skipping {task_id}: no run directory found under {runs_root}")
                continue
        result = evaluate_task(
            task_id=task_id,
            run_dir=run_dir,
            dataset_root=dataset_root,
            task_metadata=task_metadata,
            judge_mode=args.judge_mode,
            llm_provider=args.llm_provider,
            llm_model=args.llm_model,
            llm_base_url=args.llm_base_url,
            llm_api_key=args.llm_api_key,
        )
        results.append(result)
        print(json.dumps(result, ensure_ascii=False, indent=2))

    completion_rates = [item["evaluation_results"]["completion_rate"] for item in results]
    payload = {
        "evaluated_at_utc": utc_timestamp(),
        "judge_mode": args.judge_mode,
        "primary_metric": "completion_rate",
        "mean_completion_rate": mean(completion_rates) if completion_rates else 0.0,
        "results": results,
    }
    if args.output:
        output_path = Path(args.output)
    else:
        task_eval_dir = runs_root  / "evaluation_results"
        task_eval_dir.mkdir(parents=True, exist_ok=True)
        output_path = task_eval_dir / f"evaluation_{task_id}.json"
    output_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Saved evaluation summary to: {output_path}")
    return 0


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