"""Frozen DEV selection and final evaluation for the longer German NeMo Gym run.

This module prepares specifications and verifies evidence; it performs no model
inference. Optimization and the existing evaluators retain their own boundaries.
"""

from __future__ import annotations

import argparse
from collections import defaultdict
from datetime import datetime, timezone
import hashlib
import json
import math
from pathlib import Path
import random
import re
from statistics import mean

from . import agent_rl_evaluation as agent
from . import evaluation as numeric
from .agent_rl_tasks import generate_tasks


PROTOCOL = "german-nemo-long-20261020-v1"
CHECKPOINTS = [64, 128, 256]
AGENT_SEEDS = [42, 1729, 20261020]
MATH_SEEDS = [42, 20261020]
DEV_SEED = 42
DEV_TEMPLATES = [
    "Ermittle den vollständigen Auftragswert für die Kennung {record_id}.",
    "Welcher Gesamtpreis gehört zum Datensatz {record_id}?",
    "Bestimme die Kosten aller Einheiten im Auftrag {record_id}.",
    "Wie viel ist für die Bestellung mit Kennung {record_id} insgesamt zu zahlen?",
]
TEST_TEMPLATES = [
    "Gib den gesamten Warenwert zum Schlüssel {record_id} an.",
    "Welche Endsumme ergibt sich für den hinterlegten Vorgang {record_id}?",
    "Berechne den Betrag der kompletten Lieferung {record_id}.",
    "Wie hoch sind die gesamten Einkaufskosten für {record_id}?",
    "Nenne die Summe für sämtliche bestellten Stücke unter {record_id}.",
    "Welcher Gesamtbetrag ist für den Eintrag {record_id} fällig?",
    "Bestimme den Wert der vollständigen Bestellung zum Registereintrag {record_id}.",
    "Wie viel kosten alle im Beleg {record_id} verzeichneten Einheiten zusammen?",
]
SAMPLING = {"temperature": 0.6, "top_p": 0.95, "top_k": 20, "min_p": 0.0}


def _check(condition, message):
    if not condition:
        raise ValueError(message)


def _sha(raw):
    return hashlib.sha256(raw).hexdigest()


def _read(path):
    raw = Path(path).read_bytes()
    return json.loads(raw), _sha(raw)


def _rows(path):
    raw = Path(path).read_bytes()
    rows = [json.loads(line) for line in raw.splitlines() if line.strip()]
    _check(rows and len({row["id"] for row in rows}) == len(rows), "empty/duplicate dataset IDs")
    return rows, _sha(raw)


def _write(path, value):
    Path(path).write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n")


