engines/fixture_builder.py

from __future__ import annotations

import json
import re
from functools import lru_cache
from pathlib import Path
from typing import Any, Mapping

from engines.compare import build_summary_sentence
from engines.presets import MODEL_BY_KEY, PRESET_BY_KEY

FIXTURE_VERSION = "v7"

REPO_ROOT = Path(__file__).resolve().parents[1]
BUNDLE_ROOT = REPO_ROOT / "data" / "benchmark_bundle"
COMPACT_BUNDLE_PATH = BUNDLE_ROOT / "space_benchmark_bundle.json"
BACKEND_TRUTH_BUNDLE_PATH = BUNDLE_ROOT / "backend_truth_source.json"

MODEL_DIR_BY_KEY = {
    "qwen35_4b_hf": "qwen35_4b",
    "qwen35_9b_hf": "qwen35_9b",
    "qwen35_27b_hf": "qwen35_27b",
}

COMPACT_TASK_ROWS: dict[str, dict[int, dict[str, Any]]] = {
    "qwen35_4b_hf": {
        1024: {"task_name": "instruction_constraints", "task_family": "instruction", "task_variant": "constraints"},
        2048: {"task_name": "reasoning_arithmetic", "task_family": "reasoning", "task_variant": "arithmetic"},
    },
    "qwen35_9b_hf": {
        1024: {"task_name": "instruction_constraints", "task_family": "instruction", "task_variant": "constraints"},
        2048: {"task_name": "retrieval_passkey", "task_family": "retrieval", "task_variant": "passkey"},
    },
    "qwen35_27b_hf": {
        1024: {"task_name": "instruction_constraints", "task_family": "instruction", "task_variant": "constraints"},
        2048: {"task_name": "reasoning_arithmetic", "task_family": "reasoning", "task_variant": "arithmetic"},
    },
}

LONG_BENCH_ROWS: dict[str, dict[int, dict[str, Any]]] = {
    model_key: {
        4096: {"dataset": "hotpotqa", "row_index": 0, "prompt_id": "hotpot_case_order"},
        8192: {"dataset": "hotpotqa", "row_index": 0, "prompt_id": "hotpot_case_order"},
    }
    for model_key in MODEL_DIR_BY_KEY
}

MODE_TO_PROFILE = {
    "dense": None,
    "m0": "quality",
    "m0_selective_k": "systems",
}

BACKEND_MODE_TO_VARIANT = {
    "dense": None,
    "m0": "shortlist_base",
    "m0_selective_k": "learned_selector",
}

TASK_PROFILE_TO_SUMMARY_FIELD = {
    "exact": "exact_decode_ms_per_step",
    "quality": "quality_decode_ms_per_step",
    "systems": "systems_decode_ms_per_step",
}

TASK_PROFILE_TO_SUCCESS_FIELD = {
    "exact": "exact_success",
    "quality": "quality_success",
    "systems": "systems_success",
}


def mib_to_bytes(value: float) -> int:
    return int(round(float(value) * 1024.0 * 1024.0))


def _model_dir(model_key: str) -> str:
    try:
        return MODEL_DIR_BY_KEY[model_key]
    except KeyError as exc:
        raise ValueError(f"Unsupported benchmark bundle model: {model_key}") from exc


@lru_cache(maxsize=None)
def _space_bundle() -> dict[str, Any]:
    return json.loads(COMPACT_BUNDLE_PATH.read_text(encoding="utf-8"))


@lru_cache(maxsize=None)
def _backend_truth_bundle() -> dict[str, Any]:
    if BACKEND_TRUTH_BUNDLE_PATH.exists():
        return json.loads(BACKEND_TRUTH_BUNDLE_PATH.read_text(encoding="utf-8"))
    return _space_bundle().get("backend_truth", {})


def _clean_text(value: str) -> str:
    cleaned = str(value or "").replace("\r\n", "\n")
    cleaned = re.sub(r"(?is)<think>.*?</think>", " ", cleaned)
    cleaned = cleaned.replace("<|im_start|>", " ").replace("<|im_end|>", " ")
    return " ".join(cleaned.split()).strip()


