scripts/build_l3_development_bundle.py

#!/usr/bin/env python3
"""Build a hash-linked L3 train/dev bundle from authorized gold receipts."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import sys
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import pandas as pd

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from telen.experiment import (  # noqa: E402
    PILOT_SEEDS,
    load_l3_training_bundle,
)
from telen.learned_evaluation import DevelopmentEvaluationConfig  # noqa: E402


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def portable_path(path: Path, base: Path) -> str:
    try:
        return str(path.resolve().relative_to(base.resolve()))
    except ValueError:
        return str(path.resolve())


def load_receipt(path: Path, artifact_type: str) -> dict:
    value = json.loads(path.read_text(encoding="utf-8"))
    if value.get("artifact_type") != artifact_type:
        raise ValueError(f"{path} has the wrong artifact_type")
    return value


def question_strata(questions: pd.DataFrame) -> dict[str, str]:
    mapping: dict[str, str] = {}
    for row in questions.to_dict(orient="records"):
        if row.get("annotation_status") != "adjudicated_gold":
            continue
        suffixes = ["affected-1", "affected-2"]
        if str(row.get("affected_question_3", "")).strip():
            suffixes.append("affected-3")
        for suffix in suffixes:
            mapping[f"{row['query_group_id']}:{suffix}"] = "affected"
        mapping[f"{row['query_group_id']}:unaffected-control"] = "control"
    return mapping


def parse_ids(value: str, label: str) -> list[str]:
    try:
        parsed = json.loads(value)
    except json.JSONDecodeError as exc:
        raise ValueError(f"{label} must be a JSON array") from exc
    if (
        not isinstance(parsed, list)
        or not parsed
        or any(not isinstance(item, str) or not item for item in parsed)
        or len(parsed) != len(set(parsed))
    ):
        raise ValueError(f"{label} must contain unique non-empty string ids")
    return parsed


def build(args: argparse.Namespace) -> dict[str, object]:
    annotation = load_receipt(args.annotation_audit, "expert_annotation_gate_audit")
    split = load_receipt(args.split_audit, "telen_l3_event_split_audit")
    baseline = load_receipt(args.baseline_gate, "telen_gold_baseline_gate")
    if annotation.get("learned_residual_gate", {}).get("authorized") is not True:
        raise ValueError("expert annotation gate is not authorized")
    if split.get("authorized") is not True:
        raise ValueError("event split gate is not authorized")
    if baseline.get("authorized") is not True:
        raise ValueError("gold baseline/oracle gate is not authorized")
    annotation_hash = sha256_file(args.annotation_audit)
    split_hash = sha256_file(args.split_audit)
    if split.get("inputs", {}).get("annotation_audit_sha256") != annotation_hash:
        raise ValueError("split receipt belongs to another annotation audit")
    if (
        baseline.get("inputs", {}).get("annotation_audit_sha256") != annotation_hash
        or baseline.get("inputs", {}).get("split_audit_sha256") != split_hash
    ):
        raise ValueError("baseline receipt is not linked to annotation/split receipts")

    events_annotation = pd.read_csv(args.annotation_events, dtype=str).fillna("")
    questions = pd.read_csv(args.questions, dtype=str).fillna("")
    qrels = pd.read_csv(args.qrels, dtype=str).fillna("")
    base_scores = pd.read_csv(args.base_scores, dtype=str).fillna("")
    assignments = pd.read_csv(args.split_assignments, dtype=str).fillna("")
    legal_events = pd.read_csv(args.legal_events, dtype=str).fillna("")
    interventions = pd.read_csv(args.interventions, dtype=str).fillna("")
    annotation_inputs = annotation.get("inputs", {})
    expected_annotation_hashes = {
        "events_sha256": sha256_file(args.annotation_events),
        "questions_sha256": sha256_file(args.questions),
        "qrels_sha256": sha256_file(args.qrels),
    }
    for name, observed in expected_annotation_hashes.items():
        if annotation_inputs.get(name) != observed:
            raise ValueError(f"annotation receipt does not authorize {name}")
    if split.get("inputs", {}).get("assignments_sha256") != sha256_file(
        args.split_assignments
    ):
        raise ValueError("split assignment file differs from its receipt")
    baseline_inputs = baseline.get("inputs", {})
    if (
        baseline_inputs.get("qrels_sha256") != sha256_file(args.qrels)
        or baseline_inputs.get("questions_sha256") != sha256_file(args.questions)
        or baseline_inputs.get("scores_sha256") != sha256_file(args.base_scores)
    ):
        raise ValueError("baseline gate input hashes differ from bundle inputs")

    required_qrels = {
        "event_group_id",
        "query_group_id",
        "query_id",
        "snapshot_phase",
        "snapshot_date",
        "candidate_version_id",
        "relevance_grade_adjudicated",
        "judgment_status",
    }
    required_scores = {
        "query_id",
        "snapshot_date",
        "system",
        "candidate_version_id",
        "score",
    }
    required_interventions = {
        "query_instance_id",
        "event_group_id",
        "query_id",
        "snapshot_date",
        "stale_as_of",
        "causal_event_ids",
        "wrong_event_ids",
        "unrelated_event_ids",
    }
    for frame, required, label in (
        (qrels, required_qrels, "qrels"),
        (base_scores, required_scores, "base scores"),
        (interventions, required_interventions, "interventions"),
    ):
        missing = sorted(required - set(frame.columns))
        if missing:
            raise ValueError(f"{label} are missing columns: {missing}")
    if not qrels["judgment_status"].eq("adjudicated").all():
        raise ValueError("all qrels must be adjudicated")
    qrels["grade"] = pd.to_numeric(
        qrels["relevance_grade_adjudicated"], errors="coerce"
    )
    base_scores["score_numeric"] = pd.to_numeric(base_scores["score"], errors="coerce")
    if qrels["grade"].isna().any() or not qrels["grade"].isin({0, 1, 2}).all():
        raise ValueError("qrels contain invalid grades")
    if base_scores["score_numeric"].isna().any() or not np.isfinite(
        base_scores["score_numeric"]
    ).all():
        raise ValueError("base scores contain invalid values")

    strata = question_strata(questions)
    if set(qrels["query_id"]) != set(strata):
        raise ValueError("qrel query ids differ from adjudicated question instances")
    assignment_by_event = dict(
        zip(assignments["event_group_id"], assignments["split"], strict=True)
    )
    if set(assignment_by_event.values()) - {"train", "dev"}:
        raise ValueError("bundle accepts only train/dev event assignments")
    event_ids = set(legal_events["event_id"])
    if legal_events["event_id"].duplicated().any():
        raise ValueError("legal event catalog contains duplicate event IDs")

    intervention_by_instance: dict[tuple[str, str], dict[str, object]] = {}
    for row in interventions.to_dict(orient="records"):
        instance_key = (row["query_id"], row["snapshot_date"])
        if instance_key in intervention_by_instance:
            raise ValueError(f"duplicate intervention row for {instance_key}")
        expected_instance = f"{row['query_id']}@{row['snapshot_date']}"
        if row["query_instance_id"] != expected_instance:
            raise ValueError("query_instance_id must equal query_id@snapshot_date")
        causal = parse_ids(row["causal_event_ids"], "causal_event_ids")
        wrong = parse_ids(row["wrong_event_ids"], "wrong_event_ids")
        unrelated = parse_ids(row["unrelated_event_ids"], "unrelated_event_ids")
        if set(causal) & set(wrong) or set(causal) & set(unrelated) or set(wrong) & set(unrelated):
            raise ValueError("intervention event sets must be disjoint")
        if not set(causal + wrong + unrelated).issubset(event_ids):
            raise ValueError("intervention plan references an unknown event")
        intervention_by_instance[instance_key] = {
            **row,
            "causal": causal,
            "wrong": wrong,
            "unrelated": unrelated,
        }

    locked_system = baseline["systems"]["locked_rule"]
    locked_scores = base_scores.loc[base_scores["system"].eq(locked_system)].copy()
    score_lookup = {
        (row["query_id"], row["snapshot_date"], row["candidate_version_id"]): float(
            row["score_numeric"]
        )
        for row in locked_scores.to_dict(orient="records")
    }
    qrel_keys = {
        (row["query_id"], row["snapshot_date"], row["candidate_version_id"])
        for row in qrels.to_dict(orient="records")
    }
    if set(score_lookup) != qrel_keys:
        raise ValueError("locked rule scores do not cover exactly all qrel candidates")

    query_instances = {
        (row["query_id"], row["snapshot_date"]): {
            "event_group_id": row["event_group_id"],
            "query_group_id": row["query_group_id"],
        }
        for row in qrels.to_dict(orient="records")
    }
    if set(query_instances) != set(intervention_by_instance):
        raise ValueError("intervention plan must cover every qrel query snapshot")
    control_positives: dict[tuple[str, str], list[str]] = {}
    for (query_id, snapshot), frame in qrels.groupby(["query_id", "snapshot_date"]):
        if strata[query_id] != "control":
            continue
        group_values = set(frame["query_group_id"])
        if len(group_values) != 1:
            raise ValueError("one control query spans multiple query groups")
        positives = sorted(frame.loc[frame["grade"].ge(1), "candidate_version_id"])
        if not positives:
            raise ValueError(f"control query {query_id}@{snapshot} has no positive")
        control_positives[(next(iter(group_values)), snapshot)] = positives

    examples: list[dict[str, str]] = []
    development_queries: list[dict[str, str]] = []
    development_candidates: list[dict[str, str]] = []
    for (query_id, snapshot), frame in qrels.groupby(
        ["query_id", "snapshot_date"], sort=True
    ):
        metadata = query_instances[(query_id, snapshot)]
        event_group_id = metadata["event_group_id"]
        query_group_id = metadata["query_group_id"]
        split_name = assignment_by_event.get(event_group_id)
        if split_name is None:
            raise ValueError(f"query uses unassigned event group {event_group_id}")
        intervention = intervention_by_instance[(query_id, snapshot)]
        if intervention["event_group_id"] != event_group_id:
            raise ValueError("intervention event group differs from qrels")
        query_instance_id = f"{query_id}@{snapshot}"
        if split_name == "dev":
            phase_values = set(frame["snapshot_phase"])
            if len(phase_values) != 1:
                raise ValueError("one query instance spans multiple snapshot phases")
            snapshot_phase = next(iter(phase_values))
            if snapshot_phase not in {"pre", "post"}:
                raise ValueError("snapshot_phase must be pre or post")
            query_slice = (
                "control"
                if strata[query_id] == "control"
                else "affected_post"
                if snapshot_phase == "post"
                else "pre_update"
            )
            development_queries.append(
                {
                    "query_instance_id": query_instance_id,
                    "query_id": query_id,
                    "event_group_id": event_group_id,
                    "query_group_id": query_group_id,
                    "query_slice": query_slice,
                    "split": "dev",
                    "as_of": snapshot,
                    "stale_as_of": str(intervention["stale_as_of"]),
                    "causal_event_ids": json.dumps(
                        intervention["causal"], ensure_ascii=False
                    ),
                    "wrong_event_ids": json.dumps(
                        intervention["wrong"], ensure_ascii=False
                    ),
                    "unrelated_event_ids": json.dumps(
                        intervention["unrelated"], ensure_ascii=False
                    ),
                }
            )
            for candidate_id in sorted(frame["candidate_version_id"]):
                development_candidates.append(
                    {
                        "query_instance_id": query_instance_id,
                        "query_id": query_id,
                        "candidate_version_id": candidate_id,
                        "base_score": repr(
                            score_lookup[(query_id, snapshot, candidate_id)]
                        ),
                    }
                )
        if (
            strata[query_id] != "affected"
            or set(frame["snapshot_phase"]) != {"post"}
        ):
            continue
        positives = frame.loc[frame["grade"].ge(1)].copy()
        negatives = frame.loc[frame["grade"].eq(0)].copy()
        if positives.empty or negatives.empty:
            raise ValueError(
                f"affected query {query_instance_id} needs positive and negative qrels"
            )
        negative_ids = sorted(
            negatives["candidate_version_id"],
            key=lambda item: (-score_lookup[(query_id, snapshot, item)], item),
        )
        hardest_negative = negative_ids[0]
        unaffected_ids = control_positives.get((query_group_id, snapshot))
        if not unaffected_ids:
            raise ValueError(
                f"affected query {query_instance_id} lacks matched control positives"
            )
        positive_ids = sorted(
            positives["candidate_version_id"],
            key=lambda item: (
                -int(positives.loc[
                    positives["candidate_version_id"].eq(item), "grade"
                ].iloc[0]),
                -score_lookup[(query_id, snapshot, item)],
                item,
            ),
        )
        for position, positive_id in enumerate(positive_ids, start=1):
            stable_ids = [
                item
                for item in unaffected_ids
                if item not in {positive_id, hardest_negative}
            ]
            if not stable_ids:
                raise ValueError(
                    f"affected query {query_instance_id} has no distinct control version "
                    "for stability supervision"
                )
            examples.append(
                {
                    "example_id": f"{query_instance_id}:positive-{position}",
                    "event_group_id": event_group_id,
                    "query_group_id": query_group_id,
                    "query_id": query_id,
                    "split": split_name,
                    "as_of": snapshot,
                    "positive_version_id": positive_id,
                    "negative_version_id": hardest_negative,
                    "base_positive_score": repr(
                        score_lookup[(query_id, snapshot, positive_id)]
                    ),
                    "base_negative_score": repr(
                        score_lookup[(query_id, snapshot, hardest_negative)]
                    ),
                    "causal_event_ids": json.dumps(
                        intervention["causal"], ensure_ascii=False
                    ),
                    "unaffected_version_ids": json.dumps(
                        stable_ids, ensure_ascii=False
                    ),
                }
            )
    if not examples or not development_queries or not development_candidates:
        raise ValueError("bundle generation produced an empty required artifact")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    examples_path = args.output_dir / "l3-examples.csv"
    development_queries_path = args.output_dir / "l3-development-queries.csv"
    development_candidates_path = args.output_dir / "l3-development-candidates.csv"
    pd.DataFrame(examples).to_csv(examples_path, index=False)
    pd.DataFrame(development_queries).to_csv(development_queries_path, index=False)
    pd.DataFrame(development_candidates).to_csv(
        development_candidates_path, index=False
    )

    base_receipt_path = args.output_dir / "locked-base-score-receipt.json"
    base_receipt = {
        "schema_version": 1,
        "artifact_type": "telen_locked_base_scores",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "authorized": True,
        "system_id": locked_system,
        "expert_annotation_audit_sha256": annotation_hash,
        "event_split_audit_sha256": split_hash,
        "gold_baseline_gate_sha256": sha256_file(args.baseline_gate),
        "source_scores_sha256": sha256_file(args.base_scores),
        "examples_sha256": sha256_file(examples_path),
        "development_queries_sha256": sha256_file(development_queries_path),
        "development_candidates_sha256": sha256_file(
            development_candidates_path
        ),
        "candidate_pool_frozen": True,
        "text_floor_passed": True,
        "rule_baseline_locked": True,
        "gold_oracle_utility_passed": True,
        "event_split_authorized": True,
        "final_test_accessed": False,
    }
    base_receipt_path.write_text(
        json.dumps(base_receipt, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    evaluation_plan_path = args.output_dir / "l3-development-evaluation-plan.json"
    evaluation_config = {
        "bootstrap_samples": args.bootstrap_samples,
        "bootstrap_seed": args.bootstrap_seed,
        "minimum_event_groups": args.minimum_dev_events,
        "minimum_effect": args.minimum_learned_effect,
        "control_equivalence_margin": args.control_equivalence_margin,
        "alpha": args.alpha,
    }
    DevelopmentEvaluationConfig.from_mapping(evaluation_config)
    evaluation_plan = {
        "schema_version": 1,
        "artifact_type": "telen_l3_development_evaluation_plan",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "final_test_accessed": False,
        "development_queries": str(development_queries_path),
        "development_queries_sha256": sha256_file(development_queries_path),
        "development_candidates": str(development_candidates_path),
        "development_candidates_sha256": sha256_file(
            development_candidates_path
        ),
        "evaluation": evaluation_config,
    }
    evaluation_plan_path.write_text(
        json.dumps(evaluation_plan, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    manifest_path = args.output_dir / "l3-development-training-manifest.json"
    file_paths = {
        "annotation_audit": args.annotation_audit,
        "annotation_events": args.annotation_events,
        "annotation_questions": args.questions,
        "annotation_qrels": args.qrels,
        "base_score_receipt": base_receipt_path,
        "versions": args.versions,
        "events": args.legal_events,
        "examples": examples_path,
        "version_embeddings": args.version_embeddings,
        "query_embeddings": args.query_embeddings,
    }
    manifest = {
        "schema_version": 1,
        "artifact_type": "telen_l3_training_bundle",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "protocol_version": annotation["protocol_version"],
        "embedding_dim": args.embedding_dim,
        "base_system_id": locked_system,
        "files": {
            name: {
                "path": portable_path(path, args.output_dir),
                "sha256": sha256_file(path),
            }
            for name, path in file_paths.items()
        },
        "training": {
            "stage": "development_pilot",
            "seeds": list(PILOT_SEEDS),
            "epochs": args.epochs,
            "batch_size": args.batch_size,
            "learning_rate": args.learning_rate,
            "weight_decay": args.weight_decay,
            "hidden_dim": args.hidden_dim,
            "residual_cap": args.residual_cap,
            "initial_gate_bias": args.initial_gate_bias,
            "graph_seed_k": getattr(args, "graph_seed_k", 2),
            "ranking_margin": args.ranking_margin,
            "causal_margin": args.causal_margin,
            "causal_weight": args.causal_weight,
            "stability_weight": args.stability_weight,
            "sparsity_weight": args.sparsity_weight,
            "positive_relevance_min": 1,
            "negative_relevance_max": 0,
        },
    }
    manifest_path.write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    # Reuse the same fail-closed loader as the real runner. This catches schema,
    # qrel, provenance, embedding and leakage defects before handoff.
    loaded = load_l3_training_bundle(manifest_path)

    receipt = {
        "schema_version": 1,
        "artifact_type": "telen_l3_development_bundle_build",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "authorized": True,
        "training_manifest": str(manifest_path),
        "training_manifest_sha256": sha256_file(manifest_path),
        "locked_base_score_receipt": str(base_receipt_path),
        "locked_base_score_receipt_sha256": sha256_file(base_receipt_path),
        "development_queries": str(development_queries_path),
        "development_queries_sha256": sha256_file(development_queries_path),
        "development_candidates": str(development_candidates_path),
        "development_candidates_sha256": sha256_file(
            development_candidates_path
        ),
        "evaluation_plan": str(evaluation_plan_path),
        "evaluation_plan_sha256": sha256_file(evaluation_plan_path),
        "counts": {
            "training_examples": len(loaded.train_rows),
            "development_examples": len(loaded.dev_rows),
            "development_query_instances": len(development_queries),
            "development_candidate_rows": len(development_candidates),
        },
        "final_test_accessed": False,
    }
    receipt_path = args.output_dir / "l3-development-bundle-build-receipt.json"
    receipt_path.write_text(
        json.dumps(receipt, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return receipt


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--annotation-audit", type=Path, required=True)
    parser.add_argument("--split-audit", type=Path, required=True)
    parser.add_argument("--baseline-gate", type=Path, required=True)
    parser.add_argument("--annotation-events", type=Path, required=True)
    parser.add_argument("--questions", type=Path, required=True)
    parser.add_argument("--qrels", type=Path, required=True)
    parser.add_argument("--split-assignments", type=Path, required=True)
    parser.add_argument("--base-scores", type=Path, required=True)
    parser.add_argument("--versions", type=Path, required=True)
    parser.add_argument("--legal-events", type=Path, required=True)
    parser.add_argument("--interventions", type=Path, required=True)
    parser.add_argument("--version-embeddings", type=Path, required=True)
    parser.add_argument("--query-embeddings", type=Path, required=True)
    parser.add_argument("--embedding-dim", type=int, required=True)
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--learning-rate", type=float, default=1e-3)
    parser.add_argument("--weight-decay", type=float, default=1e-2)
    parser.add_argument("--hidden-dim", type=int, default=128)
    parser.add_argument("--residual-cap", type=float, default=0.5)
    parser.add_argument("--initial-gate-bias", type=float, default=-2.0)
    parser.add_argument("--graph-seed-k", type=int, default=2)
    parser.add_argument("--ranking-margin", type=float, default=0.2)
    parser.add_argument("--causal-margin", type=float, default=0.05)
    parser.add_argument("--causal-weight", type=float, default=1.0)
    parser.add_argument("--stability-weight", type=float, default=0.25)
    parser.add_argument("--sparsity-weight", type=float, default=0.01)
    parser.add_argument("--bootstrap-samples", type=int, default=10_000)
    parser.add_argument("--bootstrap-seed", type=int, default=20260829)
    parser.add_argument("--minimum-dev-events", type=int, default=5)
    parser.add_argument("--minimum-learned-effect", type=float, default=0.02)
    parser.add_argument("--control-equivalence-margin", type=float, default=0.02)
    parser.add_argument("--alpha", type=float, default=0.05)
    parser.add_argument("--output-dir", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    result = build(parse_args())
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()