def _write_rows(path, rows):
    Path(path).write_text(
        "".join(json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n" for row in rows)
    )


def _sources():
    root = Path(__file__).parent
    names = {"gym_experiment_evaluation.py", "evaluation.py", *agent.source_hashes()}
    return {name: numeric.sha256_file(root / name) for name in sorted(names)}


def fresh_agent_tasks(split):
    """New IDs/question wordings; preserve the old scaffold and hidden-fact ranges."""
    _check(split in {"dev", "test"}, "agent split must be dev or test")
    count, seed, templates = (
        (64, 20261021, DEV_TEMPLATES) if split == "dev" else (256, 20261020, TEST_TEMPLATES)
    )
    rows = generate_tasks(split, count, seed=seed)
    for index, row in enumerate(rows):
        suffix = row["prompt"].split("\n\nWerkzeugprotokoll:", 1)[1]
        row["prompt"] = (
            templates[index % len(templates)].format(record_id=row["task"]["record_id"])
            + "\n\nWerkzeugprotokoll:"
            + suffix
        )
        row["task"]["template_id"] = f"{PROTOCOL}-{split}-{index % len(templates)}"
    return rows


def prepare(baseline, tokenizer, math_dev, mgsm, output, *, sft_train=None):
    """Prepare CPU-only frozen data/protocol. math_dev must be SFT-source disjoint."""
    output = Path(output).resolve()
    _check(not output.exists(), "choose a fresh protocol output directory")
    paths = {
        name: Path(value).resolve()
        for name, value in (
            ("baseline", baseline),
            ("tokenizer", tokenizer),
            ("math_dev", math_dev),
            ("mgsm", mgsm),
        )
    }
    paths["sft_train"] = (
        Path(sft_train).resolve() if sft_train else paths["math_dev"].parent / "train.jsonl"
    )
    dev, dev_sha = _rows(paths["math_dev"])
    train, train_sha = _rows(paths["sft_train"])
    test, test_sha = _rows(paths["mgsm"])
    _check(len(dev) == 128 and len(test) == 250, "require 128 math DEV and all 250 MGSM items")
    for rows in (train, dev, test):
        _check(len({row["source_id"] for row in rows}) == len(rows), "duplicate source identities")
    _check(
        not ({row["source_id"] for row in train} & {row["source_id"] for row in dev}),
        "math DEV overlaps original SFT training sources",
    )
    _check(
        all(row["source_id"].startswith("gsm8k/train/") for row in dev),
        "math DEV must use train-source German math, never MGSM",
    )
    _check(
        {row["source_id"] for row in test} == {f"mgsm/de/test/{i:03d}" for i in range(250)},
        "full MGSM source identity coverage required",
    )
    math_rows = []
    for row in dev:
        answer = row.get("solution", row.get("answer"))
        _check(
            isinstance(answer, str) and re.fullmatch(r"-?\d+", answer),
            "math DEV needs canonical integer answers",
        )
        math_rows.append(
            {
                "id": row["id"],
                "source_id": row["source_id"],
                "prompt": row["prompt"],
                "answer": answer,
                "scorer": "final_numeric",
                "category": "German math DEV",
            }
        )
    _check(
        all(row.get("scorer") == "final_numeric" for row in test),
        "MGSM must use strict final_numeric scorer",
    )
    template_file = paths["tokenizer"] / "chat_template.jinja"
    template = (
        template_file.read_text()
        if template_file.exists()
        else _read(paths["tokenizer"] / "tokenizer_config.json")[0].get("chat_template")
    )
    _check(
        isinstance(template, str) and template.strip(), "frozen tokenizer chat template required"
    )
    for name in ("baseline", "tokenizer"):
        _check(
            paths[name].is_dir() and not (paths[name] / "adapter_config.json").exists(),
            "use full baseline and frozen tokenizer directories",
        )
    identities = {
        name: numeric.checkpoint_sha256(paths[name]) for name in ("baseline", "tokenizer")
    }
    output.mkdir(parents=True)
    data = {
        "agent_dev": fresh_agent_tasks("dev"),
        "agent_test": fresh_agent_tasks("test"),
        "math_dev": math_rows,
        "math_test": test,
    }
    frozen = {}
    for key, rows in data.items():
        target = output / f"{key}.jsonl"
        _write_rows(target, rows)
        frozen[key] = {
            "path": str(target),
            "sha256": numeric.sha256_file(target),
            "count": len(rows),
        }
    plan = {
        "schema_version": 1,
        "protocol": PROTOCOL,
        "status": "prepared",
        "release_eligible": False,
        "baseline": {"path": str(paths["baseline"]), "sha256": identities["baseline"]},
        "tokenizer": {"path": str(paths["tokenizer"]), "sha256": identities["tokenizer"]},
        "chat_template": template,
        "chat_template_kwargs": {"enable_thinking": True},
        "sampling": SAMPLING,
        "max_completion_length": 768,
        "max_tool_calling_iterations": 3,
        "checkpoints": CHECKPOINTS,
        "dev_seeds": [DEV_SEED],
        "test_seeds": {"agent": AGENT_SEEDS, "math": MATH_SEEDS},
        "datasets": frozen,
        "source_inputs": {
            name: {"path": str(paths[name]), "sha256": digest}
            for name, digest in (
                ("math_dev", dev_sha),
                ("sft_train", train_sha),
                ("mgsm", test_sha),
            )
        },
        "source_sha256": _sources(),
        "selection_rule": "Only DEV: math correct >= SFT baseline AND strict agent success > baseline; rank by agent successes, fewer tool errors, earliest step. Otherwise recommend retaining SFT baseline but assess the preregistered step256 diagnostic candidate once; final tests cannot promote it.",
        "final_gate": "Strict observed MGSM zero drop pooled AND each fixed seed; passing does not prove no harm.",
        "bootstrap": {
            "method": "paired item percentile; seeds grouped by item",
            "seed": 20261020,
            "resamples": 10000,
        },
        "language_audit": {
            "status": "planned_unassessed",
            "seed": 42,
            "agent_ids": [row["id"] for row in data["agent_test"][:10]],
            "math_ids": [row["id"] for row in test[:10]],
            "rubric": "Inspect natural-language reasoning separately from quoted input/tool JSON/numeric final answers; record dominant language, closed thinking, and answer format. No inferred fluency or explanation-correctness score; identify agent versus human reviewer.",
        },
        "limitations": [
            "One synthetic scaffolded lookup/calculator family with only 80 ordered operand pairs; same numeric ranges and tools, new question templates/IDs, no broad generalization claim.",
            "Single training seed. Item bootstrap excludes training-seed variance and unmeasured generation-seed variance; related synthetic items remain correlated.",
            "MGSM64 from the earlier pilot was already inspected. The full250 set is not wholly unseen; no MGSM results are used for selection or changes.",
            "Strict final-answer protocol scores reflect format and token limits as well as accepted answer correctness.",
            "DEV math retains imperfect translated explanations; selection uses final integer gold, not translation quality.",
            "The retained 768-token agent stream includes tool observations/framing. Native training may physically generate and discard extra tokens; final evaluation samples only the remaining stream budget. Retained budgets do not establish equal physical generation compute.",
        ],
    }
    _write(output / "plan.json", plan)
    return plan


def verify_plan(path):
    plan, digest = _read(path)
    _check(
        plan.get("protocol") == PROTOCOL and plan["source_sha256"] == _sources(),
        "evaluation protocol/source identity mismatch",
    )
    _check(
        plan["checkpoints"] == CHECKPOINTS
        and plan["test_seeds"] == {"agent": AGENT_SEEDS, "math": MATH_SEEDS}
        and plan["dev_seeds"] == [DEV_SEED],
        "frozen checkpoint/seed schedule mismatch",
    )
    _check(
        plan["sampling"] == SAMPLING
        and plan["max_completion_length"] == 768
        and plan["max_tool_calling_iterations"] == 3
        and plan["chat_template_kwargs"] == {"enable_thinking": True},
        "frozen generation protocol mismatch",
    )
    _check(
        plan["bootstrap"]
        == {
            "method": "paired item percentile; seeds grouped by item",
            "seed": 20261020,
            "resamples": 10000,
        },
        "frozen bootstrap protocol mismatch",
    )
    for split in ("dev", "test"):
        rows, _ = _rows(plan["datasets"][f"agent_{split}"]["path"])
        _check(
            rows == fresh_agent_tasks(split), "frozen new task IDs/templates differ from protocol"
        )
    for value in [*plan["datasets"].values(), *plan["source_inputs"].values()]:
        _check(numeric.sha256_file(value["path"]) == value["sha256"], "frozen input changed")
    return plan, digest


def write_specs(plan_path, checkpoint, label, split="dev", *, selection_path=None):
    """Return/write the list manifest consumed by an external GPU orchestrator."""
    plan, digest = verify_plan(plan_path)
    _check(
        split in {"dev", "test"} and re.fullmatch(r"[A-Za-z0-9_-]+", label), "invalid split/label"
    )
    checkpoint = Path(checkpoint).resolve()
    _check(
        checkpoint.is_dir() and not (checkpoint / "adapter_config.json").exists(),
        "evaluation requires a standalone full/merged checkpoint",
    )
    model_hash = numeric.checkpoint_sha256(checkpoint)
    selection_hash = None
    if split == "test":
        _check(selection_path is not None, "final test specs require frozen DEV selection first")
        decision, selection_hash = _verified_selection(plan_path, selection_path)
        _check(
            model_hash in {plan["baseline"]["sha256"], decision["selected"]["sha256"]},
            "test model is neither baseline nor DEV-selected checkpoint",
        )
    root = Path(plan_path).resolve().parent
    spec_dir = root / "specifications" / label / split
    _check(not spec_dir.exists(), "specification manifest already exists")
    spec_dir.mkdir(parents=True)
    entries = []
    for kind in ("agent", "math"):
        seeds = plan["dev_seeds"] if split == "dev" else plan["test_seeds"][kind]
        data = plan["datasets"][f"{kind}_{split}"]
        for seed in seeds:
            sampling = plan["sampling"] | {"seed": seed}
            if kind == "agent":
                spec = {
                    "schema_version": 1,
                    "model": {"path": str(checkpoint), "sha256": model_hash},
                    "tokenizer": plan["tokenizer"],
                    "tasks": {k: data[k] for k in ("path", "sha256")},
                    "sampling": sampling,
                    "max_completion_length": plan["max_completion_length"],
                    "max_tool_calling_iterations": plan["max_tool_calling_iterations"],
                    "chat_template_kwargs": plan["chat_template_kwargs"],
                    "trials": 1,
                    "device": "cuda",
                    "dtype": "bfloat16",
                }
            else:
                spec = {
                    "schema_version": 1,
                    "benchmark": {
                        "id": f"{PROTOCOL}-math-{split}",
                        "version": "1",
                        "status": "research-pilot",
                        **data,
                    },
                    "model": {
                        "id": label,
                        "revision": model_hash,
                        "tokenizer_id": "frozen-qwen3-tokenizer",
                        "tokenizer_revision": plan["tokenizer"]["sha256"],
                        "chat_template": plan["chat_template"],
                        "chat_template_kwargs": plan["chat_template_kwargs"],
                    },
                    "sampling": sampling | {"max_new_tokens": plan["max_completion_length"]},
                    "trials": 1,
                    "scorer_version": "1",
                    "harness_version": "1",
                    "adapter": {
                        "type": "transformers",
                        "path": str(checkpoint),
                        "sha256": model_hash,
                        "tokenizer_path": plan["tokenizer"]["path"],
                        "tokenizer_sha256": plan["tokenizer"]["sha256"],
                        "device": "cuda",
                        "dtype": "bfloat16",
                    },
                }
            target = spec_dir / f"{kind}-{seed}.json"
            _write(target, spec)
            entries.append(
                {
                    "kind": kind,
                    "seed": seed,
                    "path": str(
                        root / "evaluations" / label / split / f"{kind}-{seed}" / "report.json"
                    ),
                    "spec_path": str(target),
                    "spec_sha256": numeric.sha256_file(target),
                    "split": split,
                    "plan_sha256": digest,
                    "selection_sha256": selection_hash,
                }
            )
    _write(spec_dir / "manifest.json", entries)
    return entries


def _load_suite(plan, plan_digest, manifest_path, split, *, selection_digest=None):
    entries, manifest_digest = _read(manifest_path)
    _check(isinstance(entries, list), "report manifest must be a list")
    expected = {
        (kind, seed)
        for kind in ("agent", "math")
        for seed in (plan["dev_seeds"] if split == "dev" else plan["test_seeds"][kind])
    }
    _check(
        len(entries) == len(expected)
        and {(entry["kind"], entry["seed"]) for entry in entries} == expected,
        "report manifest missing, duplicate or wrong fixed seeds",
    )
    reports, evidence, model_ids = {}, [], set()
    for entry in entries:
        _check(
            entry.get("split") == split and entry.get("plan_sha256") == plan_digest,
            "report manifest uses wrong split/protocol",
        )
        _check(
            entry.get("selection_sha256") == selection_digest,
            "report manifest selection identity mismatch",
        )
        spec, spec_digest = _read(entry["spec_path"])
        _check(spec_digest == entry["spec_sha256"], "frozen specification changed")
        report, report_digest = _read(entry["path"])
        kind, seed = entry["kind"], entry["seed"]
        _check(
            report["provenance"]["spec"] == spec
            and spec["sampling"]["seed"] == seed
            and spec["trials"] == 1,
            "report actual seed/specification mismatch",
        )
        data = plan["datasets"][f"{kind}_{split}"]
        dataset_rows, _ = _rows(data["path"])
        raw_results = report["episodes"] if kind == "agent" else report["results"]
        _check(
            len(raw_results) == data["count"]
            and {(row["id"], row["trial"]) for row in raw_results}
            == {(row["id"], 0) for row in dataset_rows},
            "report is missing full frozen task coverage",
        )
        _check(
            spec["sampling"]
            == plan["sampling"]
            | {"seed": seed}
            | ({"max_new_tokens": 768} if kind == "math" else {}),
            "report sampling protocol mismatch",
        )
        if kind == "agent":
            agent.verify_report(report)
            _check(
                all(
                    row["seed"] == (seed + int(_sha(row["id"].encode())[:8], 16)) % (2**32)
                    for row in raw_results
                ),
                "agent captured episode seed differs from fixed seed policy",
            )
            _check(
                all(
                    math.isfinite(row["wall_seconds"]) and row["wall_seconds"] >= 0
                    for row in raw_results
                ),
                "invalid agent timing evidence",
            )
            _check(
                report["source_kind"] == "captured",
                "fixture agent reports cannot select or qualify a model",
            )
            _check(
                spec["tasks"]["sha256"] == data["sha256"]
                and spec["tokenizer"] == plan["tokenizer"]
                and spec["max_completion_length"] == 768
                and spec["max_tool_calling_iterations"] == 3
                and spec["chat_template_kwargs"] == plan["chat_template_kwargs"]
                and spec["device"] == "cuda"
                and spec["dtype"] == "bfloat16",
                "agent protocol identity mismatch",
            )
            _check(
                report["provenance"]["runtime"]["device"] == "cuda"
                and report["provenance"]["runtime"]["dtype"] == "bfloat16",
                "actual agent runtime differs from frozen CUDA/BF16 protocol",
            )
            model_ids.add((spec["model"]["path"], spec["model"]["sha256"]))
        else:
            numeric._verify_report(report)
            _check(
                report["provenance"]["execution_mode"] == "transformers",
                "fixture math reports cannot select or qualify a model",
            )
            _check(
                spec["benchmark"]["sha256"] == data["sha256"]
                and spec["benchmark"]["count"] == data["count"]
                and spec["model"]["tokenizer_revision"] == plan["tokenizer"]["sha256"]
                and spec["model"]["chat_template"] == plan["chat_template"]
                and spec["model"]["chat_template_kwargs"] == plan["chat_template_kwargs"]
                and spec["adapter"]["tokenizer_sha256"] == plan["tokenizer"]["sha256"]
                and spec["adapter"]["sha256"] == spec["model"]["revision"]
                and spec["adapter"]["device"] == "cuda"
                and spec["adapter"]["dtype"] == "bfloat16",
                "math protocol identity mismatch",
            )
            model_ids.add((spec["adapter"]["path"], spec["model"]["revision"]))
        reports[kind, seed] = report
        evidence.append(
            {"kind": kind, "seed": seed, "path": entry["path"], "sha256": report_digest}
        )
    _check(len(model_ids) == 1, "one suite must evaluate exactly one checkpoint")
    model_path, model_hash = next(iter(model_ids))
    return {
        "reports": reports,
        "model": {"path": model_path, "sha256": model_hash},
        "evidence": evidence,
        "manifest_sha256": manifest_digest,
    }


def _compatible(suites):
    """Compare actual runtime/parser/package identities across models and seeds."""
    for kind in ("agent", "math"):
        signatures = []
        for suite in suites:
            for (name, _), report in suite["reports"].items():
                if name != kind:
                    continue
                provenance = report["provenance"]
                signature = {"packages": provenance["packages"]}
                if kind == "agent":
                    signature.update(
                        {
                            key: provenance[key]
                            for key in ("effective_template_sha256", "tool_schemas", "sources")
                        }
                    )
                else:
                    signature["scorer_sha256"] = provenance["scorer_sha256"]
                signatures.append(signature)
        _check(
            signatures and all(value == signatures[0] for value in signatures),
            "incompatible actual evaluator/template/package identities",
        )


def _values(suite, kind):
    result = {}
    for (name, seed), report in suite["reports"].items():
        if name != kind:
            continue
        rows = report["episodes"] if kind == "agent" else report["results"]
        for row in rows:
            result[row["id"], seed] = (
                bool(row["outcome"]["success"]) if kind == "agent" else row["outcome"] == "correct"
            )
    return result


def _choose(plan_path, baseline_manifest, candidate_manifests):
    plan, digest = verify_plan(plan_path)
    _check(
        {int(step) for step in candidate_manifests} == set(CHECKPOINTS)
        and len(candidate_manifests) == 3,
        "selection requires only all preregistered DEV checkpoints 64/128/256",
    )
    baseline = _load_suite(plan, digest, baseline_manifest, "dev")
    _check(baseline["model"]["sha256"] == plan["baseline"]["sha256"], "DEV baseline model mismatch")
    base_agent, base_math = (
        sum(_values(baseline, "agent").values()),
        sum(_values(baseline, "math").values()),
    )
    candidates, eligible = [], []
    for step in CHECKPOINTS:
        manifest = candidate_manifests.get(step, candidate_manifests.get(str(step)))
        suite = _load_suite(plan, digest, manifest, "dev")
        _compatible([baseline, suite])
        correct, math_correct = (
            sum(_values(suite, "agent").values()),
            sum(_values(suite, "math").values()),
        )
        errors = sum(
            report["metrics"]["tool_errors"]
            for (kind, _), report in suite["reports"].items()
            if kind == "agent"
        )
        row = {
            "step": step,
            "model": suite["model"],
            "agent_correct": correct,
            "math_correct": math_correct,
            "tool_errors": errors,
            "admissible": math_correct >= base_math,
            "agent_dev_improvement": correct > base_agent,
            "manifest_path": str(Path(manifest).resolve()),
            "manifest_sha256": suite["manifest_sha256"],
            "evidence": suite["evidence"],
        }
        candidates.append(row)
        if row["admissible"] and row["agent_dev_improvement"]:
            eligible.append(row)
    chosen = (
        min(eligible, key=lambda row: (-row["agent_correct"], row["tool_errors"], row["step"]))
        if eligible
        else None
    )
    return {
        "schema_version": 1,
        "protocol": PROTOCOL,
        "plan_sha256": digest,
        "release_eligible": False,
        "baseline_manifest": str(Path(baseline_manifest).resolve()),
        "baseline_manifest_sha256": baseline["manifest_sha256"],
        "baseline_dev": {
            "agent_correct": base_agent,
            "math_correct": base_math,
            "evidence": baseline["evidence"],
        },
        "candidate_manifests": {
            str(k): str(Path(v).resolve()) for k, v in candidate_manifests.items()
        },
        "candidates": candidates,
        "selected_step": chosen["step"] if chosen else 256,
        "selected": chosen["model"] if chosen else candidates[-1]["model"],
        "fallback_to_baseline": chosen is None,
        "retain_baseline": chosen is None,
        "dev_acceptance": chosen is not None,
        "recommendation": chosen["model"] if chosen else plan["baseline"],
        "assessment_reason": "DEV-admissible improvement"
        if chosen
        else "preregistered step256 diagnostic; retain baseline regardless of final test outcomes",
        "reason": "DEV-only selection; MGSM and final agent results were not inputs",
    }


def select(plan_path, baseline_manifest, candidate_manifests, output):
    output = Path(output).resolve()
    _check(not output.exists(), "choose a fresh selection output directory")
    decision = _choose(plan_path, baseline_manifest, candidate_manifests)
    output.mkdir(parents=True)
    _write(output / "selection.json", decision)
    return decision


def _verified_selection(plan_path, selection_path):
    decision, digest = _read(selection_path)
    reproduced = _choose(plan_path, decision["baseline_manifest"], decision["candidate_manifests"])
    _check(
        decision == reproduced, "DEV selection changed or differs from verified deterministic rule"
    )
    return decision, digest


def _paired(a, b, bootstrap):
    _check(a.keys() == b.keys() and a, "paired task/seed identities differ")
    by_item = defaultdict(list)
    for key in sorted(a):
        by_item[key[0]].append(int(b[key]) - int(a[key]))
    differences = [mean(values) for values in by_item.values()]
    rng, count = random.Random(bootstrap["seed"]), bootstrap["resamples"]
    draws = sorted(mean(rng.choices(differences, k=len(differences))) for _ in range(count))
    seeds = sorted({key[1] for key in a})
    return {
        "items": len(by_item),
        "paired_trials": len(a),
        "baseline_correct": sum(a.values()),
        "candidate_correct": sum(b.values()),
        "baseline_rate": mean(a.values()),
        "candidate_rate": mean(b.values()),
        "wins": sum(b[key] and not a[key] for key in a),
        "losses": sum(a[key] and not b[key] for key in a),
        "both_correct": sum(a[key] and b[key] for key in a),
        "both_incorrect": sum(not a[key] and not b[key] for key in a),
        "delta": mean(differences),
        "per_seed": {
            str(seed): {
                "total": sum(key[1] == seed for key in a),
                "baseline_correct": sum(value for key, value in a.items() if key[1] == seed),
                "candidate_correct": sum(value for key, value in b.items() if key[1] == seed),
                "wins": sum(b[key] and not a[key] for key in a if key[1] == seed),
                "losses": sum(a[key] and not b[key] for key in a if key[1] == seed),
            }
            for seed in seeds
        },
        "bootstrap": {
            **bootstrap,
            "ci_95": [draws[int(0.025 * (count - 1))], draws[int(0.975 * (count - 1))]],
            "limitation": "Descriptive item resampling groups fixed seeds together; it excludes training-seed variance, unmeasured generation seeds and correlations between related task items.",
        },
    }


def _counts(suite, kind, seed=None):
    reports = [
        report
        for (name, report_seed), report in suite["reports"].items()
        if name == kind and (seed is None or seed == report_seed)
    ]
    if kind == "agent":
        keys = (
            "episodes",
            "successful",
            "tool_calls",
            "tool_errors",
            "malformed_final_outputs",
            "turn_cap_episodes",
            "token_cap_episodes",
            "assistant_tokens",
            "prompt_tokens",
            "completion_stream_tokens",
            "wall_seconds",
        )
        result = {key: sum(report["metrics"][key] for report in reports) for key in keys}
        result["workflow_successes"] = sum(
            episode["outcome"]["workflow_valid"]
            for report in reports
            for episode in report["episodes"]
        )
        final_correct = 0
        for report in reports:
            tasks = {
                row["id"]: row["task"]
                for row in (
                    json.loads(line)
                    for line in report["provenance"]["tasks_snapshot"].splitlines()
                    if line.strip()
                )
            }
            final_correct += sum(
                episode["outcome"]["final_answer"]
                == str(
                    tasks[episode["id"]]["private_facts"]["quantity"]
                    * tasks[episode["id"]]["private_facts"]["unit_price"]
                )
                for episode in report["episodes"]
            )
        result["correct_accepted_final_answers"] = final_correct
        result["final_answer_note"] = (
            "Accepted final numeric answer independently of required tool workflow; never replaces strict workflow success."
        )
        return result
    rows = [row for report in reports for row in report["results"]]
    _check(
        all(isinstance(row.get("usage"), dict) for row in rows),
        "local math token usage must be captured",
    )
    _check(
        all(
            type(row["usage"].get(key)) is int and row["usage"][key] >= 0
            for row in rows
            for key in ("prompt_tokens", "completion_tokens")
        )
        and all(
            row["usage"]["completion_tokens"] <= 768
            and math.isfinite(row["wall_seconds"])
            and row["wall_seconds"] >= 0
            for row in rows
        ),
        "invalid math usage/timing evidence",
    )
    return {
        "trials": len(rows),
        "correct": sum(row["outcome"] == "correct" for row in rows),
        "malformed_outputs": sum(numeric._final_numeric(row["output"]) is None for row in rows),
        "wrong_accepted_answers": sum(
            row["outcome"] != "correct" and numeric._final_numeric(row["output"]) is not None
            for row in rows
        ),
        "hit_token_cap_trials": sum(row["usage"]["completion_tokens"] >= 768 for row in rows),
        "prompt_tokens": sum(row["usage"]["prompt_tokens"] for row in rows),
        "completion_tokens": sum(row["usage"]["completion_tokens"] for row in rows),
        "generation_seconds": sum(row["wall_seconds"] for row in rows),
        "token_cap_note": "Hit the declared cap; absent finish reasons this does not prove truncation.",
    }


def summarize(plan_path, selection_path, baseline_manifest, candidate_manifest, output):
    output = Path(output).resolve()
    _check(not output.exists(), "choose a fresh summary output directory")
    plan, digest = verify_plan(plan_path)
    selection, selection_digest = _verified_selection(plan_path, selection_path)
    suites = [
        _load_suite(plan, digest, path, "test", selection_digest=selection_digest)
        for path in (baseline_manifest, candidate_manifest)
    ]
    _compatible(suites)
    _check(
        suites[0]["model"]["sha256"] == plan["baseline"]["sha256"]
        and suites[1]["model"] == selection["selected"],
        "final evidence is not baseline versus DEV-selected checkpoint",
    )
    results = {
        kind: _paired(*[_values(suite, kind) for suite in suites], plan["bootstrap"])
        for kind in ("agent", "math")
    }
    retention = results["math"]
    tasks, _ = _rows(plan["datasets"]["agent_test"]["path"])
    templates = {row["id"]: row["task"]["template_id"] for row in tasks}
    agent_values = [_values(suite, "agent") for suite in suites]
    template_breakdown = {}
    for template in sorted(set(templates.values())):
        keys = [key for key in agent_values[0] if templates[key[0]] == template]
        old, new = [sum(values[key] for key in keys) for values in agent_values]
        template_breakdown[template] = {
            "paired_trials": len(keys),
            "baseline_correct": old,
            "candidate_correct": new,
            "delta": (new - old) / len(keys),
        }
    seed_gates = {
        seed: row["candidate_correct"] >= row["baseline_correct"]
        for seed, row in retention["per_seed"].items()
    }
    observed_gate = retention["candidate_correct"] >= retention["baseline_correct"] and all(
        seed_gates.values()
    )
    overall_acceptance = (
        selection["dev_acceptance"] and results["agent"]["delta"] > 0 and observed_gate
    )
    summary = {
        "schema_version": 1,
        "protocol": PROTOCOL,
        "status": "complete",
        "release_eligible": False,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "plan_sha256": digest,
        "selection_sha256": selection_digest,
        "selected_step": selection["selected_step"],
        "fallback_to_baseline": selection["fallback_to_baseline"],
        "dev_acceptance": selection["dev_acceptance"],
        "dev_retain_baseline": selection["retain_baseline"],
        "retain_baseline": not overall_acceptance,
        "recommendation": selection["selected"] if overall_acceptance else plan["baseline"],
        "assessment_reason": selection["assessment_reason"],
        "selected": selection["selected"],
        "agent": results["agent"],
        "agent_template_breakdown": template_breakdown,
        "math": retention,
        "observed_zero_drop_gate": {
            "passed": observed_gate,
            "per_seed": seed_gates,
            "note": "Observed pooled and per-seed zero drop only; a pass does not establish no harm or equivalence.",
        },
        "observed_joint_goal_met": results["agent"]["delta"] > 0 and observed_gate,
        "overall_acceptance": overall_acceptance,
        "agent_improvement_uncertainty": "positive descriptive interval"
        if results["agent"]["bootstrap"]["ci_95"][0] > 0
        else "inconclusive or no observed gain",
        "counts": {
            label: {kind: _counts(suite, kind) for kind in ("agent", "math")}
            for label, suite in zip(("baseline", "candidate"), suites, strict=True)
        },
        "per_seed_counts": {
            label: {
                kind: {str(seed): _counts(suite, kind, seed) for seed in plan["test_seeds"][kind]}
                for kind in ("agent", "math")
            }
            for label, suite in zip(("baseline", "candidate"), suites, strict=True)
        },
        "evidence": {
            label: suite["evidence"]
            for label, suite in zip(("baseline", "candidate"), suites, strict=True)
        },
        "language_audit": plan["language_audit"],
        "limitations": plan["limitations"],
    }
    output.mkdir(parents=True)
    _write(output / "summary.json", summary)
    lines = [
        "# Longer German NeMo Gym experiment",
        "",
        f"Assessment optimizer step: {selection['selected_step']}; DEV acceptance: {selection['dev_acceptance']}; retain baseline: {not overall_acceptance}.",
        selection["assessment_reason"],
        "",
        "| Final metric | Baseline | Candidate | Paired delta |",
        "| --- | ---: | ---: | ---: |",
    ]
    for kind, name in (
        ("agent", "Strict scaffolded agent workflow success"),
        ("math", "MGSM strict final-answer protocol accuracy"),
    ):
        row = results[kind]
        lines.append(
            f"| {name} | {row['baseline_correct']}/{row['paired_trials']} | {row['candidate_correct']}/{row['paired_trials']} | {row['delta']:+.3%} |"
        )
    lines += [
        "",
        f"Observed MGSM pooled/per-seed zero-drop gate: **{'passed' if observed_gate else 'regression'}**. Passing observed counts does not prove no harm.",
        "",
        "Descriptive paired item bootstrap intervals group all fixed generation seeds by item; they exclude training-seed variance. See summary.json for per-seed wins/losses, intervals, usage and caps.",
        "",
        "Language inspection remains planned/unassessed until the fixed samples receive a documented reviewer audit. No release eligibility or general German reasoning claim.",
    ]
    (output / "summary.md").write_text("\n".join(lines) + "\n")
    return summary


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    prep = commands.add_parser("prepare")
    for key in ("baseline", "tokenizer", "math-dev", "mgsm", "output"):
        prep.add_argument("--" + key, type=Path, required=True)
    prep.add_argument("--sft-train", type=Path)
    specs = commands.add_parser("specs")
    for key in ("plan", "checkpoint"):
        specs.add_argument("--" + key, type=Path, required=True)
    specs.add_argument("--label", required=True)
    specs.add_argument("--split", choices=["dev", "test"], default="dev")
    specs.add_argument("--selection", type=Path)
    choose = commands.add_parser("select")
    for key in ("plan", "baseline-manifest", "candidates", "output"):
        choose.add_argument("--" + key, type=Path, required=True)
    report = commands.add_parser("summarize")
    for key in ("plan", "selection", "baseline-manifest", "candidate-manifest", "output"):
        report.add_argument("--" + key, type=Path, required=True)
    args = parser.parse_args()
    if args.command == "prepare":
        result = prepare(
            args.baseline,
            args.tokenizer,
            args.math_dev,
            args.mgsm,
            args.output,
            sft_train=args.sft_train,
        )
    elif args.command == "specs":
        result = write_specs(
            args.plan, args.checkpoint, args.label, args.split, selection_path=args.selection
        )
    elif args.command == "select":
        result = select(args.plan, args.baseline_manifest, _read(args.candidates)[0], args.output)
    else:
        result = summarize(
            args.plan, args.selection, args.baseline_manifest, args.candidate_manifest, args.output
        )
    print(json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False))


if __name__ == "__main__":
    main()