def _pick_response_text(record: Mapping[str, Any]) -> str:
    task_lines = record.get("task_generated_lines")
    if isinstance(task_lines, list) and task_lines:
        first_lines = [str(line).strip() for line in task_lines[:2] if str(line).strip()]
        if first_lines:
            return "\n".join(first_lines)
    task_value = _clean_text(str(record.get("task_generated_value") or ""))
    if task_value:
        return task_value
    for key in (
        "task_generated_text_cleaned",
        "longbench_generated_text_cleaned",
        "dotcache_text",
        "dense_text",
    ):
        text = _clean_text(str(record.get(key) or ""))
        if text:
            return text
    return ""


def _backend_output_text(*, variant_label: str, prompt_length: int, generated_ids: list[int]) -> str:
    token_count = len(generated_ids)
    if token_count <= 0:
        return f"{variant_label} replay on the exact-length backend-truth prompt at {prompt_length} tokens."
    return (
        f"{variant_label} replay on the exact-length backend-truth prompt at {prompt_length} tokens. "
        f"This benchmark row records an {token_count}-token decode sample."
    )


def _make_payload(*, text: str, latency_ms: float, kv_bytes: int, trace: list[dict[str, Any]]) -> dict[str, Any]:
    tok_per_sec = 0.0 if latency_ms <= 0.0 else round(1000.0 / latency_ms, 3)
    return {
        "text": text,
        "tok_per_sec": tok_per_sec,
        "latency_ms_per_token": round(float(latency_ms), 3),
        "kv_bytes": int(kv_bytes),
        "trace": trace,
    }


def _compact_task_result(request: Mapping[str, Any], *, model_key: str, context_length: int, mode: str) -> dict[str, Any]:
    row_bundle = _space_bundle()["compact_task"].get(model_key, {}).get(str(context_length))
    if row_bundle is None:
        raise ValueError(f"No compact-task benchmark row is bundled for model={model_key} at {context_length} tokens.")
    task_name = str(row_bundle["task_name"])
    candidate_profile = MODE_TO_PROFILE[mode]
    baseline_profile = row_bundle["profiles"]["exact"]
    candidate_bundle = baseline_profile if candidate_profile is None else row_bundle["profiles"][candidate_profile]
    baseline_latency = float(baseline_profile["latency_ms"])
    candidate_latency = baseline_latency if candidate_profile is None else float(candidate_bundle["latency_ms"])
    baseline_kv_bytes = int(baseline_profile["resident_bytes"])
    candidate_kv_bytes = baseline_kv_bytes if candidate_profile is None else int(candidate_bundle["resident_bytes"])
    baseline_text = str(baseline_profile["text"])
    candidate_text = baseline_text if candidate_profile is None else str(candidate_bundle["text"])
    exact_success = bool(baseline_profile["success"])
    candidate_success = exact_success if candidate_profile is None else bool(candidate_bundle["success"])
    agreement = 1.0 if exact_success == candidate_success else 0.0
    speedup = 1.0 if candidate_profile is None else baseline_latency / max(candidate_latency, 1e-8)
    memory_reduction = 1.0 if candidate_profile is None else baseline_kv_bytes / max(candidate_kv_bytes, 1)
    candidate_label = {
        None: "Exact reference",
        "quality": "Quality profile",
        "systems": "Systems profile",
    }[candidate_profile]
    task_label = task_name.replace("_", " ")
    comparison = {
        "agreement": agreement,
        "speedup": speedup,
        "memory_reduction": memory_reduction,
        "summary": build_summary_sentence(agreement, speedup, memory_reduction),
        "notes": [
            "paper_fixture",
            "benchmark_bundle:qwen_results_matrix_20260404",
            f"fixture_version:{FIXTURE_VERSION}",
            "benchmark_section:compact_task",
            f"task_name:{task_name}",
        ],
        "paper_fixture": True,
        "paper_section": "Compact-task matrix",
        "paper_headline": f"Benchmark-backed compact-task replay for {MODEL_BY_KEY[model_key].label} on {task_label}.",
        "paper_subtitle": (
            f"Representative {task_label} row bundled from the Qwen compact-task benchmark at {context_length} tokens."
        ),
        "paper_summary": (
            f"Exact latency {baseline_latency:.1f} ms/token versus {candidate_label.lower()} "
            f"{candidate_latency:.1f} ms/token on the bundled benchmark row."
        ),
        "paper_metric_badge": "Task success unchanged" if agreement >= 1.0 else "Task outcome changed",
        "baseline_label": "Exact reference",
        "candidate_label": candidate_label,
    }
    baseline = _make_payload(
        text=baseline_text,
        latency_ms=baseline_latency,
        kv_bytes=baseline_kv_bytes,
        trace=[
            {"name": "paper_fixture", "value": 1, "unit": "flag"},
            {"name": "benchmark_section", "value": "compact_task", "unit": "label"},
            {"name": "task_name", "value": task_name, "unit": "label"},
        ],
    )
    candidate = _make_payload(
        text=candidate_text,
        latency_ms=candidate_latency,
        kv_bytes=candidate_kv_bytes,
        trace=[
            {"name": "paper_fixture", "value": 1, "unit": "flag"},
            {"name": "benchmark_section", "value": "compact_task", "unit": "label"},
            {"name": "task_name", "value": task_name, "unit": "label"},
            {"name": "profile", "value": candidate_label, "unit": "label"},
        ],
    )
    return {
        "request": dict(request),
        "baseline": baseline,
        "candidate": candidate,
        "comparison": comparison,
    }


def _longbench_result(request: Mapping[str, Any], *, model_key: str, context_length: int, mode: str) -> dict[str, Any]:
    row_bundle = _space_bundle()["longbench_mini"].get(model_key, {}).get(str(context_length))
    if row_bundle is None:
        raise ValueError(f"No LongBench benchmark row is bundled for model={model_key} at {context_length} tokens.")
    candidate_profile = MODE_TO_PROFILE[mode]
    baseline_profile = row_bundle["profiles"]["exact"]
    candidate_bundle = baseline_profile if candidate_profile is None else row_bundle["profiles"][candidate_profile]
    baseline_latency = float(baseline_profile["latency_ms"])
    candidate_latency = baseline_latency if candidate_profile is None else float(candidate_bundle["latency_ms"])
    baseline_kv_bytes = int(baseline_profile["resident_bytes"])
    candidate_kv_bytes = baseline_kv_bytes if candidate_profile is None else int(candidate_bundle["resident_bytes"])
    baseline_text = str(baseline_profile["text"])
    candidate_text = baseline_text if candidate_profile is None else str(candidate_bundle["text"])
    baseline_score = float(baseline_profile.get("qa_f1") or 0.0)
    candidate_score = baseline_score if candidate_profile is None else float(candidate_bundle.get("qa_f1") or 0.0)
    agreement = 1.0 if abs(baseline_score - candidate_score) < 1e-9 else max(0.0, 1.0 - abs(baseline_score - candidate_score))
    speedup = 1.0 if candidate_profile is None else baseline_latency / max(candidate_latency, 1e-8)
    memory_reduction = 1.0 if candidate_profile is None else baseline_kv_bytes / max(candidate_kv_bytes, 1)
    candidate_label = {
        None: "Exact reference",
        "quality": "Quality profile",
        "systems": "Systems profile",
    }[candidate_profile]
    comparison = {
        "agreement": agreement,
        "speedup": speedup,
        "memory_reduction": memory_reduction,
        "summary": build_summary_sentence(agreement, speedup, memory_reduction),
        "notes": [
            "paper_fixture",
            "benchmark_bundle:qwen_results_matrix_20260404",
            f"fixture_version:{FIXTURE_VERSION}",
            "benchmark_section:longbench_mini",
            f"prompt_id:{row_bundle['prompt_id']}",
        ],
        "paper_fixture": True,
        "paper_section": "LongBench QA mini-pack",
        "paper_headline": f"Benchmark-backed LongBench replay for {MODEL_BY_KEY[model_key].label}.",
        "paper_subtitle": (
            f"Representative LongBench QA row bundled from {row_bundle['dataset']} at {context_length} prompt tokens."
        ),
        "paper_summary": (
            f"Mean QA F1 stays at {baseline_score:.3f} while latency moves from {baseline_latency:.1f} ms/token "
            f"to {candidate_latency:.1f} ms/token."
        ),
        "paper_metric_badge": f"QA F1 {baseline_score:.3f} -> {candidate_score:.3f}",
        "baseline_label": "Exact reference",
        "candidate_label": candidate_label,
    }
    baseline = _make_payload(
        text=baseline_text,
        latency_ms=baseline_latency,
        kv_bytes=baseline_kv_bytes,
        trace=[
            {"name": "paper_fixture", "value": 1, "unit": "flag"},
            {"name": "benchmark_section", "value": "longbench_mini", "unit": "label"},
            {"name": "prompt_id", "value": str(row_bundle["prompt_id"]), "unit": "label"},
        ],
    )
    candidate = _make_payload(
        text=candidate_text,
        latency_ms=candidate_latency,
        kv_bytes=candidate_kv_bytes,
        trace=[
            {"name": "paper_fixture", "value": 1, "unit": "flag"},
            {"name": "benchmark_section", "value": "longbench_mini", "unit": "label"},
            {"name": "prompt_id", "value": str(row_bundle["prompt_id"]), "unit": "label"},
            {"name": "profile", "value": candidate_label, "unit": "label"},
        ],
    )
    return {
        "request": dict(request),
        "baseline": baseline,
        "candidate": candidate,
        "comparison": comparison,
    }


def _backend_truth_result(request: Mapping[str, Any], *, model_key: str, context_length: int, mode: str) -> dict[str, Any]:
    row_bundle = _backend_truth_bundle().get(model_key, {}).get(str(context_length))
    if row_bundle is None:
        raise ValueError(f"No backend-truth benchmark row is bundled for model={model_key} at {context_length} tokens.")
    candidate_variant = BACKEND_MODE_TO_VARIANT[mode]
    exact_stats = dict(row_bundle["profiles"]["exact"])
    if candidate_variant is None:
        candidate_stats = exact_stats
    else:
        candidate_stats = dict(row_bundle["profiles"][candidate_variant])
    baseline_latency = float(exact_stats["latency_ms"])
    candidate_latency = baseline_latency if candidate_variant is None else float(candidate_stats["latency_ms"])
    baseline_kv_bytes = int(exact_stats["resident_bytes"])
    candidate_kv_bytes = baseline_kv_bytes if candidate_variant is None else int(candidate_stats["resident_bytes"])
    exact_token_count = int(exact_stats.get("token_count") or 0)
    candidate_token_count = exact_token_count if candidate_variant is None else int(candidate_stats.get("token_count") or 0)
    baseline_text = _clean_text(str(exact_stats.get("text") or "")) or _backend_output_text(
        variant_label="Exact reference",
        prompt_length=context_length,
        generated_ids=[0] * exact_token_count,
    )
    candidate_text = (
        baseline_text
        if candidate_variant is None
        else _clean_text(str(candidate_stats.get("text") or ""))
        or _backend_output_text(
            variant_label={
                "shortlist_base": "Shortlist baseline",
                "learned_selector": "Learned selector",
            }[candidate_variant],
            prompt_length=context_length,
            generated_ids=[0] * candidate_token_count,
        )
    )
    agreement = 1.0 if candidate_variant is None else float(bool(row_bundle["output_match"][candidate_variant]))
    speedup = 1.0 if candidate_variant is None else baseline_latency / max(candidate_latency, 1e-8)
    memory_reduction = 1.0 if candidate_variant is None else baseline_kv_bytes / max(candidate_kv_bytes, 1)
    candidate_label = {
        None: "Exact reference",
        "shortlist_base": "Shortlist baseline",
        "learned_selector": "Learned selector",
    }[candidate_variant]
    comparison = {
        "agreement": agreement,
        "speedup": speedup,
        "memory_reduction": memory_reduction,
        "summary": build_summary_sentence(agreement, speedup, memory_reduction),
        "notes": [
            "paper_fixture",
            "benchmark_bundle:qwen_results_matrix_20260404",
            f"fixture_version:{FIXTURE_VERSION}",
            "benchmark_section:backend_truth",
            f"prompt_length:{context_length}",
        ],
        "paper_fixture": True,
        "paper_section": "Backend-truth decode",
        "paper_headline": f"Benchmark-backed backend-truth replay for {MODEL_BY_KEY[model_key].label}.",
        "paper_subtitle": f"Exact-length serving benchmark at {context_length} prompt tokens.",
        "paper_summary": (
            f"Decode latency moves from {baseline_latency:.1f} ms/token to {candidate_latency:.1f} ms/token "
            f"with resident KV {baseline_kv_bytes / (1024 ** 2):.1f} MiB versus {candidate_kv_bytes / (1024 ** 2):.1f} MiB. "
            f"The recorded {exact_token_count}-token decode sample is identical across exact, shortlist, and learned rows on this benchmark."
        ),
        "paper_metric_badge": (
            f"M3 fraction {float(candidate_stats.get('m3_fraction') or 0.0):.3f}, "
            f"selector {float(candidate_stats.get('selector_us') or 0.0):.1f} us"
        ),
        "baseline_label": "Exact reference",
        "candidate_label": candidate_label,
    }
    baseline = _make_payload(
        text=baseline_text,
        latency_ms=baseline_latency,
        kv_bytes=baseline_kv_bytes,
        trace=[
            {"name": "paper_fixture", "value": 1, "unit": "flag"},
            {"name": "benchmark_section", "value": "backend_truth", "unit": "label"},
            {"name": "prompt_length", "value": context_length, "unit": "tokens"},
            {"name": "decode_sample_tokens", "value": exact_token_count, "unit": "tokens"},
        ],
    )
    candidate = _make_payload(
        text=candidate_text,
        latency_ms=candidate_latency,
        kv_bytes=candidate_kv_bytes,
        trace=[
            {"name": "paper_fixture", "value": 1, "unit": "flag"},
            {"name": "benchmark_section", "value": "backend_truth", "unit": "label"},
            {"name": "prompt_length", "value": context_length, "unit": "tokens"},
            {"name": "decode_sample_tokens", "value": candidate_token_count, "unit": "tokens"},
            {"name": "profile", "value": candidate_label, "unit": "label"},
        ],
    )
    return {
        "request": dict(request),
        "baseline": baseline,
        "candidate": candidate,
        "comparison": comparison,
    }


def build_fixture_result(request: Mapping[str, Any]) -> dict[str, Any]:
    preset_key = str(request.get("preset") or "")
    model_key = str(request.get("model") or "")
    context_length = int(request.get("context_length") or 0)
    mode = str(request.get("mode") or "dense")

    if preset_key not in PRESET_BY_KEY:
        raise ValueError(f"Unsupported benchmark-backed preset: {preset_key}")
    if model_key not in MODEL_BY_KEY:
        raise ValueError(f"Unsupported benchmark-backed model: {model_key}")

    if preset_key == "compact_task":
        return _compact_task_result(request, model_key=model_key, context_length=context_length, mode=mode)
    if preset_key == "backend_truth":
        return _backend_truth_result(request, model_key=model_key, context_length=context_length, mode=mode)
    if preset_key == "longbench_mini":
        return _longbench_result(request, model_key=model_key, context_length=context_length, mode=mode)
    raise ValueError(f"Preset '{preset_key}' is not backed by the bundled benchmark data.")