"""
NovelAI sampling-preset optimizer — **optimize_fin** (classic 1-phase gen + windowed stylo).

Loads exam.txt + exam.toml, downloads cache/tokenizer.json from tokenizer_url when
missing or when the URL changes (tracked via cache/tokenizer.json.meta), generates
n_samples completions per preset via NovelAI streaming
(countdown / resume / 429 backoff), scores with multi-metric geometric-median
match and chunked embedding cosine match, then optimizes with Optuna
OptunaHub AutoSampler (picks TPE / CMA-ES / GP / … by search space).
Stops when every study has SINCE_BEST_CAP work trials without a new best
(or Ctrl+C). Work is metered per study (not shared trial counts) so peer
imports cannot stall the scheduler.

**One-phase generation (full answer after prefill):**

  prompt = system_prefix + human_prefill  (prefill_tokens)
  → generate answer_tokens AI tokens

  Scored text = full AI completion. Human reference = next answer_tokens of
  exam.txt after prefill (windowed stylo @ 256/16 + emb 32/256).

**exam.toml layout / sample count:** ``n_samples`` (default 11) is the number of
completions per trial; it is fingerprinted into the study key.

**exam.toml fixed params (optional):** set any of top_k, top_p, temperature,
logit_tolerance to a constant to pin it (not registered as Optuna params);
omit the key to suggest it as a normal search float/int. Example::

  top_p = 0.5
  logit_tolerance = 4.158883
  n_samples = 11

Per exam fingerprint (sha256 of exam.toml + exam.txt + layout) the DB holds three
parallel **single-objective** studies named ``{study_key}:{i}``:

  :0  maximize mean(stylo, 256 emb)
  :1  maximize mean(stylo, 32 emb)
  :2  maximize 32 emb

They start with the same COMPLETE trials (same params + component scores;
each study's Optuna ``value`` is that study's scalar objective). Each step:
pick the study with the lowest ``last_budget_for_new_best`` (still under
SINCE_BEST_CAP), pull peer COMPLETE trials whose params it lacks (same exam
hash only; re-map value to this study's objective), then run one new trial
on that study. Hard-stops when every study hits SINCE_BEST_CAP.

Search params: top_k, top_p, temperature, and optionally logit_tolerance.
min_p is never free: min_p = exp(-logit_tolerance / temperature) (nats).
If exam.toml sets logit_tolerance, that value is fixed; if omitted, Optuna
tunes it in ~[2.77, 5.55] (nats ≡ bits 4..8). Same math as 2**(-bits/T).
Derived min_p + logit_tolerance are trial user_attrs for reference.

Raw components (stored as trial user_attrs on every study):
  stylo:   256-token sliding windows (stride 16, same spans as 256 emb) —
           per window: text-metric vector (stylo.py); per window index,
           geometric median across samples → scale-normalized L2 match vs human
           window; stylo = median over windows. Metric extracts fire live as
           each full 256-tok window completes (overlap generation).
  32 emb:  32-token sliding windows (stride 16) — per window index, geometric
           median of sample embeddings → cosine vs human; median over windows
  256 emb: same with 256-token windows, stride 16

Embeddings are requested incrementally while each sample streams: completed
32-tok and 256-tok windows are queued (mixed in one buffer) and flushed as
OpenRouter embed batches of up to EMBED_INPUT_MAX (1024) inputs per HTTP call.
Finalize flushes any remainder. Stylo metric extracts still fire per 256-tok
window (local CPU). By generation end, scores are mostly ready (only a
possible final covering window + in-flight batch futures remain).
"""
from __future__ import annotations

import gc
import hashlib
import json
import math
import os
import sys
import threading
import time
import tomllib
import traceback
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Any

# Must run before numpy / spaCy (via stylo) pull in multi-threaded BLAS.
for _k, _v in (
    ("OMP_NUM_THREADS", "1"),
    ("OPENBLAS_NUM_THREADS", "1"),
    ("MKL_NUM_THREADS", "1"),
    ("NUMEXPR_NUM_THREADS", "1"),
    ("VECLIB_MAXIMUM_THREADS", "1"),
    ("BLIS_NUM_THREADS", "1"),
):
    os.environ.setdefault(_k, _v)

import numpy as np
import optuna
import requests
from dotenv import load_dotenv
from openai import OpenAI
import optunahub
from optuna.distributions import FloatDistribution, IntDistribution
from optuna.samplers import BaseSampler
from optuna.trial import TrialState, create_trial
from tokenizers import Tokenizer

from stylo import (
    FEATURE_NAMES,
    extract_vector,
    extract_vectors,
    metric_pool_info,
    reap_stale_metric_workers,
    shutdown_metric_pool,
    submit_extract_vector,
    warmup_metric_workers,
)

# HuggingFace tokenizers is NOT thread-safe for concurrent encode/decode on one
# instance. Live streams + embed batches + stylo all call decode in parallel —
# without a lock that often SIGSEGV (exit 139).
_TOK_LOCK = threading.RLock()

ROOT = Path(__file__).resolve().parent
NAI_URL = "https://text.novelai.net/oa/v1/completions"
EMBED_MODEL = "qwen/qwen3-embedding-8b"
EMBED_DIM = 4096  # max for qwen3-embedding-8b (not 4100)
# Chunk sizes for embedding objectives (tokens, NovelAI tokenizer).
# Sliding windows with fixed stride (same for both sizes).
CHUNK_32 = 32
CHUNK_32_STRIDE = 16  # 50% overlap on 32-tok windows
CHUNK_256 = 256
CHUNK_256_STRIDE = 16  # dense 256-tok windows
# Stylo uses the same token windows as 256 emb (size + stride + span rules).
STYLO_CHUNK = CHUNK_256
STYLO_STRIDE = CHUNK_256_STRIDE
# Default completions per preset; overridden by exam.toml ``n_samples`` at load.
DEFAULT_N_SAMPLES = 11
N_SAMPLES = DEFAULT_N_SAMPLES
BASE_DELAY = 1.0
MAX_DELAY = 128.0
# Bump when stylo metrics / chunk embedding / stylo-window semantics change.
METRICS_VERSION = "v7-fullgen-stylo256-stride16"
# NovelAI GLM/Xialong rejects prompt token ids >= this (HTTP 400
# "token beyond max token for GLM: 151524"). Exclusive upper bound.
GLM_MAX_TOKEN_ID = 151524
# Practical top_k ceiling (must stay well under GLM_MAX_TOKEN_ID so Optuna
# never samples a value that looks like an OOB token id).
TOP_K_MAX = 25000
# Xialong/GLM on NovelAI website context (prompt + max_tokens budget).
NAI_CONTEXT_LIMIT = 28672
# Soft-cut / empty loops are prevented by min_p logic; no HTTP attempt cap.
MAX_HTTP_ATTEMPTS = None  # unlimited stream POSTs per trial (None = no cap)
MAX_CONCURRENT_STREAMS = 6  # NovelAI stream POSTs in flight at once (≤ N_SAMPLES)
MAX_TRIAL_SECONDS = 3 * 60 * 60  # wall-clock budget per generate_five
# Per-study work budget (optimize steps) without a new best.
# Scheduler always works the study with the lowest budget still under the cap.
# Hard stop when every study hits SINCE_BEST_CAP without a new best.
SINCE_BEST_CAP = 200

# AutoSampler (optunahub samplers/auto_sampler) chooses an internal algorithm
# from the search space / trial history. SEARCH_SPACE bounds must match the
# objective suggest_* calls. No fixed seed — a seed makes every restart
# re-suggest the same next point after FAIL/RUNNING purge.
#
# Free unless pinned in exam.toml: top_k, top_p, temperature, logit_tolerance.
# min_p is never free: min_p = exp(-logit_tolerance / temperature).
# Same curve as 2**(-bits/T) with bits = logit_tolerance / ln(2).
SEARCH_SPACE: dict[str, IntDistribution | FloatDistribution] = {
    "top_k": IntDistribution(6, TOP_K_MAX, log=True),
    "top_p": FloatDistribution(0.05, 1.0),
    "temperature": FloatDistribution(0.8, 10.0, log=True),
    # Only suggested when exam.toml omits logit_tolerance.
    # Range ≡ former bits [4, 8] as nats: 4*ln2 .. 8*ln2.
    "logit_tolerance": FloatDistribution(
        4.0 * math.log(2.0), 8.0 * math.log(2.0)
    ),
}
# Default: 6*ln(2) ≈ 4.158883 nats ≡ bits=6 ≡ 2**(-6/T) (CSV ridge).
LOGIT_TOLERANCE_DEFAULT = 6.0 * math.log(2.0)
MIN_P_MAX = 0.95
PARAM_LOGIT_TOLERANCE = "logit_tolerance"


def min_p_from_temperature(
    temperature: float,
    logit_tolerance: float | None = None,
) -> float:
    """min_p = exp(-logit_tolerance / temperature), clipped to [0, MIN_P_MAX].

    ``logit_tolerance`` is temperature-invariant headroom in nats under the max
    logit. Default 6*ln(2) → same as 2**(-6/T).
    """
    t = float(temperature)
    if t <= 0.0:
        raise ValueError(f"temperature must be > 0, got {t}")
    tol = (
        float(logit_tolerance)
        if logit_tolerance is not None
        else LOGIT_TOLERANCE_DEFAULT
    )
    if tol <= 0.0:
        raise ValueError(f"logit_tolerance must be > 0, got {tol}")
    return float(min(MIN_P_MAX, math.exp(-tol / t)))

# Three parallel single-objective studies per exam fingerprint (toml+txt hash).
# Same starting trials (params + component scores); each study maximizes one
# scalar. Work goes to the study with the lowest per-study work budget since
# its best improved (round-robin starvation avoidance). Before each trial,
# that study pulls COMPLETE trials with unseen params from sibling studies
# only (never across a different exam hash).
N_STUDY_SHARDS = 3
SHARD_SEP = ":"
# trial.system_attrs — COMPLETE trial copied from a peer study.
SYS_IMPORTED_FROM = "imported_from"
# study.user_attrs — which scalar this study maximizes (0, 1, or 2).
ATTR_OBJECTIVE_INDEX = "objective_index"
# study.user_attrs — optimize steps since last new best (own work only).
ATTR_BUDGET_SINCE_BEST = "last_budget_for_new_best"

optuna.logging.set_verbosity(optuna.logging.WARNING)


class LineProgress:
    """Single-line terminal progress (\\r updates). Thread-safe."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._width = 0
        self._active = False

    def update(self, msg: str) -> None:
        with self._lock:
            try:
                pad = max(0, self._width - len(msg))
                sys.stdout.write("\r" + msg + (" " * pad))
                sys.stdout.flush()
                self._width = len(msg)
                self._active = True
            except Exception:
                pass

    def end(self, msg: str | None = None) -> None:
        """Finish the line: optional final message, always ends with newline."""
        with self._lock:
            try:
                if msg is not None:
                    pad = max(0, self._width - len(msg))
                    sys.stdout.write("\r" + msg + (" " * pad) + "\n")
                elif self._active:
                    sys.stdout.write("\n")
                sys.stdout.flush()
            except Exception:
                pass
            self._width = 0
            self._active = False


def _valid_token_id(tid: int) -> bool:
    return 0 <= int(tid) < GLM_MAX_TOKEN_ID


def _sanitize_token_ids(ids: list[Any] | None) -> tuple[list[int], list[Any]]:
    """Split into (kept in-range ids, dropped invalid/OOB values)."""
    if not ids:
        return [], []
    kept: list[int] = []
    dropped: list[Any] = []
    for t in ids:
        try:
            ti = int(t)
        except (TypeError, ValueError):
            dropped.append(t)
            continue
        if _valid_token_id(ti):
            kept.append(ti)
        else:
            dropped.append(ti)
    return kept, dropped


def _finite(x: float, default: float = 0.0) -> float:
    v = float(x)
    return v if np.isfinite(v) else default


# ---------------------------------------------------------------------------
# config / exam load
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class ExamConfig:
    model: str
    token_bias: dict[str, float]
    system_prefix: list[int]
    prefill_tokens: int
    answer_tokens: int
    tokenizer_url: str
    # Completions per preset trial (from exam.toml n_samples).
    n_samples: int
    # Nats in min_p = exp(-logit_tolerance / T). None → Optuna tunes it.
    logit_tolerance: float | None
    # Optional fixed sampling params from exam.toml (None → Optuna suggests).
    fixed_top_k: int | None = None
    fixed_top_p: float | None = None
    fixed_temperature: float | None = None


@dataclass(frozen=True)
class ExamBundle:
    config: ExamConfig
    prefill_ids: list[int]
    prompt_ids: list[int]  # system_prefix + prefill
    answer_text: str  # human reference = next answer_tokens after prefill
    toml_hash: str
    exam_hash: str
    study_key: str  # sha256(toml || exam || layout) — Optuna study name

    @property
    def total_answer_tokens(self) -> int:
        """AI completion length (= answer_tokens)."""
        return int(self.config.answer_tokens)


def make_study_key(toml_hash: str, exam_hash: str, cfg: "ExamConfig") -> str:
    """Study identity: exam.toml + exam.txt + n_samples + full-gen layout.

    Fixed params from exam.toml are fingerprinted so pinning top_p etc. never
    mixes Optuna history with a free-parameter run. Human cache is keyed by
    this study_key so a layout change triggers a recompute.
    """
    fixed_fp = (
        f"fk={cfg.fixed_top_k}"
        f":fp={cfg.fixed_top_p}"
        f":fT={cfg.fixed_temperature}"
        f":ftol={cfg.logit_tolerance}"
    )
    return hashlib.sha256(
        f"{toml_hash}:{exam_hash}:n_samples={cfg.n_samples}"
        f":mode=fullgen:answer={cfg.answer_tokens}"
        f":pin=nosuggest"  # fixed knobs not in trial.params (AutoSampler-safe)
        f":{fixed_fp}"
        f":chunk32={CHUNK_32}/{CHUNK_32_STRIDE}"
        f":chunk256={CHUNK_256}/{CHUNK_256_STRIDE}"
        f":stylo={STYLO_CHUNK}/{STYLO_STRIDE}".encode("utf-8")
    ).hexdigest()


def load_exam_toml(path: Path) -> tuple[ExamConfig, str]:
    raw = path.read_bytes()
    toml_hash = hashlib.sha256(raw).hexdigest()
    data = tomllib.loads(raw.decode("utf-8"))
    prefix = [int(x) for x in data["system_prefix"]]
    bias_raw = data.get("token_bias") or {}
    bias = {str(k): float(v) for k, v in bias_raw.items()}
    url = str(data.get("tokenizer_url") or "").strip()
    if not url:
        raise SystemExit("exam.toml missing tokenizer_url")
    oob_prefix = [t for t in prefix if not _valid_token_id(t)]
    if oob_prefix:
        raise SystemExit(
            f"system_prefix has token id(s) >= {GLM_MAX_TOKEN_ID}: {oob_prefix[:8]}"
        )
    oob_bias = [k for k in bias if not _valid_token_id(int(k))]
    if oob_bias:
        raise SystemExit(
            f"token_bias has key(s) >= {GLM_MAX_TOKEN_ID}: {oob_bias[:8]}"
        )
    # Optional fixed params. Absent key → Optuna searches SEARCH_SPACE.
    # logit_tolerance (nats). Legacy: min_p_constant, min_p_bits → nats.
    raw_tol = data.get(PARAM_LOGIT_TOLERANCE, None)
    if raw_tol is None and "min_p_constant" in data:
        raw_tol = data["min_p_constant"]
        print(
            f"exam.toml: legacy min_p_constant → {PARAM_LOGIT_TOLERANCE}"
        )
    if raw_tol is None and "min_p_bits" in data:
        try:
            bits = float(data["min_p_bits"])
            raw_tol = bits * math.log(2.0)
            print(
                f"exam.toml: legacy min_p_bits={bits:g} → "
                f"{PARAM_LOGIT_TOLERANCE}={float(raw_tol):g} nats"
            )
        except (TypeError, ValueError):
            raw_tol = None
    if raw_tol is None:
        logit_tolerance: float | None = None
    else:
        try:
            logit_tolerance = float(raw_tol)
        except (TypeError, ValueError) as e:
            raise SystemExit(
                f"exam.toml {PARAM_LOGIT_TOLERANCE} must be a positive float "
                f"or omitted, got {raw_tol!r}"
            ) from e
        if not math.isfinite(logit_tolerance) or logit_tolerance <= 0.0:
            raise SystemExit(
                f"exam.toml {PARAM_LOGIT_TOLERANCE} must be finite and > 0, "
                f"got {logit_tolerance}"
            )

    def _opt_float(key: str) -> float | None:
        if key not in data:
            return None
        try:
            v = float(data[key])
        except (TypeError, ValueError) as e:
            raise SystemExit(
                f"exam.toml {key} must be a float or omitted, got {data[key]!r}"
            ) from e
        if not math.isfinite(v):
            raise SystemExit(f"exam.toml {key} must be finite, got {v}")
        return v

    def _opt_int(key: str) -> int | None:
        if key not in data:
            return None
        try:
            v = int(data[key])
        except (TypeError, ValueError) as e:
            raise SystemExit(
                f"exam.toml {key} must be an int or omitted, got {data[key]!r}"
            ) from e
        return v

    fixed_top_k = _opt_int("top_k")
    fixed_top_p = _opt_float("top_p")
    fixed_temperature = _opt_float("temperature")

    # n_samples: required layout knob (default matches historical constant).
    if "n_samples" in data:
        try:
            n_samples = int(data["n_samples"])
        except (TypeError, ValueError) as e:
            raise SystemExit(
                f"exam.toml n_samples must be an int >= 1, got {data['n_samples']!r}"
            ) from e
    else:
        n_samples = DEFAULT_N_SAMPLES
    if n_samples < 1:
        raise SystemExit(f"exam.toml n_samples must be >= 1, got {n_samples}")

    # Validate fixed values sit inside SEARCH_SPACE (when present).
    if fixed_top_k is not None:
        d = SEARCH_SPACE["top_k"]
        if not (d.low <= fixed_top_k <= d.high):
            raise SystemExit(
                f"exam.toml top_k={fixed_top_k} outside SEARCH_SPACE "
                f"[{d.low}, {d.high}]"
            )
    if fixed_top_p is not None:
        d = SEARCH_SPACE["top_p"]
        if not (d.low <= fixed_top_p <= d.high):
            raise SystemExit(
                f"exam.toml top_p={fixed_top_p} outside SEARCH_SPACE "
                f"[{d.low}, {d.high}]"
            )
    if fixed_temperature is not None:
        d = SEARCH_SPACE["temperature"]
        if not (d.low <= fixed_temperature <= d.high):
            raise SystemExit(
                f"exam.toml temperature={fixed_temperature} outside "
                f"SEARCH_SPACE [{d.low}, {d.high}]"
            )
    if logit_tolerance is not None:
        d = SEARCH_SPACE[PARAM_LOGIT_TOLERANCE]
        # Fixed tol may be outside the free-search band (e.g. classic 6*ln2);
        # only warn if wildly non-positive already checked.
        if not (d.low <= logit_tolerance <= d.high):
            print(
                f"exam.toml: {PARAM_LOGIT_TOLERANCE}={logit_tolerance:g} "
                f"outside free-search band [{d.low:g}, {d.high:g}] "
                f"(ok when pinned)"
            )

    cfg = ExamConfig(
        model=str(data["model"]),
        token_bias=bias,
        system_prefix=prefix,
        prefill_tokens=int(data["prefill_tokens"]),
        answer_tokens=int(data["answer_tokens"]),
        tokenizer_url=url,
        n_samples=n_samples,
        logit_tolerance=logit_tolerance,
        fixed_top_k=fixed_top_k,
        fixed_top_p=fixed_top_p,
        fixed_temperature=fixed_temperature,
    )
    if cfg.prefill_tokens < 1 or cfg.answer_tokens < 1:
        raise SystemExit("prefill_tokens and answer_tokens must be >= 1")
    # Keep module-level N_SAMPLES in sync so generation/scoring paths use
    # the exam.toml value without threading n_samples through every call.
    global N_SAMPLES
    N_SAMPLES = cfg.n_samples
    pinned = []
    if cfg.fixed_top_k is not None:
        pinned.append(f"top_k={cfg.fixed_top_k}")
    if cfg.fixed_top_p is not None:
        pinned.append(f"top_p={cfg.fixed_top_p}")
    if cfg.fixed_temperature is not None:
        pinned.append(f"temperature={cfg.fixed_temperature}")
    if cfg.logit_tolerance is not None:
        pinned.append(f"logit_tolerance={cfg.logit_tolerance:g}")
    if pinned:
        print(f"exam.toml pinned: {', '.join(pinned)}")
    print(f"exam.toml n_samples={cfg.n_samples}")
    return cfg, toml_hash


def _tokenizer_meta_path(path: Path) -> Path:
    """Sidecar recording which URL produced tokenizer.json."""
    return path.with_name(path.name + ".meta")


def _read_tokenizer_meta(path: Path) -> dict[str, str] | None:
    meta_path = _tokenizer_meta_path(path)
    if not meta_path.is_file():
        return None
    try:
        data = json.loads(meta_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    if not isinstance(data, dict):
        return None
    url = data.get("url")
    if not isinstance(url, str) or not url.strip():
        return None
    out = {"url": url.strip()}
    sha = data.get("sha256")
    if isinstance(sha, str) and sha:
        out["sha256"] = sha
    return out


def _write_tokenizer_meta(path: Path, *, url: str, sha256: str) -> None:
    meta = {"url": url.strip(), "sha256": sha256}
    meta_path = _tokenizer_meta_path(path)
    tmp = meta_path.with_suffix(meta_path.suffix + ".part")
    tmp.write_text(json.dumps(meta, indent=2) + "\n", encoding="utf-8")
    tmp.replace(meta_path)


def _tokenizer_cache_ok(path: Path, url: str) -> bool:
    """True if tokenizer.json exists and was downloaded from this exact url."""
    if not path.is_file():
        return False
    meta = _read_tokenizer_meta(path)
    if meta is None:
        return False
    if meta["url"] != url.strip():
        return False
    # Optional integrity: if we recorded sha256, file must still match.
    want = meta.get("sha256")
    if want:
        got = hashlib.sha256(path.read_bytes()).hexdigest()
        if got != want:
            return False
    return True


def ensure_tokenizer(path: Path, url: str) -> None:
    """
    Ensure tokenizer.json matches tokenizer_url.

    Downloads when missing, when sidecar meta is missing/corrupt, when the
    stored URL differs from exam.toml, or when the file hash no longer matches
    the sidecar. Does not re-fetch on every run if the URL is unchanged.
    """
    url = url.strip()
    if not url:
        raise SystemExit("tokenizer_url is empty")
    path.parent.mkdir(parents=True, exist_ok=True)
    if _tokenizer_cache_ok(path, url):
        return

    meta = _read_tokenizer_meta(path)
    if path.is_file() and meta is None:
        print(
            f"Tokenizer {path.name} has no .meta sidecar — "
            f"re-downloading to bind it to tokenizer_url"
        )
    elif path.is_file() and meta is not None and meta["url"] != url:
        print(
            f"tokenizer_url changed\n"
            f"  was: {meta['url']}\n"
            f"  now: {url}\n"
            f"Re-downloading {path.name} …"
        )
    else:
        print(f"Downloading tokenizer from {url} …")

    try:
        with requests.get(url, stream=True, timeout=300) as resp:
            resp.raise_for_status()
            tmp = path.with_suffix(path.suffix + ".part")
            h = hashlib.sha256()
            with open(tmp, "wb") as f:
                for chunk in resp.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        h.update(chunk)
                        f.write(chunk)
            tmp.replace(path)
            sha = h.hexdigest()
    except requests.RequestException as e:
        raise SystemExit(f"failed to download tokenizer: {e}") from e

    if not path.is_file() or path.stat().st_size < 100:
        raise SystemExit(f"tokenizer download looks empty/corrupt: {path}")
    try:
        Tokenizer.from_file(str(path))
    except Exception as e:
        path.unlink(missing_ok=True)
        meta_p = _tokenizer_meta_path(path)
        if meta_p.is_file():
            meta_p.unlink(missing_ok=True)
        raise SystemExit(f"tokenizer.json failed to load after download: {e}") from e

    _write_tokenizer_meta(path, url=url, sha256=sha)
    print(f"Saved {path.name} + {_tokenizer_meta_path(path).name} (sha256={sha[:12]}…)")


def load_tokenizer(path: Path) -> Tokenizer:
    if not path.is_file():
        raise SystemExit(f"tokenizer not found: {path}")
    try:
        return Tokenizer.from_file(str(path))
    except Exception as e:
        raise SystemExit(f"failed to load tokenizer {path}: {e}") from e


def build_exam(
    cfg: ExamConfig,
    exam_txt: Path,
    tok: Tokenizer,
    *,
    toml_hash: str,
) -> ExamBundle:
    """Classic 1-phase layout: prompt = system + prefill; gen answer_tokens.

    Human reference for scoring is the next answer_tokens of exam.txt after
    the prefill slice (windowed stylo + emb).
    """
    raw = exam_txt.read_bytes()
    exam_hash = hashlib.sha256(raw).hexdigest()
    text = raw.decode("utf-8")
    if not text:
        raise SystemExit("exam.txt is empty")
    ids = _encoding_ids(tok, text)
    need = cfg.prefill_tokens + cfg.answer_tokens
    if len(ids) < need:
        raise SystemExit(
            f"exam.txt has {len(ids)} tokens; need prefill+answer={need} "
            f"(prefill={cfg.prefill_tokens}, answer_tokens={cfg.answer_tokens})"
        )
    prefill = ids[: cfg.prefill_tokens]
    answer_ids = ids[cfg.prefill_tokens : need]
    answer_text = _tok_decode(tok, answer_ids)
    prompt = list(cfg.system_prefix) + list(prefill)
    bad = [t for t in prompt if not _valid_token_id(t)]
    if bad:
        raise SystemExit(
            f"exam prompt has {len(bad)} token id(s) >= {GLM_MAX_TOKEN_ID} "
            f"(NovelAI GLM max); first few: {bad[:8]}"
        )
    study_key = make_study_key(toml_hash, exam_hash, cfg)
    ctx_need = len(prompt) + cfg.answer_tokens
    if ctx_need > NAI_CONTEXT_LIMIT:
        raise SystemExit(
            f"exam context budget {ctx_need} exceeds NovelAI limit "
            f"{NAI_CONTEXT_LIMIT} (prompt={len(prompt)} + answer={cfg.answer_tokens})"
        )
    return ExamBundle(
        config=cfg,
        prefill_ids=prefill,
        prompt_ids=prompt,
        answer_text=answer_text,
        toml_hash=toml_hash,
        exam_hash=exam_hash,
        study_key=study_key,
    )


def _sanitize_metric_row(v: np.ndarray) -> np.ndarray:
    """Finite + hard cap so one bad library value cannot nuke scaled L2."""
    v = np.asarray(v, dtype=float)
    if v.shape != (len(FEATURE_NAMES),):
        raise RuntimeError(
            f"metric vector shape {v.shape}, expected ({len(FEATURE_NAMES)},)"
        )
    v = np.where(np.isfinite(v), v, 0.0)
    return np.clip(v, -1e6, 1e6)


def _sanitize_metric_matrix(mat: np.ndarray) -> np.ndarray:
    """Stack of window metric rows → finite, clipped (n_windows, n_feat)."""
    mat = np.asarray(mat, dtype=float)
    n_feat = len(FEATURE_NAMES)
    if mat.ndim != 2 or mat.shape[1] != n_feat:
        raise RuntimeError(
            f"metric matrix shape {mat.shape}, expected (n, {n_feat})"
        )
    if mat.shape[0] == 0:
        return np.zeros((0, n_feat), dtype=float)
    return np.stack([_sanitize_metric_row(mat[i]) for i in range(mat.shape[0])], axis=0)


def _metric_vector(text: str) -> np.ndarray:
    """Fixed-order fiction metrics for one window/string; non-finite → 0.

    Empty/whitespace → zeros (no text to score). Extract failures **raise** —
    never silent zero-fill of real windows (that poisons stylo scores).
    """
    t = (text or "").strip()
    if not t:
        return np.zeros(len(FEATURE_NAMES), dtype=float)
    try:
        v = np.asarray(extract_vector(t), dtype=float)
    except Exception as e:
        raise RuntimeError(
            f"metric extract failed ({type(e).__name__}: {e})"
        ) from e
    return _sanitize_metric_row(v)


def _metric_matrix(texts: list[str]) -> np.ndarray:
    """Batch metrics for several window texts → (n, n_feat).

    Empty/whitespace windows stay zero rows. Non-empty extract failures raise
    (no silent zero-fill).
    """
    n_feat = len(FEATURE_NAMES)
    if not texts:
        return np.zeros((0, n_feat), dtype=float)
    cleaned = [(t if isinstance(t, str) else str(t)) for t in texts]
    # Empty / whitespace windows → zero row (do not abort the whole batch).
    nonempty_idx = [i for i, t in enumerate(cleaned) if t.strip()]
    out = np.zeros((len(cleaned), n_feat), dtype=float)
    if not nonempty_idx:
        return out
    batch = [cleaned[i] for i in nonempty_idx]
    try:
        mat = np.asarray(extract_vectors(batch), dtype=float)
    except Exception as e:
        # Fall back per-window so one bad string can surface a clear error
        # without losing which text failed (extract_vectors is all-or-nothing).
        try:
            for i in nonempty_idx:
                out[i] = _metric_vector(cleaned[i])
            return out
        except Exception:
            raise RuntimeError(
                f"metric extract batch failed ({type(e).__name__}: {e})"
            ) from e
    if mat.ndim != 2 or mat.shape[0] != len(batch):
        raise RuntimeError(
            f"metric matrix shape {mat.shape}, expected ({len(batch)}, n_feat)"
        )
    for row, i in enumerate(nonempty_idx):
        out[i] = _sanitize_metric_row(mat[row])
    return out


def _submit_window_metrics(text: str) -> Future[np.ndarray]:
    """Queue one 256-tok window metric extract on the shared stylo pool."""
    return submit_extract_vector(text, n_texts_hint=N_SAMPLES)


def metric_token_chunks(
    text: str,
    tok: Tokenizer,
    *,
    chunk_size: int = STYLO_CHUNK,
    stride: int | None = None,
) -> np.ndarray:
    """Text-metric matrix for sliding token windows → (n_windows, n_feat)."""
    chunks = token_chunks(text, tok, chunk_size, stride=stride)
    if not chunks:
        return np.zeros((0, len(FEATURE_NAMES)), dtype=float)
    return _metric_matrix(chunks)


def _metric_scale(human_metrics: np.ndarray) -> np.ndarray:
    """Global per-feature scale from human window matrix (or single vector).

    scale_j = max(max_w |human[w, j]|, 1) so all window comparisons share one
    scale and count-like dims do not dominate L2 / geometric median.
    """
    h = np.asarray(human_metrics, dtype=float)
    if h.ndim == 1:
        return np.maximum(np.abs(h), 1.0)
    if h.ndim != 2 or h.shape[0] == 0:
        return np.ones(len(FEATURE_NAMES), dtype=float)
    return np.maximum(np.max(np.abs(h), axis=0), 1.0)


def _scaled_metric_distance(
    a: np.ndarray, b: np.ndarray, scale: np.ndarray
) -> float:
    return _finite(float(np.linalg.norm((a - b) / scale)))


# ---------------------------------------------------------------------------
# geometric median (Weiszfeld / Vardi–Zhang) — numpy only
# ---------------------------------------------------------------------------

def geometric_median(X: np.ndarray, *, tol: float = 1e-8, max_iter: int = 500) -> np.ndarray:
    X = np.asarray(X, dtype=float)
    if X.ndim != 2:
        raise ValueError("X must be 2-D")
    n = X.shape[0]
    if n == 0:
        raise ValueError("empty")
    if n == 1:
        return X[0].copy()
    if n == 2:
        return 0.5 * (X[0] + X[1])

    y = np.median(X, axis=0).astype(float, copy=True)
    scale = float(np.mean(np.linalg.norm(X - y, axis=1)))
    if not np.isfinite(scale) or scale <= 0.0:
        return y
    eps = tol * max(scale, 1.0)

    for _ in range(max_iter):
        diffs = X - y
        dists = np.linalg.norm(diffs, axis=1)
        at = dists <= eps
        n_at = int(np.count_nonzero(at))
        far = ~at
        if not np.any(far):
            return y
        inv = 1.0 / dists[far]
        inv_sum = float(np.sum(inv))
        if inv_sum <= 0.0 or not np.isfinite(inv_sum):
            return y
        T = (X[far] * inv[:, None]).sum(axis=0) / inv_sum
        if n_at == 0:
            y_new = T
        else:
            r_vec = (diffs[far] * inv[:, None]).sum(axis=0)
            r = float(np.linalg.norm(r_vec))
            if not np.isfinite(r) or r <= float(n_at):
                return y
            alpha = float(n_at) / r
            y_new = (1.0 - alpha) * T + alpha * y
        if float(np.linalg.norm(y_new - y)) <= eps:
            return y_new
        y = y_new
    return y


# ---------------------------------------------------------------------------
# embeddings (OpenRouter) + token chunking
# ---------------------------------------------------------------------------

def make_embed_client() -> OpenAI:
    key = os.getenv("OPENROUTER_API_KEY")
    if not key:
        raise SystemExit("OPENROUTER_API_KEY missing in .env")
    # Per-request timeout. Small (1-window) calls finish in seconds; mega-batches
    # need longer — raise EMBED_HTTP_TIMEOUT_S if you use EMBED_FLUSH_SIZE=1024.
    timeout_s = float(os.getenv("EMBED_HTTP_TIMEOUT_S") or "60")
    return OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=key,
        timeout=timeout_s,
        max_retries=0,  # we retry ourselves in _embed_texts_one_batch
    )


# Per-thread clients: sharing one OpenAI/httpx client across many embed-pool
# threads has been observed to hang with no request visible on OpenRouter.
_embed_tls = threading.local()


def _thread_local_embed_client(_ignored: OpenAI | None = None) -> OpenAI:
    """Return an OpenAI client local to this worker thread."""
    c = getattr(_embed_tls, "client", None)
    if c is None:
        c = make_embed_client()
        _embed_tls.client = c
    return c


# OpenRouter API max inputs per embeddings.create (hard limit).
EMBED_INPUT_MAX = 1024
# How many *windows* we pack into one live HTTP call before flush.
# Default **1** = one request per window (small payloads; historically more
# reliable than 1024×256-tok mega-batches ~200k+ tokens). API still allows
# up to EMBED_INPUT_MAX. Override: EMBED_FLUSH_SIZE=32|256|1024
def _embed_flush_size() -> int:
    raw = (os.getenv("EMBED_FLUSH_SIZE") or "64").strip()
    try:
        n = int(raw)
    except ValueError:
        n = 64
    return max(1, min(n, EMBED_INPUT_MAX))


EMBED_FLUSH_SIZE = _embed_flush_size()


# Concurrent embed HTTP calls (live pool + human parallel path).
# Override: EMBED_BATCH_PARALLEL=N / EMBED_POOL_WORKERS=N
def _embed_batch_parallel() -> int:
    raw = (os.getenv("EMBED_BATCH_PARALLEL") or "").strip()
    if raw:
        try:
            return max(1, int(raw))
        except ValueError:
            pass
    return 1024


def embed_texts(
    client: OpenAI,
    texts: list[str],
    *,
    parallel: bool = True,
    inter_batch_delay_s: float = 0.0,
    batch_size: int | None = None,
) -> np.ndarray:
    """Embed texts; force max dimensions (4096). Empty → (0, d).

    Splits into batches of ``batch_size`` (default EMBED_INPUT_MAX=1024) per
    HTTP call.

    ``parallel=True`` (default, live/general): multiple batches may run at once.

    ``parallel=False`` + ``inter_batch_delay_s``: one request at a time, sleep
    between batches (used for **human cache** so OpenRouter is not flooded).
    """
    if not texts:
        return np.zeros((0, EMBED_DIM), dtype=float)
    # OpenRouter / gateways occasionally hand back null/missing data; never
    # let that become a bare TypeError: 'NoneType' object is not iterable.
    clean_texts = [
        t if isinstance(t, str) else ("" if t is None else str(t)) for t in texts
    ]
    max_n = int(batch_size) if batch_size is not None else EMBED_INPUT_MAX
    max_n = max(1, min(max_n, EMBED_INPUT_MAX))
    if len(clean_texts) <= max_n:
        return _embed_texts_one_batch(client, clean_texts)

    batches = [
        clean_texts[start : start + max_n]
        for start in range(0, len(clean_texts), max_n)
    ]
    delay = max(0.0, float(inter_batch_delay_s))
    use_parallel = bool(parallel) and delay <= 0.0
    workers = min(len(batches), _embed_batch_parallel()) if use_parallel else 1

    if workers <= 1:
        parts: list[np.ndarray] = []
        for i, b in enumerate(batches):
            if i > 0 and delay > 0.0:
                time.sleep(delay)
            parts.append(_embed_texts_one_batch(client, b))
    else:
        with ThreadPoolExecutor(
            max_workers=workers, thread_name_prefix="embed-batch"
        ) as pool:
            # map preserves batch order.
            parts = list(
                pool.map(
                    lambda b: _embed_texts_one_batch(client, b),
                    batches,
                )
            )
    arr = np.vstack(parts)
    if arr.shape != (len(clean_texts), EMBED_DIM):
        raise RuntimeError(
            f"expected embedding shape ({len(clean_texts)}, {EMBED_DIM}), got {arr.shape}"
        )
    return arr


# OpenRouter occasionally returns data=None / empty / 429 / timeouts under load.
# No attempt cap: without embeds the trial cannot score, so wait it out.
_EMBED_RETRY_BASE_S = float(os.getenv("EMBED_RETRY_BASE_S") or "1.5")


def _embed_texts_one_batch(
    client: OpenAI | None, clean_texts: list[str]
) -> np.ndarray:
    """Single embeddings.create call for ≤ EMBED_INPUT_MAX texts.

    Retries forever on transient API errors (timeout, 429, empty data, …)
    with exponential backoff. Always uses a **thread-local** OpenAI client
    (ignores shared ``client`` for the HTTP call) so concurrent pool workers
    do not hang on one httpx session while OpenRouter shows zero active
    requests.
    """
    if not clean_texts:
        return np.zeros((0, EMBED_DIM), dtype=float)
    if len(clean_texts) > EMBED_INPUT_MAX:
        raise RuntimeError(
            f"internal: batch size {len(clean_texts)} > EMBED_INPUT_MAX={EMBED_INPUT_MAX}"
        )

    cl = _thread_local_embed_client(client)
    attempt = 0
    while True:
        attempt += 1
        try:
            resp = cl.embeddings.create(
                model=EMBED_MODEL,
                input=clean_texts,
                dimensions=EMBED_DIM,
                encoding_format="float",
            )
            raw = getattr(resp, "data", None)
            if raw is None:
                # Some gateways put error text on the response object.
                err_bits = []
                for attr in ("error", "message", "model", "id"):
                    v = getattr(resp, attr, None)
                    if v is not None:
                        err_bits.append(f"{attr}={v!r}")
                # Also try model_dump / dict if present.
                dump = None
                for meth in ("model_dump", "to_dict", "dict"):
                    fn = getattr(resp, meth, None)
                    if callable(fn):
                        try:
                            dump = fn()
                            break
                        except Exception:
                            pass
                if dump is not None:
                    err_bits.append(f"body={str(dump)[:300]}")
                raise RuntimeError(
                    "embedding API returned data=None"
                    + (f" ({', '.join(err_bits)})" if err_bits else "")
                    + f" n={len(clean_texts)} attempt={attempt}"
                )
            try:
                data = sorted(raw, key=lambda d: getattr(d, "index", 0))
            except TypeError as e:
                raise RuntimeError(
                    f"embedding API returned non-iterable data: {type(raw).__name__}"
                ) from e
            if len(data) != len(clean_texts):
                raise RuntimeError(
                    f"embedding count mismatch: got {len(data)}, "
                    f"expected {len(clean_texts)}"
                )
            vectors: list[list[float]] = []
            for i, d in enumerate(data):
                emb = getattr(d, "embedding", None)
                if emb is None:
                    raise RuntimeError(
                        f"embedding API item {i} has embedding=None"
                    )
                try:
                    vec = [float(x) for x in emb]
                except TypeError as e:
                    raise RuntimeError(
                        f"embedding API item {i} embedding not iterable: "
                        f"{type(emb).__name__}"
                    ) from e
                vectors.append(vec)
            arr = np.asarray(vectors, dtype=float)
            if arr.ndim != 2 or arr.shape != (len(clean_texts), EMBED_DIM):
                raise RuntimeError(
                    f"expected embedding shape ({len(clean_texts)}, {EMBED_DIM}), "
                    f"got {arr.shape}"
                )
            if not np.all(np.isfinite(arr)):
                raise RuntimeError("embedding contains non-finite values")
            if attempt > 1:
                print(
                    f"  embed batch ok on retry {attempt} "
                    f"(n={len(clean_texts)})",
                    flush=True,
                )
            return arr
        except Exception as e:
            msg = str(e).lower()
            # Non-retryable client bugs (programming error, not flaky API).
            if "internal: batch size" in msg:
                raise
            delay = _EMBED_RETRY_BASE_S * (2.0 ** min(attempt - 1, 5))
            # Cap wait so we don't sleep minutes between retries.
            delay = min(delay, 30.0)
            print(
                f"  embed batch retry {attempt} in {delay:.1f}s: "
                f"{type(e).__name__}: {str(e)[:160]}",
                flush=True,
            )
            time.sleep(delay)


def embed_one(client: OpenAI, text: str) -> np.ndarray:
    """Single-text embed → (EMBED_DIM,) vector. One HTTP request."""
    arr = embed_texts(client, [text])
    if arr.shape != (1, EMBED_DIM):
        raise RuntimeError(f"embed_one expected (1, {EMBED_DIM}), got {arr.shape}")
    return arr[0]


def _encoding_ids(tok: Tokenizer, text: str) -> list[int]:
    """tokenizers Encoding.ids → list[int]; never iterate a bare None.

    Serialized: HF Tokenizer encode is not thread-safe.
    """
    with _TOK_LOCK:
        enc = tok.encode(text if isinstance(text, str) else (text or ""))
        raw = getattr(enc, "ids", None)
        if raw is None:
            return []
        return list(raw)


def _tok_decode(tok: Tokenizer, ids: list[int] | Any) -> str:
    """Thread-safe tokenizer.decode (shared instance across stream workers)."""
    if not ids:
        return ""
    with _TOK_LOCK:
        return tok.decode(ids)


def chunk_stride_for(chunk_size: int) -> int:
    """Default stride for a window size (CHUNK_32 / CHUNK_256 use fixed strides)."""
    if chunk_size == CHUNK_32:
        return CHUNK_32_STRIDE
    if chunk_size == CHUNK_256:
        return CHUNK_256_STRIDE
    if chunk_size < 1:
        raise ValueError(f"chunk_size must be >= 1, got {chunk_size}")
    return max(1, min(chunk_size, CHUNK_32_STRIDE))


def chunk_spans(
    n_tokens: int, chunk_size: int, stride: int | None = None
) -> list[tuple[int, int]]:
    """Half-open [start, end) spans for overlapping sliding windows.

    Full windows advance by ``stride``. If the last full window does not reach
    the end, one final window ending at n_tokens is added (full-size when
    n_tokens >= chunk_size, else a single short span).
    """
    if chunk_size < 1:
        raise ValueError(f"chunk_size must be >= 1, got {chunk_size}")
    if n_tokens <= 0:
        return []
    step = chunk_stride_for(chunk_size) if stride is None else int(stride)
    if step < 1:
        raise ValueError(f"stride must be >= 1, got {step}")
    if n_tokens <= chunk_size:
        return [(0, n_tokens)]
    spans: list[tuple[int, int]] = []
    start = 0
    while start + chunk_size <= n_tokens:
        spans.append((start, start + chunk_size))
        start += step
    if spans[-1][1] < n_tokens:
        tail_start = n_tokens - chunk_size
        if tail_start != spans[-1][0]:
            spans.append((tail_start, n_tokens))
    return spans


def token_chunks(
    text: str,
    tok: Tokenizer,
    chunk_size: int,
    stride: int | None = None,
) -> list[str]:
    """Split text into overlapping token windows (see chunk_spans).

    Each window is decoded back to a string so the embed API only ever sees
    ~chunk_size tokens of content per request item.
    """
    ids = _encoding_ids(tok, text or "")
    return token_id_chunks(ids, tok, chunk_size, stride=stride)


def token_id_chunks(
    ids: list[int],
    tok: Tokenizer,
    chunk_size: int,
    stride: int | None = None,
) -> list[str]:
    """Same as token_chunks but from pre-tokenized ids (streaming path)."""
    if not ids:
        return []
    out: list[str] = []
    for a, b in chunk_spans(len(ids), chunk_size, stride=stride):
        out.append(_tok_decode(tok, ids[a:b]))
    return out


def cosine_sim(a: np.ndarray, b: np.ndarray) -> float:
    a = np.asarray(a, dtype=float).ravel()
    b = np.asarray(b, dtype=float).ravel()
    na = float(np.linalg.norm(a))
    nb = float(np.linalg.norm(b))
    if na <= 0.0 or nb <= 0.0 or not np.isfinite(na) or not np.isfinite(nb):
        return 0.0
    return _finite(float(np.dot(a, b) / (na * nb)))


def embed_token_chunks(
    client: OpenAI, text: str, tok: Tokenizer, chunk_size: int
) -> np.ndarray:
    """Embed one text as (n_chunks, EMBED_DIM) via fixed-size token windows."""
    chunks = token_chunks(text, tok, chunk_size)
    return embed_texts(client, chunks)


def _stack_embs(rows: list[np.ndarray]) -> np.ndarray:
    if not rows:
        return np.zeros((0, EMBED_DIM), dtype=float)
    arr = np.stack([np.asarray(r, dtype=float).ravel() for r in rows], axis=0)
    if arr.ndim != 2 or arr.shape[1] != EMBED_DIM:
        raise RuntimeError(f"bad emb stack shape {arr.shape}")
    if not np.all(np.isfinite(arr)):
        raise RuntimeError("embedding contains non-finite values")
    return arr


class LiveChunkEmbedder:
    """Queue embeds + stylo metrics as sliding windows complete during stream.

    Full windows (CHUNK_32 / CHUNK_256 with stride 16) are decoded as soon as
    enough tokens exist. **Embeds are not sent one window per HTTP call**:
    32-tok and 256-tok windows share one pending buffer and flush as batches of
    up to ``EMBED_INPUT_MAX`` (1024) inputs per OpenRouter request. Mixing
    window sizes in a batch is intentional (API only sees text strings).

    Each full 256-tok window also queues a local stylo metric extract (same
    decoded text / spans as 256 emb). A final covering window (if needed) is
    enqueued in finalize(), which force-flushes any partial embed batch.
    Soft-cut resume reuses the same instance so already-queued starts are not
    re-done.
    """

    def __init__(
        self,
        client: OpenAI,
        tokenizer: Tokenizer,
        executor: ThreadPoolExecutor,
        *,
        token_ids: list[int] | None = None,
    ) -> None:
        self._client = client
        self._tok = tokenizer
        self._pool = executor
        self.token_ids: list[int] = list(token_ids or [])
        # Next window start not yet enqueued (full windows only).
        self._next_start_32 = 0
        self._next_start_256 = 0
        # Starts already enqueued (so finalize can skip duplicates).
        self._starts_32: set[int] = set()
        self._starts_256: set[int] = set()
        self._starts_stylo: set[int] = set()
        # Pending embed items: (kind "32"|"256", start, text). Mixed sizes OK.
        self._pending_embed: list[tuple[str, int, str]] = []
        # Flushed batches: ([(kind, start), ...], Future → (n, EMBED_DIM)).
        self._batch_jobs: list[
            tuple[list[tuple[str, int]], Future[np.ndarray]]
        ] = []
        # Ordered span starts (enqueue order == span order for each kind).
        self._order_32: list[int] = []
        self._order_256: list[int] = []
        # Stylo futures keyed by window start so order always matches emb_256 /
        # human windows even if a mid-stream submit fails and is retried later.
        self._futs_stylo: dict[int, Future[np.ndarray]] = {}
        # Harvested successful rows — survive finalize retries. Never cleared
        # except on abandon. This is why e/s must not drop while sample is in w.
        self._ready_emb_32: dict[int, np.ndarray] = {}
        self._ready_emb_256: dict[int, np.ndarray] = {}
        self._ready_stylo: dict[int, np.ndarray] = {}
        self._lock = threading.Lock()
        self._finalized = False  # True only after finalize() fully succeeds
        self._pool_dead = False
        self._final_result: (
            tuple[np.ndarray, np.ndarray, np.ndarray] | None
        ) = None
        # How many times finalize fully failed for this sample (park→retry).
        self.score_fail_count = 0
        # Kick any full windows already present (resume path).
        if self.token_ids:
            self.set_tokens(self.token_ids)

    def set_tokens(self, token_ids: list[int]) -> None:
        """Update id buffer; queue embeds + stylo for newly completed windows.

        After a **successful** finalize (``_final_result`` set), no-op — sample
        score is frozen. After abandon (``_pool_dead``), no-op. Failed finalize
        leaves both false so soft retry can enqueue remaining windows without
        wiping ready caches.
        """
        with self._lock:
            if self._pool_dead:
                return
            if self._finalized and self._final_result is not None:
                return
            self.token_ids = list(token_ids)
            self._enqueue_full_unlocked()

    def _try_submit_stylo(self, text: str) -> Future[np.ndarray] | None:
        """Submit one window metric extract; None if the metric pool rejects."""
        try:
            return _submit_window_metrics(text)
        except RuntimeError:
            return None

    def _flush_embed_batches_unlocked(self, *, force: bool) -> bool:
        """Send pending embeds in chunks of EMBED_INPUT_MAX.

        When ``force`` is False, only flush full batches (len >= max).
        When ``force`` is True (finalize), flush all remaining pending items.
        Returns False if the embed pool is shut down.
        """
        if self._pool_dead:
            return False
        # Live flush size (default 1 window/window). Cap at API max 1024.
        max_n = max(1, min(int(EMBED_FLUSH_SIZE), EMBED_INPUT_MAX))
        while self._pending_embed and (
            force or len(self._pending_embed) >= max_n
        ):
            take = min(len(self._pending_embed), max_n)
            batch = self._pending_embed[:take]
            del self._pending_embed[:take]
            texts = [t for _k, _s, t in batch]
            meta = [(k, s) for k, s, _t in batch]
            try:
                # Thread-local client inside worker (shared client can hang).
                fut = self._pool.submit(_embed_texts_one_batch, None, texts)
            except RuntimeError:
                # Put items back so caller can surface a clear error.
                self._pending_embed = batch + self._pending_embed
                self._pool_dead = True
                return False
            self._batch_jobs.append((meta, fut))
            if not force:
                # Only one full batch per non-force call is enough; loop if
                # pending still ≥ max (e.g. catch-up after long pause).
                continue
        return True

    def _queue_embed_unlocked(self, kind: str, start: int, text: str) -> bool:
        """Enqueue one window for the next embed batch. False if pool dead."""
        if kind == "32":
            if start in self._starts_32:
                return True
            self._starts_32.add(start)
            # Re-queue after failed batch must not duplicate order entries.
            if start not in self._order_32:
                self._order_32.append(start)
        elif kind == "256":
            if start in self._starts_256:
                return True
            self._starts_256.add(start)
            if start not in self._order_256:
                self._order_256.append(start)
        else:
            raise ValueError(f"unknown embed kind {kind!r}")
        self._pending_embed.append((kind, start, text))
        return self._flush_embed_batches_unlocked(force=False)

    def _enqueue_full_unlocked(self) -> None:
        ids = self.token_ids
        n = len(ids)
        while self._next_start_32 + CHUNK_32 <= n:
            start = self._next_start_32
            text = _tok_decode(self._tok, ids[start : start + CHUNK_32])
            if not self._queue_embed_unlocked("32", start, text):
                return
            self._next_start_32 += CHUNK_32_STRIDE
        while self._next_start_256 + CHUNK_256 <= n:
            start = self._next_start_256
            text = _tok_decode(self._tok, ids[start : start + CHUNK_256])
            if not self._queue_embed_unlocked("256", start, text):
                return
            # Stylo shares 256-tok windows; still one local extract per window.
            if start not in self._starts_stylo:
                mfut = self._try_submit_stylo(text)
                if mfut is not None:
                    self._futs_stylo[start] = mfut
                    self._starts_stylo.add(start)
                # If submit fails, leave start free — finalize retries.
            self._next_start_256 += CHUNK_256_STRIDE

    def progress(self) -> tuple[int, int]:
        """(completed window embeds, total window embeds) including pending.

        Counts harvested ready rows + in-flight batch futures + pending.
        Failed futures are **not** counted as done (they look stuck otherwise).
        """
        with self._lock:
            ready_32 = set(self._ready_emb_32)
            ready_256 = set(self._ready_emb_256)
            pending = len(self._pending_embed)
            jobs = list(self._batch_jobs)
        done = len(ready_32) + len(ready_256)
        in_flight = 0
        for meta, fut in jobs:
            # Skip starts already harvested into ready (avoid double-count).
            n_left = 0
            for kind, start in meta:
                if kind == "32" and start in ready_32:
                    continue
                if kind == "256" and start in ready_256:
                    continue
                n_left += 1
            if n_left <= 0:
                continue
            if not fut.done():
                in_flight += n_left
                continue
            # Done future: only count as done if it succeeded.
            try:
                if fut.exception() is not None:
                    # Failed — will be re-queued; count as still outstanding.
                    in_flight += n_left
                else:
                    done += n_left
            except Exception:
                in_flight += n_left
        total = done + in_flight + pending
        return done, total

    def stylo_progress(self) -> tuple[int, int]:
        """(completed, submitted) stylo windows including harvested ready rows."""
        with self._lock:
            ready = len(self._ready_stylo)
            futs = [
                (s, f)
                for s, f in self._futs_stylo.items()
                if s not in self._ready_stylo
            ]
        done_f = sum(1 for _s, f in futs if f.done())
        pend_f = sum(1 for _s, f in futs if not f.done())
        done = ready + done_f
        return done, done + pend_f

    def abandon(self) -> None:
        """Kill this sample's emb+stylo forever.

        **Only** call when the sample or whole trial is dying:
        shutdown, empty attempt (never had tokens), surplus after a successful
        score, or hard stop. Soft-cut / fin-err / retry must never call this.
        """
        with self._lock:
            self._finalized = True
            self._pool_dead = True
            self._final_result = None
            for fut in self._futs_stylo.values():
                if not fut.done():
                    fut.cancel()
            self._futs_stylo.clear()
            self._pending_embed.clear()
            for _meta, fut in self._batch_jobs:
                if not fut.done():
                    fut.cancel()
            self._batch_jobs.clear()
            # Only place ready caches are cleared.
            self._ready_emb_32.clear()
            self._ready_emb_256.clear()
            self._ready_stylo.clear()

    def _queue_span_unlocked(
        self, ids: list[int], start: int, end: int, *, kind: str
    ) -> bool:
        """Enqueue one embed span if start not yet done. False if pool down."""
        if kind == "32" and start in self._ready_emb_32:
            return True
        if kind == "256" and start in self._ready_emb_256:
            return True
        starts = self._starts_32 if kind == "32" else self._starts_256
        if start in starts:
            return True
        text = _tok_decode(self._tok, ids[start:end])
        return self._queue_embed_unlocked(kind, start, text)

    def _submit_stylo_span_unlocked(
        self, ids: list[int], start: int, end: int
    ) -> bool:
        """Submit one stylo window if start not yet done. False if pool down."""
        if start in self._ready_stylo:
            return True
        if start in self._starts_stylo and start in self._futs_stylo:
            return True
        # Failed prior attempt cleared fut but left start — allow re-submit.
        if start in self._starts_stylo and start not in self._futs_stylo:
            self._starts_stylo.discard(start)
        text = _tok_decode(self._tok, ids[start:end])
        fut = self._try_submit_stylo(text)
        if fut is None:
            return False
        self._futs_stylo[start] = fut
        self._starts_stylo.add(start)
        return True

    def _drop_failed_stylo_unlocked(self, start: int) -> None:
        """Forget a failed stylo *future* so finalize can re-submit.

        Never touches ``_ready_stylo`` — successful windows are immortal until
        ``abandon()``.
        """
        self._futs_stylo.pop(start, None)
        self._starts_stylo.discard(start)

    def _wait_fut(
        self,
        fut: Future[np.ndarray],
        *,
        on_tick: Any | None = None,
        max_wait_s: float | None = None,
    ) -> np.ndarray:
        """Wait with short timeouts so UI/Ctrl+C can run (no multi-min freeze).

        Default: wait indefinitely (embed HTTP retries forever with backoff).
        ``max_wait_s``: if set, raise TimeoutError after that many seconds so a
        caller can fail/retry instead of hanging on a truly stuck future.
        """
        t0 = time.perf_counter()
        while True:
            try:
                return fut.result(timeout=0.25)
            except TimeoutError:
                if max_wait_s is not None and time.perf_counter() - t0 >= float(
                    max_wait_s
                ):
                    try:
                        fut.cancel()
                    except Exception:
                        pass
                    raise TimeoutError(
                        f"future hung >{max_wait_s:.0f}s (likely stuck embed HTTP)"
                    )
                if on_tick is not None:
                    try:
                        on_tick()
                    except Exception:
                        pass

    def _harvest_embed_batches(self, *, on_tick: Any | None = None) -> None:
        """Pull completed batch futures into ready caches (keeps progress)."""
        with self._lock:
            jobs = list(self._batch_jobs)
        failed: list[tuple[list[tuple[str, int]], BaseException]] = []
        succeeded: set[int] = set()  # id(meta)
        for meta, fut in jobs:
            if not fut.done():
                try:
                    arr = np.asarray(
                        self._wait_fut(fut, on_tick=on_tick), dtype=float
                    )
                except Exception as e:
                    failed.append((meta, e))
                    continue
            else:
                try:
                    arr = np.asarray(fut.result(), dtype=float)
                except Exception as e:
                    failed.append((meta, e))
                    continue
            if arr.ndim != 2 or arr.shape[0] != len(meta) or arr.shape[1] != EMBED_DIM:
                failed.append(
                    (
                        meta,
                        RuntimeError(
                            f"embed batch shape {arr.shape}, expected "
                            f"({len(meta)}, {EMBED_DIM})"
                        ),
                    )
                )
                continue
            with self._lock:
                for i, (kind, start) in enumerate(meta):
                    row = arr[i]
                    if kind == "32":
                        self._ready_emb_32[start] = row
                    elif kind == "256":
                        self._ready_emb_256[start] = row
            succeeded.add(id(meta))

        failed_ids = {id(m) for m, _ in failed}
        with self._lock:
            new_jobs: list[
                tuple[list[tuple[str, int]], Future[np.ndarray]]
            ] = []
            for meta, fut in self._batch_jobs:
                mid = id(meta)
                if mid in failed_ids:
                    # Free starts so finalize can re-queue (never pop ready).
                    for kind, start in meta:
                        if kind == "32":
                            if start not in self._ready_emb_32:
                                self._starts_32.discard(start)
                        else:
                            if start not in self._ready_emb_256:
                                self._starts_256.discard(start)
                    continue
                if mid in succeeded:
                    # Already in ready — drop job entry (keep list small).
                    continue
                # Still running (shouldn't remain after sequential wait) or other.
                new_jobs.append((meta, fut))
            self._batch_jobs = new_jobs

        if failed:
            meta0, err0 = failed[0]
            raise RuntimeError(
                f"embed batch failed ({type(err0).__name__}: {err0})"
            ) from err0

    def finalize(
        self, *, on_tick: Any | None = None
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Cover remaining spans, flush, wait → (emb_32, emb_256, metrics).

        Safe to call multiple times: successful rows stay in ready caches so
        e/s progress does not collapse on retry. ``_finalized`` only after
        full success (so soft-cut park can still reuse this embedder).

        ``on_tick`` is called every ~0.25s while blocked on a future so the
        progress line keeps moving (avoids multi-minute “dead terminal”).
        """
        with self._lock:
            if self._finalized and self._final_result is not None:
                return self._final_result
            if self._pool_dead:
                raise RuntimeError("LiveChunkEmbedder abandoned")
            self._enqueue_full_unlocked()
            ids = list(self.token_ids)
            n = len(ids)
            for start, end in chunk_spans(n, CHUNK_32, CHUNK_32_STRIDE):
                if not self._queue_span_unlocked(ids, start, end, kind="32"):
                    raise RuntimeError(
                        "embed pool shut down before final window embed (32)"
                    )
            for start, end in chunk_spans(n, CHUNK_256, CHUNK_256_STRIDE):
                if not self._queue_span_unlocked(ids, start, end, kind="256"):
                    raise RuntimeError(
                        "embed pool shut down before final window embed (256)"
                    )
                if not self._submit_stylo_span_unlocked(ids, start, end):
                    raise RuntimeError(
                        "metric pool shut down before final stylo window"
                    )
            if not self._flush_embed_batches_unlocked(force=True):
                raise RuntimeError(
                    "embed pool shut down before flushing final embed batch"
                )
            order_32 = list(self._order_32)
            order_256 = list(self._order_256)

        # Harvest embeds into ready caches (partial success kept on failure).
        self._harvest_embed_batches(on_tick=on_tick)

        # Stylo: harvest each window into ready; failed window is dropped for
        # re-submit on next finalize call (does not clear other ready rows).
        with self._lock:
            futs_stylo = dict(self._futs_stylo)
            ready_stylo = dict(self._ready_stylo)
        for s in order_256:
            if s in ready_stylo:
                continue
            fut = futs_stylo.get(s)
            if fut is None:
                with self._lock:
                    self._drop_failed_stylo_unlocked(s)
                raise RuntimeError(f"missing stylo future for window start={s}")
            try:
                row = _sanitize_metric_row(
                    np.asarray(self._wait_fut(fut, on_tick=on_tick), dtype=float)
                )
            except Exception as e:
                with self._lock:
                    self._drop_failed_stylo_unlocked(s)
                raise RuntimeError(
                    f"stylo window start={s} failed ({type(e).__name__}: {e})"
                ) from e
            with self._lock:
                self._ready_stylo[s] = row
            if on_tick is not None:
                try:
                    on_tick()
                except Exception:
                    pass

        with self._lock:
            by_32 = dict(self._ready_emb_32)
            by_256 = dict(self._ready_emb_256)
            by_sty = dict(self._ready_stylo)

        if any(s not in by_32 for s in order_32):
            missing = [s for s in order_32 if s not in by_32]
            raise RuntimeError(f"missing 32-tok embed starts: {missing[:8]}…")
        if any(s not in by_256 for s in order_256):
            missing = [s for s in order_256 if s not in by_256]
            raise RuntimeError(f"missing 256-tok embed starts: {missing[:8]}…")
        if any(s not in by_sty for s in order_256):
            missing = [s for s in order_256 if s not in by_sty]
            raise RuntimeError(f"missing stylo starts: {missing[:8]}…")

        emb_32 = _stack_embs([by_32[s] for s in order_32])
        emb_256 = _stack_embs([by_256[s] for s in order_256])
        metrics = np.stack([by_sty[s] for s in order_256], axis=0)

        expect_32 = len(chunk_spans(n, CHUNK_32, CHUNK_32_STRIDE))
        expect_256 = len(chunk_spans(n, CHUNK_256, CHUNK_256_STRIDE))
        if emb_32.shape[0] != expect_32 or emb_256.shape[0] != expect_256:
            raise RuntimeError(
                f"live embed count mismatch: "
                f"32={emb_32.shape[0]}/{expect_32} "
                f"256={emb_256.shape[0]}/{expect_256} (n_tokens={n})"
            )
        if metrics.shape[0] != expect_256:
            raise RuntimeError(
                f"live stylo count mismatch: "
                f"stylo={metrics.shape[0]}/{expect_256} (n_tokens={n})"
            )
        result = (emb_32, emb_256, metrics)
        with self._lock:
            self._finalized = True
            self._final_result = result
        return result


def chunk_geom_median_cos_scores(
    ai_embs: list[np.ndarray],
    human_emb: np.ndarray,
) -> tuple[float, list[float]]:
    """Position-aligned chunk embedding match.

    For each human chunk index i:
      - stack sample embeddings at i (samples missing i are skipped)
      - geometric median → 4096-d vector
      - cosine similarity vs human chunk i
    Aggregate objective = median of those cosines (robust to a few bad chunks).

    Per-sample score = median cosine of that sample's chunks vs human (no
    geometric median across samples), used only for score_std diagnostics.
    """
    human_emb = np.asarray(human_emb, dtype=float)
    if human_emb.ndim != 2 or human_emb.shape[1] != EMBED_DIM:
        raise ValueError(
            f"human_emb shape {getattr(human_emb, 'shape', None)}, "
            f"expected (n, {EMBED_DIM})"
        )
    n_h = int(human_emb.shape[0])
    if n_h == 0:
        zeros = [0.0] * len(ai_embs)
        return 0.0, zeros

    chunk_scores: list[float] = []
    for i in range(n_h):
        rows: list[np.ndarray] = []
        for emb in ai_embs:
            emb = np.asarray(emb, dtype=float)
            if emb.ndim == 2 and i < emb.shape[0]:
                rows.append(emb[i])
        if not rows:
            chunk_scores.append(0.0)
            continue
        med = geometric_median(np.stack(rows, axis=0))
        chunk_scores.append(cosine_sim(med, human_emb[i]))

    agg_score = _finite(float(np.median(chunk_scores)))

    per_sample: list[float] = []
    for emb in ai_embs:
        emb = np.asarray(emb, dtype=float)
        if emb.ndim != 2 or emb.shape[0] == 0:
            per_sample.append(0.0)
            continue
        n = min(n_h, int(emb.shape[0]))
        sims = [cosine_sim(emb[i], human_emb[i]) for i in range(n)]
        # Penalize missing tail chunks so shorter AI answers do not look perfect.
        if n < n_h:
            sims.extend([0.0] * (n_h - n))
        per_sample.append(_finite(float(np.median(sims))))

    return agg_score, per_sample


def chunk_geom_median_metric_scores(
    ai_metrics: list[np.ndarray],
    human_metrics: np.ndarray,
    scale: np.ndarray,
) -> tuple[float, list[float]]:
    """Position-aligned windowed stylo match (mirrors chunk_geom_median_cos_scores).

    For each human window index i:
      - stack sample metric vectors at i (samples missing i are skipped)
      - geometric median in human-scaled feature space
      - scale-normalized L2 distance vs human window i → 1/(1+d)
    Aggregate stylo = median of those window match scores.

    Per-sample score = median of that sample's per-window match scores
    (missing tail windows scored as 0), used only for score_std diagnostics.
    """
    human_metrics = np.asarray(human_metrics, dtype=float)
    scale = np.asarray(scale, dtype=float).ravel()
    n_feat = len(FEATURE_NAMES)
    if human_metrics.ndim != 2 or human_metrics.shape[1] != n_feat:
        raise ValueError(
            f"human_metrics shape {getattr(human_metrics, 'shape', None)}, "
            f"expected (n, {n_feat})"
        )
    if scale.shape != (n_feat,):
        raise ValueError(f"scale shape {scale.shape}, expected ({n_feat},)")
    n_h = int(human_metrics.shape[0])
    if n_h == 0:
        zeros = [0.0] * len(ai_metrics)
        return 0.0, zeros

    human_scaled = human_metrics / scale
    chunk_scores: list[float] = []
    for i in range(n_h):
        rows: list[np.ndarray] = []
        for mat in ai_metrics:
            mat = np.asarray(mat, dtype=float)
            if mat.ndim == 2 and i < mat.shape[0] and mat.shape[1] == n_feat:
                rows.append(mat[i] / scale)
        if not rows:
            chunk_scores.append(0.0)
            continue
        med = geometric_median(np.stack(rows, axis=0))
        dist = _finite(float(np.linalg.norm(med - human_scaled[i])))
        chunk_scores.append(_per_sample_match_score(dist))

    agg_score = _finite(float(np.median(chunk_scores)))

    per_sample: list[float] = []
    for mat in ai_metrics:
        mat = np.asarray(mat, dtype=float)
        if mat.ndim != 2 or mat.shape[0] == 0 or mat.shape[1] != n_feat:
            per_sample.append(0.0)
            continue
        n = min(n_h, int(mat.shape[0]))
        sims = [
            _per_sample_match_score(
                _scaled_metric_distance(mat[i], human_metrics[i], scale)
            )
            for i in range(n)
        ]
        if n < n_h:
            sims.extend([0.0] * (n_h - n))
        per_sample.append(_finite(float(np.median(sims))))

    return agg_score, per_sample


# ---------------------------------------------------------------------------
# scoring
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class HumanCache:
    # Stylo: one metric vector per 256-tok window (aligned with emb_256 spans).
    metrics: np.ndarray  # (n_chunks_256, n_features)
    emb_32: np.ndarray  # (n_chunks_32, EMBED_DIM)
    emb_256: np.ndarray  # (n_chunks_256, EMBED_DIM)


def arithmetic_mean2(a: float, b: float) -> float:
    """Arithmetic mean of two scores."""
    return 0.5 * (float(a) + float(b))


# Trial user_attr keys for the three raw component scores (not Optuna objectives).
ATTR_STYLO = "stylo"
ATTR_EMB_32 = "32 emb"
ATTR_EMB_256 = "256 emb"
COMPONENT_ATTR_KEYS = (ATTR_STYLO, ATTR_EMB_32, ATTR_EMB_256)
# Derived sampling params (min_p never free; window free only if exam omits it).
ATTR_MIN_P = "min_p"
ATTR_LOGIT_TOLERANCE = "logit_tolerance"


# Human-readable names for the three single-objective studies (:0 / :1 / :2).
OBJECTIVE_NAMES = (
    "stylo+256 mean",
    "stylo+32 mean",
    "32 emb",
)
# Short labels for progress / status lines (keep ≤ ~5 chars).
OBJECTIVE_SHORT = ("s+256", "s+32", "32e")


def _obj_short(obj_index: int) -> str:
    if 0 <= obj_index < len(OBJECTIVE_SHORT):
        return OBJECTIVE_SHORT[obj_index]
    return f"obj{obj_index}"


def _fmt_count(n: int | float) -> str:
    """Compact integer for progress (e.g. 45000 → 45k)."""
    n = int(n)
    if abs(n) >= 10_000:
        return f"{n / 1000:.0f}k"
    if abs(n) >= 1000:
        return f"{n / 1000:.1f}k".rstrip("0").rstrip(".")
    return str(n)


def _trial_prefix(
    trial_number: int,
    obj_index: int,
    *,
    top_k: int,
    top_p: float,
    temperature: float,
) -> str:
    """Compact live-line prefix (fits narrow terminals with status fields)."""
    return (
        f"t{trial_number} {_obj_short(obj_index)} "
        f"k{_fmt_count(top_k)} p{top_p:.1f} T{temperature:.1f} "
    )


def objectives_from_components(
    stylo: float, emb_32: float, emb_256: float
) -> tuple[float, float, float]:
    """Map raw components → the three scalar objectives (one per study)."""
    return (
        arithmetic_mean2(stylo, emb_256),  # 0: stylo + 256 emb mean
        arithmetic_mean2(stylo, emb_32),  # 1: stylo + 32 emb mean
        float(emb_32),  # 2: 32 emb only
    )


def objective_value_from_components(
    stylo: float, emb_32: float, emb_256: float, obj_index: int
) -> float:
    """Single scalar for study shard ``obj_index``."""
    if not 0 <= obj_index < len(OBJECTIVE_NAMES):
        raise ValueError(f"obj_index out of range: {obj_index}")
    return float(objectives_from_components(stylo, emb_32, emb_256)[obj_index])


@dataclass(frozen=True)
class ScoreResult:
    """All three scalars plus raw components and per-sample std."""

    # (stylo+256 mean, stylo+32 mean, 32 emb) — study i maximizes values[i].
    values: tuple[float, float, float]
    # Raw component scores stored as trial user_attrs.
    stylo: float
    emb_32: float
    emb_256: float
    # Std of per-sample match scores — higher = more disagreement between samples.
    score_std: tuple[float, float, float]


@dataclass(frozen=True)
class GeneratedSample:
    """One finished completion with live window embeds + stylo metrics."""

    text: str
    emb_32: np.ndarray  # (n_chunks_32, EMBED_DIM)
    emb_256: np.ndarray  # (n_chunks_256, EMBED_DIM)
    # Stylo: metric vector per 256-tok window (aligned with emb_256).
    metrics: np.ndarray  # (n_chunks_256, n_features)


# Human embed: same batch size as live flush by default. Parallel via
# EMBED_BATCH_PARALLEL. Optional delay: HUMAN_EMBED_DELAY_S.
_HUMAN_EMBED_DELAY_S = float(os.getenv("HUMAN_EMBED_DELAY_S") or "0")
# Default 64 windows per HTTP (override HUMAN_EMBED_BATCH).
_HUMAN_EMBED_BATCH = int(os.getenv("HUMAN_EMBED_BATCH") or "64")


def compute_human_cache(
    answer_text: str, embed_client: OpenAI, tokenizer: Tokenizer
) -> HumanCache:
    """Build human baselines: emb_32, emb_256, and stylo windows.

    Default embed style matches the setup that worked at ~100k scale:
    **1 window per request**, many concurrent HTTP workers (not one mega-batch
    of 1024×256-tok ≈ 250k tokens). Stylo runs in parallel on the local pool.
    """
    chunks_32 = token_chunks(answer_text, tokenizer, CHUNK_32)
    chunks_256 = token_chunks(answer_text, tokenizer, CHUNK_256)
    delay = max(0.0, _HUMAN_EMBED_DELAY_S)
    batch_n = max(1, min(_HUMAN_EMBED_BATCH, EMBED_INPUT_MAX))
    # Parallel many small calls unless user forced a delay (then serial).
    parallel = delay <= 0.0
    n_batches = (
        (len(chunks_32) + batch_n - 1) // batch_n
        + (len(chunks_256) + batch_n - 1) // batch_n
    )
    to_s = float(os.getenv("EMBED_HTTP_TIMEOUT_S") or "300")
    par_n = _embed_batch_parallel() if parallel else 1
    print(
        f"  human embeds: batch≤{batch_n} parallel={par_n} "
        f"delay={delay:.1f}s timeout={to_s:.0f}s "
        f"(~{n_batches} HTTP calls)",
        flush=True,
    )

    with ThreadPoolExecutor(
        max_workers=3, thread_name_prefix="human-cache"
    ) as pool:
        fut_stylo = pool.submit(
            metric_token_chunks,
            answer_text,
            tokenizer,
            chunk_size=STYLO_CHUNK,
            stride=STYLO_STRIDE,
        )
        fut_32 = pool.submit(
            embed_texts,
            embed_client,
            chunks_32,
            parallel=parallel,
            inter_batch_delay_s=delay,
            batch_size=batch_n,
        )
        fut_256 = pool.submit(
            embed_texts,
            embed_client,
            chunks_256,
            parallel=parallel,
            inter_batch_delay_s=delay,
            batch_size=batch_n,
        )
        emb_32 = fut_32.result()
        emb_256 = fut_256.result()
        metrics = fut_stylo.result()
    if metrics.shape[0] != emb_256.shape[0]:
        raise RuntimeError(
            f"human stylo/emb_256 window count mismatch: "
            f"stylo={metrics.shape[0]} emb256={emb_256.shape[0]}"
        )
    return HumanCache(metrics=metrics, emb_32=emb_32, emb_256=emb_256)


def _per_sample_match_score(delta: float) -> float:
    """Map absolute deviation → (0, 1], same form as aggregate objectives."""
    d = abs(_finite(delta, default=1e9))
    return 1.0 / (1.0 + d)


def score_answers(
    samples: list[GeneratedSample],
    human: HumanCache,
) -> ScoreResult:
    """Score using windowed stylo metrics + embeddings gathered during generation."""
    if len(samples) != N_SAMPLES:
        raise ValueError(f"need {N_SAMPLES} samples, got {len(samples)}")
    cleaned = [(s.text or "").strip() for s in samples]
    if any(not t for t in cleaned):
        raise ValueError("one or more answers are empty after decode")

    n_feat = len(FEATURE_NAMES)
    human_m = np.asarray(human.metrics, dtype=float)
    if human_m.ndim != 2 or human_m.shape[1] != n_feat:
        raise RuntimeError(
            f"human.metrics shape {human_m.shape}, expected (n, {n_feat})"
        )

    # stylo: live 256-tok window metrics (same spans as emb_256) → emb-style
    # per-window geom-median L2 match → median over windows.
    ai_metrics = [np.asarray(s.metrics, dtype=float) for s in samples]
    for i, mat in enumerate(ai_metrics):
        if mat.ndim != 2 or mat.shape[1] != n_feat:
            raise RuntimeError(f"sample {i} metrics shape {mat.shape}")
    scale = _metric_scale(human_m)
    stylo, s_stylo = chunk_geom_median_metric_scores(ai_metrics, human_m, scale)

    # 32 emb / 256 emb: embeddings were fired per-chunk during streaming.
    ai_32 = [np.asarray(s.emb_32, dtype=float) for s in samples]
    ai_256 = [np.asarray(s.emb_256, dtype=float) for s in samples]
    for i, (e32, e256) in enumerate(zip(ai_32, ai_256)):
        if e32.ndim != 2 or e32.shape[1] != EMBED_DIM:
            raise RuntimeError(f"sample {i} emb_32 shape {e32.shape}")
        if e256.ndim != 2 or e256.shape[1] != EMBED_DIM:
            raise RuntimeError(f"sample {i} emb_256 shape {e256.shape}")

    emb_32, s_32 = chunk_geom_median_cos_scores(ai_32, human.emb_32)
    emb_256, s_256 = chunk_geom_median_cos_scores(ai_256, human.emb_256)

    def _std(scores: list[float]) -> float:
        return _finite(float(np.std(scores, ddof=0)))

    components = (stylo, emb_32, emb_256)
    if not all(np.isfinite(v) for v in components):
        raise RuntimeError(f"non-finite component scores: {components}")
    values = objectives_from_components(stylo, emb_32, emb_256)
    if not all(np.isfinite(v) for v in values):
        raise RuntimeError(f"non-finite objectives from {components} → {values}")

    return ScoreResult(
        values=values,
        stylo=stylo,
        emb_32=emb_32,
        emb_256=emb_256,
        score_std=(_std(s_stylo), _std(s_32), _std(s_256)),
    )


# ---------------------------------------------------------------------------
# NovelAI streaming + countdown pool
# ---------------------------------------------------------------------------

def _parse_sse_line(line: str) -> str | None:
    """Extract SSE data payload; ignore comments / event / id fields."""
    if not line:
        return None
    if line.startswith("data:"):
        return line[5:].lstrip()
    # Keepalive comments, event names, ids — not JSON payloads.
    if line.startswith(":") or line.startswith("event:") or line.startswith("id:"):
        return None
    # Some gateways emit bare JSON lines without the data: prefix.
    if line.startswith("{"):
        return line
    return None


def _coerce_token_ids(raw_ids: Any) -> list[int]:
    """Normalize token_ids field from a stream choice (list or single)."""
    if raw_ids is None:
        return []
    if isinstance(raw_ids, (list, tuple)):
        out: list[int] = []
        for t in raw_ids:
            try:
                out.append(int(t))
            except (TypeError, ValueError):
                continue
        return out
    try:
        return [int(raw_ids)]
    except (TypeError, ValueError):
        return []


@dataclass
class Preset:
    top_k: int
    min_p: float
    top_p: float
    temperature: float

    def to_api(self) -> dict[str, float | int]:
        return {
            "top_k": int(self.top_k),
            "min_p": float(self.min_p),
            "top_p": float(self.top_p),
            "temperature": float(self.temperature),
        }


# Concurrent embed HTTP workers per trial (I/O-bound). Each job is one flush
# (default 1 window). High concurrency is fine for *small* requests; mega-batches
# of 1024 windows were the reliability problem, not uncapped workers.
# Override: EMBED_POOL_WORKERS=N
def _embed_pool_workers() -> int:
    raw = (os.getenv("EMBED_POOL_WORKERS") or "").strip()
    if raw:
        try:
            return max(1, int(raw))
        except ValueError:
            pass
    # Align with EMBED_BATCH_PARALLEL when unset.
    return _embed_batch_parallel()


EMBED_POOL_WORKERS = _embed_pool_workers()


def _is_hard_trial_error(exc: BaseException) -> bool:
    """Errors that must fail the trial (not soft-park forever).

    Soft: transient embed/decode glitches already returned as partial.
    Hard: NovelAI 4xx client errors, Optuna bugs, assertion failures, etc.
    """
    if isinstance(exc, (KeyboardInterrupt, SystemExit, MemoryError)):
        return True
    if isinstance(exc, (AssertionError, TypeError, NameError, AttributeError)):
        return True
    if isinstance(exc, RuntimeError):
        msg = str(exc).lower()
        # Soft-cut paths never raise these; real HTTP 4xx does.
        if "novelai http" in msg:
            return True
        if "http 4" in msg:
            return True
    return False


class NaiPool:
    """
    Fill N_SAMPLES completions for one preset (**classic 1-phase**).

    Prompt = system + prefill. One stream generates ``answer_tokens``. Soft-cut
    resume parks partial token ids and continues with the same prompt layout.

    Sample capacity: finished + streaming + finalizing + idle ≤ N_SAMPLES.
    Concurrent NovelAI POSTs: at most MAX_CONCURRENT_STREAMS *streaming*
    (finalizing does **not** hold a stream slot — embeds drain after the
    SSE ends so the next sample can start immediately). Prefer resuming
    idle partials over new blanks (never spawn a blank while idle remains).
    Shared countdown gates every new HTTP request: wait `delay` after each
    spawn (base 1s). On 429/err/soft-cut, wait current delay then double
    (… → 64 → 128) for *spawn pacing only*.

    **No sample wipe:** once a sample has any tokens, those tokens are kept
    until the sample is successfully finished or the whole trial shuts down.
    Soft-cut / metric fail / stream error → park the same ids (re-score or
    resume). Never blank-replace a partial or completed generation. True
    empty streams (0 tokens) free the slot without creating a blank idle.
    First stream token only resets the *delay value* to 1s — not the active
    countdown.

    While tokens stream, LiveChunkEmbedder queues completed 32/256-token
    windows and flushes OpenRouter embed batches of up to EMBED_INPUT_MAX
    (mixed sizes OK), plus one stylo metric extract per 256-tok window, so
    scoring is nearly ready at gen end. When generation finishes, the worker
    releases its stream slot, force-flushes any partial embed batch, then waits
    on remaining embed + stylo futures under a separate finalizing count.

    Terminal: only updates a single LineProgress line (no stream spam).
    """

    def __init__(
        self,
        *,
        token: str,
        model: str,
        prompt_ids: list[int],
        answer_tokens: int,
        token_bias: dict[str, float],
        preset: Preset,
        tokenizer: Tokenizer,
        embed_client: OpenAI,
        embed_pool: ThreadPoolExecutor,
        progress: LineProgress | None = None,
        progress_prefix: str = "",
    ) -> None:
        self.token = token
        self.model = model
        # system + prefill (assembled in ExamBundle).
        self.prompt_ids = list(prompt_ids)
        self.answer_tokens = int(answer_tokens)
        self.total_answer_tokens = int(answer_tokens)
        self.token_bias = dict(token_bias)
        self.preset = preset
        self.tokenizer = tokenizer
        self.embed_client = embed_client
        self.embed_pool = embed_pool
        self._progress = progress
        self._progress_prefix = progress_prefix

        self._cv = threading.Condition()
        self._finished: list[GeneratedSample] = []
        # Parked soft-cut partials keep LiveChunkEmbedder so already-submitted
        # full-window embeds continue / resume without redoing work.
        self._idle: list[LiveChunkEmbedder] = []
        # NovelAI POSTs currently open (counts against MAX_CONCURRENT_STREAMS).
        self._streaming = 0
        # Gen done, waiting on embed finalize — holds a sample slot, not a
        # stream slot. Entries: (embedder, answer_token_count for UI).
        self._finalizing: list[tuple[LiveChunkEmbedder, int]] = []
        self._delay = BASE_DELAY
        self._deadline = 0.0
        self._error: BaseException | None = None
        self._live_tokens: dict[int, int] = {}  # worker_id -> tokens so far
        self._live_embedders: dict[int, LiveChunkEmbedder] = {}
        # Every embedder created this trial — shutdown abandons *all* so hard
        # errors / stop races cannot leave orphan stylo/emb jobs on the pools.
        self._owned_embedders: list[LiveChunkEmbedder] = []
        self._next_wid = 0
        self._last_ui = 0.0
        self._stop = False  # set on run() exit so late workers go quiet
        self._http_attempts = 0
        self._t0 = time.perf_counter()
        # Peak e/s so UI never rewinds when a sample briefly leaves the sum
        # or a counter is recalculated (completed work is still kept in caches).
        self._emb_peak = (0, 0)
        self._sty_peak = (0, 0)

    def _new_embedder(
        self, token_ids: list[int] | None = None
    ) -> LiveChunkEmbedder:
        emb = LiveChunkEmbedder(
            self.embed_client,
            self.tokenizer,
            self.embed_pool,
            token_ids=token_ids,
        )
        self._owned_embedders.append(emb)
        return emb

    @staticmethod
    def _kill_sample_metrics(embedder: LiveChunkEmbedder | None) -> None:
        """Destroy emb+stylo for a sample that is dying (or trial shutdown).

        Soft-cut / fin-err / park must never call this.
        """
        if embedder is None:
            return
        try:
            embedder.abandon()
        except Exception:
            pass

    # Back-compat alias used in a few call sites during rename.
    _safe_abandon = _kill_sample_metrics

    def _decode(self, token_ids: list[int], text_parts: list[str] | None = None) -> str:
        """
        Prefer tokenizer.decode(token_ids). Text deltas are not source of truth:
        on soft-cut resume only ids are parked, so stream text is incomplete.
        """
        if token_ids:
            return _tok_decode(self.tokenizer, token_ids)
        if text_parts:
            return "".join(text_parts)
        return ""

    def _ids_from_text(self, text: str) -> list[int]:
        """Approximate token ids when the API omits token_ids (local tokenizer)."""
        if not isinstance(text, str) or not text:
            return []
        raw = _encoding_ids(self.tokenizer, text)
        kept, _dropped = _sanitize_token_ids(raw)
        return kept

    def _embed_progress_from(
        self,
        finished: list[GeneratedSample],
        live: list[LiveChunkEmbedder],
        finalizing: list[LiveChunkEmbedder],
        idle: list[LiveChunkEmbedder],
    ) -> tuple[int, int]:
        """Sum emb windows from snapshots (no pool lock held)."""
        done = 0
        total = 0
        for sample in finished:
            n = int(sample.emb_32.shape[0]) + int(sample.emb_256.shape[0])
            done += n
            total += n
        for emb in live:
            d, t = emb.progress()
            done += d
            total += t
        for emb in finalizing:
            d, t = emb.progress()
            done += d
            total += t
        for emb in idle:
            d, t = emb.progress()
            done += d
            total += t
        return done, total

    def _stylo_progress_from(
        self,
        finished: list[GeneratedSample],
        live: list[LiveChunkEmbedder],
        finalizing: list[LiveChunkEmbedder],
        idle: list[LiveChunkEmbedder],
    ) -> tuple[int, int]:
        """Sum stylo windows from snapshots (no pool lock held)."""
        done = 0
        total = 0
        for sample in finished:
            n = int(sample.metrics.shape[0])
            done += n
            total += n
        for emb in live:
            d, t = emb.stylo_progress()
            done += d
            total += t
        for emb in finalizing:
            d, t = emb.stylo_progress()
            done += d
            total += t
        for emb in idle:
            d, t = emb.stylo_progress()
            done += d
            total += t
        return done, total

    def _progress_msg(
        self,
        *,
        note: str = "",
        finished: list[GeneratedSample] | None = None,
        live_tokens: dict[int, int] | None = None,
        finalizing: list[tuple[LiveChunkEmbedder, int]] | None = None,
        idle: list[LiveChunkEmbedder] | None = None,
        live_embs: list[LiveChunkEmbedder] | None = None,
        streaming: int | None = None,
        delay: float | None = None,
        deadline: float | None = None,
        http_attempts: int | None = None,
        t0: float | None = None,
    ) -> str:
        """Build progress line. Prefer pre-snapshotted lists (no _cv held)."""
        finished = self._finished if finished is None else finished
        live_tokens = self._live_tokens if live_tokens is None else live_tokens
        finalizing = self._finalizing if finalizing is None else finalizing
        idle = self._idle if idle is None else idle
        if live_embs is None:
            live_embs = list(self._live_embedders.values())
        streaming = self._streaming if streaming is None else streaming
        delay = self._delay if delay is None else delay
        deadline = self._deadline if deadline is None else deadline
        http_attempts = (
            self._http_attempts if http_attempts is None else http_attempts
        )
        t0 = self._t0 if t0 is None else t0

        left = max(0.0, deadline - time.perf_counter())
        tgt = self.total_answer_tokens
        cur = (
            len(finished) * tgt
            + sum(live_tokens.values())
            + sum(n for _emb, n in finalizing)
            + sum(len(p.token_ids) for p in idle)
        )
        total = N_SAMPLES * tgt
        fin_embs = [e for e, _n in finalizing]
        emb_done, emb_total = self._embed_progress_from(
            finished, live_embs, fin_embs, idle
        )
        sty_done, sty_total = self._stylo_progress_from(
            finished, live_embs, fin_embs, idle
        )
        # Peaks only while NAI is still running (hide soft-cut dips). During
        # score-only (st=0, w>0) show **live** e/s so a hung batch is obvious
        # instead of a frozen peak like e24k/33k forever.
        score_only = (
            streaming == 0
            and len(finalizing) > 0
            and len(idle) == 0
            and (len(finished) + len(finalizing) >= N_SAMPLES)
        )
        if not score_only:
            pe_d, pe_t = self._emb_peak
            self._emb_peak = (max(pe_d, emb_done), max(pe_t, emb_total))
            ps_d, ps_t = self._sty_peak
            self._sty_peak = (max(ps_d, sty_done), max(ps_t, sty_total))
            emb_done, emb_total = self._emb_peak
            sty_done, sty_total = self._sty_peak
        elapsed = time.perf_counter() - t0
        http_s = (
            f"{http_attempts}/{MAX_HTTP_ATTEMPTS}"
            if MAX_HTTP_ATTEMPTS is not None
            else str(http_attempts)
        )
        # Spawn-gate countdown only matters when we may still start a NovelAI
        # stream. When only emb/stylo drain, show "score" not 48/128s.
        may_spawn_nai = streaming > 0 or len(idle) > 0 or (
            len(finished) + len(finalizing) + len(idle) + streaming < N_SAMPLES
        )
        if may_spawn_nai:
            pace = f"{left:.0f}/{delay:.0f}s"
        else:
            pace = "score"
        msg = (
            f"{self._progress_prefix}"
            f"i{len(idle)}->"
            f"{pace}->"
            f"st{streaming}->"
            f"w{len(finalizing)}->"
            f"{len(finished)}/{N_SAMPLES} "
            f"tok{_fmt_count(cur)}/{_fmt_count(total)} "
            f"e{_fmt_count(emb_done)}/{_fmt_count(emb_total)} "
            f"s{_fmt_count(sty_done)}/{_fmt_count(sty_total)} "
            f"h{http_s} {elapsed:.0f}s"
        )
        if note:
            msg = f"{msg} {note}"
        return msg

    def _ui(self, *, force: bool = False, note: str = "") -> None:
        """In-place progress line (throttled unless force).

        Snapshots under ``_cv`` then builds the message **outside** the lock so
        multi-minute stylo waits cannot freeze the spawn loop / whole TTY by
        holding the pool condition while scanning thousands of futures.
        """
        prog = self._progress
        if prog is None or self._stop:
            return
        now = time.perf_counter()
        if not force and (now - self._last_ui) < 0.25:
            return
        self._last_ui = now
        with self._cv:
            if self._stop:
                return
            finished = list(self._finished)
            live_tokens = dict(self._live_tokens)
            finalizing = list(self._finalizing)
            idle = list(self._idle)
            live_embs = list(self._live_embedders.values())
            streaming = self._streaming
            delay = self._delay
            deadline = self._deadline
            http_attempts = self._http_attempts
            t0 = self._t0
        msg = self._progress_msg(
            note=note,
            finished=finished,
            live_tokens=live_tokens,
            finalizing=finalizing,
            idle=idle,
            live_embs=live_embs,
            streaming=streaming,
            delay=delay,
            deadline=deadline,
            http_attempts=http_attempts,
            t0=t0,
        )
        prog = self._progress
        if prog is not None and not self._stop:
            prog.update(msg)

    def _arm_spawn_cooldown(self) -> None:
        """Gate the *next* HTTP start by the current delay (must hold _cv)."""
        self._deadline = time.perf_counter() + float(self._delay)

    def _shutdown(self) -> None:
        """Stop spawning, abandon embedders (cancel queued stylo/emb), free UI.

        Abandons every embedder ever created this trial (not only live/idle/
        finalizing). Workers that hit hard errors used to drop out of those
        lists without cancel → orphan pool work and ``s0`` for minutes.
        """
        with self._cv:
            self._stop = True
            self._progress = None
            emb_list = list(self._owned_embedders)
            self._owned_embedders.clear()
            self._live_embedders.clear()
            self._finalizing.clear()
            self._idle.clear()
            self._cv.notify_all()
        for emb in emb_list:
            self._safe_abandon(emb)

    def _budget_ok(self) -> None:
        """Raise if this trial has exhausted HTTP or wall-clock budget."""
        # Caller should hold _cv for consistent reads (optional).
        if MAX_HTTP_ATTEMPTS is not None and self._http_attempts >= MAX_HTTP_ATTEMPTS:
            raise RuntimeError(
                f"trial HTTP budget exhausted "
                f"({self._http_attempts}/{MAX_HTTP_ATTEMPTS} stream attempts, "
                f"done={len(self._finished)}/{N_SAMPLES})"
            )
        elapsed = time.perf_counter() - self._t0
        if elapsed >= MAX_TRIAL_SECONDS:
            raise RuntimeError(
                f"trial time budget exhausted "
                f"({elapsed:.0f}s/{MAX_TRIAL_SECONDS}s, "
                f"done={len(self._finished)}/{N_SAMPLES})"
            )

    def run(self) -> list[GeneratedSample]:
        try:
            with self._cv:
                self._delay = BASE_DELAY
                self._t0 = time.perf_counter()
                self._http_attempts = 0
                self._arm_spawn_cooldown()
            self._ui(force=True, note="starting")

            while True:
                spawn_job: tuple[LiveChunkEmbedder, int] | None = None
                with self._cv:
                    if self._error is not None:
                        raise self._error
                    if len(self._finished) >= N_SAMPLES:
                        return list(self._finished[:N_SAMPLES])
                    try:
                        self._budget_ok()
                    except RuntimeError as e:
                        self._error = e
                        raise

                    now = time.perf_counter()
                    # Stream slots: only live NovelAI POSTs (not embed finalize).
                    stream_slots = MAX_CONCURRENT_STREAMS - self._streaming
                    in_flight = (
                        len(self._finished)
                        + self._streaming
                        + len(self._finalizing)
                        + len(self._idle)
                    )
                    # Resume idle first (keeps tokens). New blanks only when
                    # idle is empty and total samples < N_SAMPLES.
                    can_resume = bool(self._idle) and stream_slots > 0
                    can_new = (
                        not self._idle
                        and stream_slots > 0
                        and in_flight < N_SAMPLES
                    )
                    if not can_resume and not can_new:
                        self._cv.wait(timeout=0.1)
                    elif now < self._deadline:
                        self._cv.wait(
                            timeout=min(0.1, max(0.0, self._deadline - now))
                        )
                    else:
                        # Budget again right before spend (HTTP attempt).
                        try:
                            self._budget_ok()
                        except RuntimeError as e:
                            self._error = e
                            raise
                        if self._idle:
                            emb = self._idle.pop(0)
                        else:
                            emb = self._new_embedder()
                        self._streaming += 1
                        self._http_attempts += 1
                        # CRITICAL: arm countdown for the *next* spawn.
                        self._arm_spawn_cooldown()
                        wid = self._next_wid
                        self._next_wid += 1
                        self._live_tokens[wid] = len(emb.token_ids)
                        self._live_embedders[wid] = emb
                        self._cv.notify_all()
                        spawn_job = (emb, wid)

                self._ui(force=spawn_job is not None)
                if spawn_job is not None:
                    emb, wid = spawn_job
                    try:
                        threading.Thread(
                            target=self._worker, args=(emb, wid), daemon=True
                        ).start()
                    except Exception as e:
                        # Thread failed to start — put sample back on idle
                        # (never wipe tokens that were already generated).
                        with self._cv:
                            self._streaming = max(0, self._streaming - 1)
                            self._live_tokens.pop(wid, None)
                            self._live_embedders.pop(wid, None)
                            self._park_sample(emb, list(emb.token_ids))
                            self._cv.notify_all()
                        self._ui(
                            force=True,
                            note=f"spawn-err:{type(e).__name__}",
                        )
        finally:
            self._shutdown()

    def _reset_delay(self) -> None:
        """On first token: restore base spacing only — never clear deadline."""
        with self._cv:
            if self._stop:
                return
            self._delay = BASE_DELAY
            # Do not clear _deadline — spawn path owns request pacing.
            self._cv.notify_all()

    def _bump_fail_delay(self) -> None:
        """Escalate spawn delay after a soft-cut / empty (caller holds _cv)."""
        wait = float(self._delay)
        self._deadline = max(self._deadline, time.perf_counter() + wait)
        if wait >= MAX_DELAY:
            self._delay = BASE_DELAY
        else:
            self._delay = min(wait * 2.0, MAX_DELAY)

    def _park_sample(
        self,
        embedder: LiveChunkEmbedder | None,
        token_ids: list[int],
    ) -> None:
        """Park sample on idle — **always reuse** embedder (keeps emb/stylo).

        Never creates a replacement embedder when one is provided: that would
        drop ready caches / futures and look like e/s wipe. Only a brand-new
        blank sample (embedder is None) may construct a new LiveChunkEmbedder.
        """
        ids = list(token_ids)
        if embedder is None:
            self._idle.append(self._new_embedder(token_ids=ids))
        else:
            # If somehow killed, still do not invent a second life that zeros
            # UI by dropping the dead object without counting it — re-queue on
            # the same object is impossible when pool_dead; only then new.
            if embedder._pool_dead:
                # Sample metrics already dead with the sample; tokens-only
                # rebirth is allowed only because kill already happened.
                self._idle.append(self._new_embedder(token_ids=ids))
            else:
                # Allow set_tokens even after a successful finalize so a
                # parked re-score path can still hold the same ready caches.
                if embedder._finalized and embedder._final_result is not None:
                    # Already fully scored — keep object as-is on idle for
                    # resume→done→return cached finalize result.
                    pass
                else:
                    embedder.set_tokens(ids)
                self._idle.append(embedder)
        self._bump_fail_delay()

    def _release_empty_attempt(
        self, embedder: LiveChunkEmbedder | None
    ) -> None:
        """True empty stream (0 tokens): sample dies, metrics die with it."""
        self._kill_sample_metrics(embedder)
        self._bump_fail_delay()

    def _set_live(self, wid: int, n: int) -> None:
        if self._stop:
            return
        with self._cv:
            if self._stop:
                return
            if wid in self._live_tokens:
                self._live_tokens[wid] = n
        self._ui()

    def _release_stream_live(self, wid: int) -> None:
        """Drop live stream bookkeeping (caller holds _cv)."""
        self._streaming = max(0, self._streaming - 1)
        self._live_tokens.pop(wid, None)
        self._live_embedders.pop(wid, None)

    def _drop_finalizing(self, embedder: LiveChunkEmbedder) -> None:
        """Remove one finalizing entry by embedder identity (caller holds _cv)."""
        for i, (emb, _n) in enumerate(self._finalizing):
            if emb is embedder:
                self._finalizing.pop(i)
                return

    def _swap_finalizing(
        self, old: LiveChunkEmbedder, new: LiveChunkEmbedder, n_tok: int
    ) -> None:
        """Replace finalizing embedder (metric re-score) under lock."""
        with self._cv:
            self._drop_finalizing(old)
            if not self._stop:
                self._finalizing.append((new, n_tok))
            self._cv.notify_all()

    def _finalize_generated(
        self,
        embedder: LiveChunkEmbedder,
        *,
        text: str,
        tokens: list[int],
    ) -> tuple[GeneratedSample | None, LiveChunkEmbedder]:
        """Wait emb+stylo; retry forever on failure until stop (keep e/s progress).

        Without scores the trial cannot finish — park/give-up is useless — so
        keep re-submitting only missing/failed emb+stylo windows on the same
        embedder. Ready caches keep completed windows. Returns None only when
        the trial is stopping or tokens are empty.
        """
        active = embedder

        def _tick() -> None:
            # Keep the progress line alive during multi-minute emb/stylo drains.
            self._ui(note="fin-wait")

        attempt = 0
        while True:
            if self._stop or not tokens:
                return None, active
            attempt += 1
            try:
                emb_32, emb_256, metrics = active.finalize(on_tick=_tick)
                active.score_fail_count = 0
                return (
                    GeneratedSample(
                        text=text,
                        emb_32=emb_32,
                        emb_256=emb_256,
                        metrics=metrics,
                    ),
                    active,
                )
            except Exception as emb_err:
                active.score_fail_count = int(
                    getattr(active, "score_fail_count", 0)
                ) + 1
                self._ui(
                    force=True,
                    note=f"fin-err{attempt}:{type(emb_err).__name__}",
                )
                # Back off so permanent local bugs do not tight-loop CPU.
                delay = min(30.0, _EMBED_RETRY_BASE_S * (2.0 ** min(attempt - 1, 5)))
                self._ui(force=True, note=f"fin-retry{attempt + 1}")
                time.sleep(delay)

    def _worker(self, embedder: LiveChunkEmbedder, wid: int) -> None:
        # Track whether this worker still holds a stream / finalize slot so
        # exception paths can release cleanly without double-decrement.
        holds_stream = True
        in_finalizing = False
        try:
            if self._stop:
                with self._cv:
                    self._release_stream_live(wid)
                    holds_stream = False
                    self._cv.notify_all()
                # Trial stopping — cancel pool work (not a mid-trial wipe).
                self._safe_abandon(embedder)
                return
            clean_partial, _dropped = _sanitize_token_ids(embedder.token_ids)
            embedder.set_tokens(clean_partial)
            self._set_live(wid, len(clean_partial))
            # One phase: stream until done / partial / empty (soft-cut park).
            outcome = "empty"
            text = ""
            tokens: list[int] = list(clean_partial)
            outcome, text, tokens = self._one_stream(embedder, wid)

            # --- Generation over: free the NovelAI stream slot immediately ---
            # Policy: if this sample has any tokens, KEEP them (finalize or
            # park). Never blank-replace. Only true 0-token attempts free the
            # slot without parking.
            if outcome == "done" and tokens:
                if not (text or "").strip():
                    text = self._decode(tokens) or "[decode-empty]"
                with self._cv:
                    self._release_stream_live(wid)
                    holds_stream = False
                    if self._stop:
                        self._cv.notify_all()
                        self._safe_abandon(embedder)
                        return
                    self._finalizing.append((embedder, len(tokens)))
                    in_finalizing = True
                    self._cv.notify_all()
                self._ui(force=True, note="fin")

                sample, active = self._finalize_generated(
                    embedder, text=text, tokens=tokens
                )
                embedder = active

                with self._cv:
                    self._drop_finalizing(embedder)
                    in_finalizing = False
                    if self._stop:
                        if sample is None:
                            self._safe_abandon(embedder)
                        self._cv.notify_all()
                        return
                    if sample is None:
                        # Only reached on trial stop / empty tokens (finalize
                        # retries emb+stylo forever otherwise). Park if any
                        # tokens remain so a later resume keeps generation.
                        if tokens:
                            self._park_sample(embedder, tokens)
                        else:
                            self._safe_abandon(embedder)
                    elif len(self._finished) < N_SAMPLES:
                        self._finished.append(sample)
                        self._delay = BASE_DELAY
                        # Matrices already copied into GeneratedSample — free
                        # ready-cache RAM under the script budget.
                        try:
                            embedder.abandon()
                        except Exception:
                            pass
                    else:
                        # Already have enough finished samples; drop surplus
                        # only after a successful score (not a wipe-for-retry).
                        self._safe_abandon(embedder)
                    self._cv.notify_all()
                self._ui(force=True)
                return

            # partial / empty
            with self._cv:
                self._release_stream_live(wid)
                holds_stream = False
                if self._stop:
                    self._cv.notify_all()
                    self._safe_abandon(embedder)
                    return
                keep_ids = list(tokens) if tokens else list(embedder.token_ids)
                if keep_ids:
                    self._park_sample(embedder, keep_ids)
                else:
                    self._release_empty_attempt(embedder)
                self._cv.notify_all()
            self._ui(force=True)
        except BaseException as e:
            # Keep tokens on soft/recoverable errors. Hard errors (bad HTTP 4xx
            # that isn't soft-cut, programming bugs) fail the trial — previously
            # ALL exceptions soft-parked and a broken trial could run forever.
            keep_ids: list[int] = []
            try:
                keep_ids = list(embedder.token_ids)
            except Exception:
                keep_ids = []
            hard = _is_hard_trial_error(e)
            with self._cv:
                if holds_stream:
                    self._release_stream_live(wid)
                    holds_stream = False
                if in_finalizing:
                    self._drop_finalizing(embedder)
                    in_finalizing = False
                if self._stop or isinstance(e, (KeyboardInterrupt, SystemExit)):
                    self._safe_abandon(embedder)
                    if self._error is None:
                        self._error = e
                elif hard:
                    # Preserve generation if any, but end the trial.
                    if keep_ids:
                        self._park_sample(embedder, keep_ids)
                    else:
                        self._release_empty_attempt(embedder)
                    if self._error is None:
                        self._error = e
                    self._ui(
                        force=True,
                        note=f"hard-err:{type(e).__name__}",
                    )
                elif keep_ids:
                    self._park_sample(embedder, keep_ids)
                    self._ui(
                        force=True,
                        note=f"soft-err:{type(e).__name__}",
                    )
                else:
                    self._release_empty_attempt(embedder)
                self._cv.notify_all()

    def _stream_prompt_and_cap(
        self, out_ids: list[int]
    ) -> tuple[list[int], int, str]:
        """Build (prompt, absolute_token_cap, finish_kind) for one-phase gen.

        prompt = (system + prefill) + AI_partial
        cap = answer_tokens; finish_kind = "done" at full AI length
        """
        total = self.total_answer_tokens
        return list(self.prompt_ids) + list(out_ids), total, "done"

    def _one_stream(
        self, embedder: LiveChunkEmbedder, wid: int
    ) -> tuple[str, str, list[int]]:
        """Single POST stream → (done|partial|empty, text, token_ids)."""
        out_ids, _dropped = _sanitize_token_ids(embedder.token_ids)
        embedder.set_tokens(out_ids)
        if len(out_ids) >= self.total_answer_tokens:
            return "done", self._decode(out_ids), out_ids

        prompt, token_cap, finish_kind = self._stream_prompt_and_cap(out_ids)
        remaining = token_cap - len(out_ids)
        if remaining <= 0:
            return finish_kind, self._decode(out_ids), out_ids

        body: dict[str, Any] = {
            "model": self.model,
            "prompt": prompt,
            "max_tokens": remaining,
            "stream": True,
            **self.preset.to_api(),
            "logit_bias": self.token_bias,
        }
        headers = {"Authorization": f"Bearer {self.token}"}
        saw_token = False
        last_ui_n = len(out_ids)

        def _finish(kind: str) -> tuple[str, str, list[int]]:
            embedder.set_tokens(out_ids)
            return kind, self._decode(out_ids), out_ids

        def _on_new_tokens() -> None:
            nonlocal saw_token, last_ui_n
            # Fire embeds for any newly completed full windows immediately.
            embedder.set_tokens(out_ids)
            if self._stop:
                return
            if not saw_token:
                saw_token = True
                self._reset_delay()
            if len(out_ids) - last_ui_n >= 25:
                last_ui_n = len(out_ids)
                self._set_live(wid, len(out_ids))

        try:
            with requests.post(
                NAI_URL, json=body, headers=headers, timeout=120.0, stream=True
            ) as resp:
                if self._stop:
                    return _finish("partial" if out_ids else "empty")
                if resp.status_code == 429 or resp.status_code >= 500:
                    return _finish("partial" if out_ids else "empty")
                if 400 <= resp.status_code < 500:
                    err_body = resp.text[:400]
                    # Bad ids: soft-cut keep what we have (never return empty
                    # with a non-empty out_ids — that wiped tok at low delay).
                    if out_ids and (
                        "token beyond max" in err_body
                        or "beyond max token" in err_body
                    ):
                        return _finish("partial")
                    raise RuntimeError(
                        f"NovelAI HTTP {resp.status_code}: {err_body}"
                    )
                if resp.status_code != 200:
                    return _finish("partial" if out_ids else "empty")

                for raw in resp.iter_lines(decode_unicode=True):
                    if self._stop:
                        return _finish("partial" if out_ids else "empty")
                    payload = _parse_sse_line(raw or "")
                    if payload is None:
                        continue
                    if payload == "[DONE]":
                        break
                    try:
                        obj = json.loads(payload)
                        if not isinstance(obj, dict):
                            continue
                        choices = obj.get("choices")
                        if not isinstance(choices, list) or not choices:
                            continue
                        choice = choices[0]
                        if not isinstance(choice, dict):
                            continue
                        # API sometimes sends text as null/non-str; only real
                        # strings are used (avoids tokenizer TypeError later).
                        raw_text = choice.get("text")
                        text_delta = raw_text if isinstance(raw_text, str) else ""
                        delta_ids = _coerce_token_ids(choice.get("token_ids"))
                    except (ValueError, KeyError, IndexError, TypeError):
                        continue

                    if delta_ids:
                        take = delta_ids[: max(0, token_cap - len(out_ids))]
                        clean, oob = _sanitize_token_ids(take)
                        if oob:
                            if clean:
                                out_ids.extend(clean)
                            return _finish("partial" if out_ids else "empty")
                        out_ids.extend(clean)
                        try:
                            _on_new_tokens()
                        except Exception:
                            # Embed submit / decode glitch mid-stream: keep tokens,
                            # stop this POST as partial so countdown can retry.
                            return _finish("partial" if out_ids else "empty")
                        if len(out_ids) >= token_cap:
                            self._set_live(wid, len(out_ids))
                            return _finish(finish_kind)
                    elif text_delta:
                        try:
                            approx = self._ids_from_text(text_delta)
                        except Exception:
                            approx = []
                        room = token_cap - len(out_ids)
                        take = approx[: max(0, room)]
                        if take:
                            out_ids.extend(take)
                            try:
                                _on_new_tokens()
                            except Exception:
                                return _finish("partial" if out_ids else "empty")
                            if len(out_ids) >= token_cap:
                                self._set_live(wid, len(out_ids))
                                return _finish(finish_kind)
                        elif not saw_token and not self._stop:
                            saw_token = True
                            self._reset_delay()

        except requests.RequestException:
            return _finish("partial" if out_ids else "empty")

        if len(out_ids) >= token_cap:
            return _finish(finish_kind)
        if out_ids:
            return _finish("partial")
        return _finish("empty")


def generate_five(
    *,
    token: str,
    exam: ExamBundle,
    preset: Preset,
    tokenizer: Tokenizer,
    embed_client: OpenAI,
    progress: LineProgress | None = None,
    progress_prefix: str = "",
) -> list[GeneratedSample]:
    """Generate N_SAMPLES full-answer completions; live emb + stylo.

    Prompt = system + prefill; stream answer_tokens.
    Per sample: live 32/256 emb windows + live 256-tok stylo metric windows;
    finalize waits on remaining futures after the stream ends.

    On KeyboardInterrupt / error: embed pool is shut down with wait=False so
    Ctrl+C is not stuck joining in-flight HTTP / spaCy workers (the default
    ``with ThreadPoolExecutor`` path uses wait=True and feels unkillable).
    """
    embed_pool = ThreadPoolExecutor(
        max_workers=EMBED_POOL_WORKERS, thread_name_prefix="embed"
    )
    nai: NaiPool | None = None
    try:
        nai = NaiPool(
            token=token,
            model=exam.config.model,
            prompt_ids=exam.prompt_ids,
            answer_tokens=exam.config.answer_tokens,
            token_bias=exam.config.token_bias,
            preset=preset,
            tokenizer=tokenizer,
            embed_client=embed_client,
            embed_pool=embed_pool,
            progress=progress,
            progress_prefix=progress_prefix,
        )
        return nai.run()
    except BaseException:
        if nai is not None:
            try:
                nai._shutdown()
            except Exception:
                pass
        raise
    finally:
        try:
            embed_pool.shutdown(wait=False, cancel_futures=True)
        except TypeError:
            try:
                embed_pool.shutdown(wait=False)
            except Exception:
                pass
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Human baseline cache (external file — NOT presets.db / Optuna user_attrs)
# ---------------------------------------------------------------------------

# Keys formerly stored in study.user_attrs; purged on startup so SQLite stays lean.
_HUMAN_STUDY_ATTR_KEYS = (
    "metrics_human",
    "human_emb_32",
    "human_emb_256",
    "metrics_version",
    "chunk_32",
    "chunk_256",
    "embed_dim",
    "study_key",
    "toml_hash",
    "exam_hash",
    "n_samples",
)

# All durable caches live under cache/ (human baselines, tokenizer, …).
CACHE_DIR = ROOT / "cache"
TOKENIZER_PATH = CACHE_DIR / "tokenizer.json"


def human_cache_path(exam: ExamBundle) -> Path:
    """One compressed npz per study_key, same level as tokenizer in cache/."""
    return CACHE_DIR / f"{exam.study_key}.npz"


def _valid_chunk_emb(arr: np.ndarray) -> bool:
    return (
        arr.ndim == 2
        and arr.shape[0] >= 1
        and arr.shape[1] == EMBED_DIM
        and np.all(np.isfinite(arr))
    )


def _load_human_cache_file(path: Path, exam: ExamBundle) -> HumanCache | None:
    """Load and validate external human cache; None if missing/stale/corrupt."""
    if not path.is_file():
        return None
    n_feat = len(FEATURE_NAMES)
    try:
        with np.load(path, allow_pickle=False) as z:
            keys = set(z.files)
            need = {
                "metrics",
                "emb_32",
                "emb_256",
                "study_key",
                "toml_hash",
                "exam_hash",
                "metrics_version",
                "chunk_32",
                "chunk_256",
                "chunk_32_stride",
                "chunk_256_stride",
                "stylo_chunk",
                "stylo_stride",
                "embed_dim",
                "n_samples",
            }
            if not need.issubset(keys):
                return None

            def _s(name: str) -> str:
                v = z[name]
                if v.shape == ():
                    return str(v.item())
                return str(v)

            if (
                _s("study_key") != exam.study_key
                or _s("toml_hash") != exam.toml_hash
                or _s("exam_hash") != exam.exam_hash
                or _s("metrics_version") != METRICS_VERSION
                or int(z["chunk_32"]) != CHUNK_32
                or int(z["chunk_256"]) != CHUNK_256
                or int(z["chunk_32_stride"]) != CHUNK_32_STRIDE
                or int(z["chunk_256_stride"]) != CHUNK_256_STRIDE
                or int(z["stylo_chunk"]) != STYLO_CHUNK
                or int(z["stylo_stride"]) != STYLO_STRIDE
                or int(z["embed_dim"]) != EMBED_DIM
                or int(z["n_samples"]) != exam.config.n_samples
            ):
                return None

            metrics_h = np.asarray(z["metrics"], dtype=float)
            emb_32 = np.asarray(z["emb_32"], dtype=float)
            emb_256 = np.asarray(z["emb_256"], dtype=float)
    except Exception as e:
        print(f"Human cache unreadable ({path.name}): {type(e).__name__}: {e}")
        return None

    if not (
        metrics_h.ndim == 2
        and metrics_h.shape[0] >= 1
        and metrics_h.shape[1] == n_feat
        and np.all(np.isfinite(metrics_h))
        and _valid_chunk_emb(emb_32)
        and _valid_chunk_emb(emb_256)
        and metrics_h.shape[0] == emb_256.shape[0]
    ):
        return None
    return HumanCache(metrics=metrics_h, emb_32=emb_32, emb_256=emb_256)


def _save_human_cache_file(path: Path, exam: ExamBundle, human: HumanCache) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    # np.savez_compressed appends .npz if the path does not already end with it.
    tmp = path.parent / f"{path.stem}.part.npz"
    np.savez_compressed(
        tmp,
        metrics=np.asarray(human.metrics, dtype=np.float64),
        emb_32=np.asarray(human.emb_32, dtype=float),
        emb_256=np.asarray(human.emb_256, dtype=float),
        study_key=np.asarray(exam.study_key),
        toml_hash=np.asarray(exam.toml_hash),
        exam_hash=np.asarray(exam.exam_hash),
        metrics_version=np.asarray(METRICS_VERSION),
        chunk_32=np.asarray(CHUNK_32, dtype=np.int32),
        chunk_256=np.asarray(CHUNK_256, dtype=np.int32),
        chunk_32_stride=np.asarray(CHUNK_32_STRIDE, dtype=np.int32),
        chunk_256_stride=np.asarray(CHUNK_256_STRIDE, dtype=np.int32),
        stylo_chunk=np.asarray(STYLO_CHUNK, dtype=np.int32),
        stylo_stride=np.asarray(STYLO_STRIDE, dtype=np.int32),
        embed_dim=np.asarray(EMBED_DIM, dtype=np.int32),
        n_samples=np.asarray(exam.config.n_samples, dtype=np.int32),
    )
    tmp.replace(path)


def purge_human_attrs_from_presets_db(db_path: Path) -> int:
    """
    Delete legacy human-baseline keys from Optuna's study_user_attributes.

    Returns number of rows deleted. Does not VACUUM (caller may).
    """
    if not db_path.is_file():
        return 0
    import sqlite3

    placeholders = ",".join("?" * len(_HUMAN_STUDY_ATTR_KEYS))
    sql = (
        f"DELETE FROM study_user_attributes WHERE key IN ({placeholders})"
    )
    try:
        with sqlite3.connect(str(db_path)) as con:
            cur = con.execute(sql, _HUMAN_STUDY_ATTR_KEYS)
            n = int(cur.rowcount if cur.rowcount is not None else 0)
            con.commit()
            return max(0, n)
    except sqlite3.Error as e:
        print(f"warning: could not purge human attrs from {db_path.name}: {e}")
        return 0


def _hard_delete_trials_by_id(
    db_path: Path, trial_ids: list[int]
) -> int:
    """Delete trials rows + dependent tables for the given trial_ids."""
    if not db_path.is_file() or not trial_ids:
        return 0
    import sqlite3

    n = 0
    try:
        with sqlite3.connect(str(db_path)) as con:
            for trial_id in trial_ids:
                for table in (
                    "trial_params",
                    "trial_values",
                    "trial_user_attributes",
                    "trial_system_attributes",
                    "trial_intermediate_values",
                    "trial_heartbeats",
                ):
                    con.execute(
                        f"DELETE FROM {table} WHERE trial_id = ?", (trial_id,)
                    )
                cur = con.execute(
                    "DELETE FROM trials WHERE trial_id = ?", (trial_id,)
                )
                n += int(cur.rowcount or 0)
            con.commit()
    except sqlite3.Error as e:
        print(f"warning: could not hard-delete trials from {db_path.name}: {e}")
        return 0
    return n


def purge_fail_running_trials(db_path: Path) -> list[tuple[int, str]]:
    """
    Hard-delete FAIL and RUNNING trials from an Optuna SQLite DB.

    Incomplete / crashed workers leave RUNNING; catch=(Exception,) can leave
    FAIL. Neither helps the sampler and both clutter the dashboard — drop them (and
    dependent trial_* rows). Returns list of (trial_number, state) removed.
    """
    if not db_path.is_file():
        return []
    import sqlite3

    try:
        with sqlite3.connect(str(db_path)) as con:
            rows = con.execute(
                "SELECT trial_id, number, state FROM trials "
                "WHERE state IN ('FAIL', 'RUNNING')"
            ).fetchall()
    except sqlite3.Error as e:
        print(f"warning: could not purge FAIL/RUNNING from {db_path.name}: {e}")
        return []
    if not rows:
        return []
    ids = [int(tid) for tid, _, _ in rows]
    removed = [(int(number), str(state)) for _, number, state in rows]
    if _hard_delete_trials_by_id(db_path, ids) <= 0:
        return []
    return removed


def ensure_human_cache(
    exam: ExamBundle,
    embed_client: OpenAI,
    tokenizer: Tokenizer,
    *,
    db_path: Path | None = None,
) -> HumanCache:
    """
    Load or compute human baselines.

    Stored in cache/<study_key>.npz — never in presets.db / Optuna
    user_attrs (embeddings are huge and bloat SQLite).
    """
    n_feat = len(FEATURE_NAMES)
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    path = human_cache_path(exam)

    if db_path is not None:
        purged = purge_human_attrs_from_presets_db(db_path)
        if purged:
            print(
                f"Purged {purged} legacy human-cache row(s) from {db_path.name} "
                f"(now using {path.name})"
            )

    cached = _load_human_cache_file(path, exam)
    if cached is not None:
        print(
            f"human cache hit {path.name} "
            f"sty{cached.metrics.shape[0]} "
            f"e32{cached.emb_32.shape[0]} e256{cached.emb_256.shape[0]}"
        )
        return cached
    if path.is_file():
        print(f"human cache stale ({path.name}) — recompute")
    else:
        print(f"human cache miss ({path.name}) — compute")

    n32 = len(token_chunks(exam.answer_text, tokenizer, CHUNK_32))
    n256 = len(token_chunks(exam.answer_text, tokenizer, CHUNK_256))
    print(f"computing human: sty{n256} e32{n32} e256{n256} …")
    human = compute_human_cache(exam.answer_text, embed_client, tokenizer)
    _save_human_cache_file(path, exam, human)
    size_mb = path.stat().st_size / (1024 * 1024) if path.is_file() else 0.0
    print(
        f"  saved {path.name} {size_mb:.1f}MiB "
        f"sty{human.metrics.shape[0]} "
        f"e32{human.emb_32.shape[0]} e256{human.emb_256.shape[0]}"
    )
    return human


def make_objective(
    exam: ExamBundle,
    nai_token: str,
    embed_client: OpenAI,
    human: HumanCache,
    tokenizer: Tokenizer,
    obj_index: int,
):
    """Build Optuna objective that returns the single scalar for ``obj_index``."""

    def objective(trial: optuna.Trial) -> float:
        cfg = exam.config
        tk = SEARCH_SPACE["top_k"]
        tp = SEARCH_SPACE["top_p"]
        temp = SEARCH_SPACE["temperature"]

        # Only free knobs are suggest_*'d (keeps AutoSampler on numerical path).
        # Values present in exam.toml are fixed constants — never Optuna params.
        if cfg.fixed_temperature is not None:
            temperature = float(cfg.fixed_temperature)
        else:
            temperature = trial.suggest_float(
                "temperature", temp.low, temp.high, log=temp.log
            )

        if cfg.logit_tolerance is not None:
            logit_tolerance = float(cfg.logit_tolerance)
        else:
            w = SEARCH_SPACE[PARAM_LOGIT_TOLERANCE]
            logit_tolerance = trial.suggest_float(
                PARAM_LOGIT_TOLERANCE, w.low, w.high
            )

        if cfg.fixed_top_k is not None:
            top_k = int(cfg.fixed_top_k)
        else:
            top_k = trial.suggest_int("top_k", tk.low, tk.high, log=tk.log)

        if cfg.fixed_top_p is not None:
            top_p = float(cfg.fixed_top_p)
        else:
            top_p = trial.suggest_float("top_p", tp.low, tp.high)

        min_p = min_p_from_temperature(temperature, logit_tolerance)
        trial.set_user_attr(ATTR_LOGIT_TOLERANCE, logit_tolerance)
        trial.set_user_attr(ATTR_MIN_P, min_p)
        # Mirror pinned knobs as user_attrs so dashboard / logs still show them.
        if cfg.fixed_top_k is not None:
            trial.set_user_attr("top_k", top_k)
        if cfg.fixed_top_p is not None:
            trial.set_user_attr("top_p", top_p)
        if cfg.fixed_temperature is not None:
            trial.set_user_attr("temperature", temperature)
        preset = Preset(
            top_k=top_k,
            min_p=min_p,
            top_p=top_p,
            temperature=temperature,
        )
        prefix = _trial_prefix(
            trial.number,
            obj_index,
            top_k=preset.top_k,
            top_p=preset.top_p,
            temperature=preset.temperature,
        )
        prog = LineProgress()
        try:
            samples = generate_five(
                token=nai_token,
                exam=exam,
                preset=preset,
                tokenizer=tokenizer,
                embed_client=embed_client,
                progress=prog,
                progress_prefix=prefix,
            )
            if len(samples) != N_SAMPLES:
                raise RuntimeError(
                    f"expected {N_SAMPLES} samples, got {len(samples)}"
                )
            # Embeds + stylo windows were extracted live during streaming;
            # scoring aggregates already-finished matrices.
            n_win = int(np.median([s.metrics.shape[0] for s in samples]))
            prog.update(f"{prefix}score {n_win}w×{len(samples)}")
            result = score_answers(samples, human)
            d_s, d_32, d_256 = result.score_std
            # Always store all three raw components so peers can re-map.
            trial.set_user_attr(ATTR_STYLO, result.stylo)
            trial.set_user_attr(ATTR_EMB_32, result.emb_32)
            trial.set_user_attr(ATTR_EMB_256, result.emb_256)
            value = float(result.values[obj_index])
            # Compact end: obj value, raw sty/e256/e32, per-sample stds.
            prog.end(
                f"{prefix}"
                f"{_obj_short(obj_index)}={value:.4f} "
                f"sty:{result.stylo:.3f} e256:{result.emb_256:.3f} "
                f"e32:{result.emb_32:.3f} "
                f"σ{d_s:.3f}/{d_256:.3f}/{d_32:.3f}"
            )
            del samples
            gc.collect()
            return value
        except KeyboardInterrupt:
            try:
                prog.end(f"{prefix}interrupted")
            except Exception:
                pass
            try:
                shutdown_metric_pool(wait=False)
            except Exception:
                pass
            raise
        except Exception as e:
            # Never TrialPruned — that left PRUNED rows and hid real failures.
            # Re-raise so Optuna records FAIL (purged like RUNNING) or bubbles up.
            err = str(e)
            if len(err) > 60:
                err = err[:57] + "…"
            prog.end(f"{prefix}err {type(e).__name__}: {err}")
            tb = "".join(
                traceback.format_exception(type(e), e, e.__traceback__, limit=6)
            )
            sys.stderr.write(tb)
            sys.stderr.flush()
            gc.collect()
            raise

    return objective


def _study_exists(storage: str, study_name: str) -> bool:
    return any(s.study_name == study_name for s in optuna.get_all_study_summaries(storage))


def shard_study_names(study_key: str) -> list[str]:
    """Optuna study names for the N single-obj studies of one exam fingerprint."""
    return [f"{study_key}{SHARD_SEP}{i}" for i in range(N_STUDY_SHARDS)]


def shard_objective_index(study_name: str) -> int:
    """Parse ``…:i`` → objective index i."""
    try:
        return int(study_name.rsplit(SHARD_SEP, 1)[-1])
    except ValueError as e:
        raise ValueError(f"not a shard study name: {study_name!r}") from e


def _short_study_name(name: str) -> str:
    """Compact label: short key + shard short objective name."""
    if SHARD_SEP in name:
        base, _, suf = name.rpartition(SHARD_SEP)
        if suf.isdigit():
            return f"{base[:8]}:{_obj_short(int(suf))}"
        return f"{base[:8]}:{suf}"
    return f"{name[:8]}…"


def _param_key(params: dict[str, Any]) -> tuple[tuple[str, Any], ...]:
    """Stable fingerprint for uniqueness across studies (exact param values)."""
    return tuple(sorted(params.items()))


def _trial_fixed_or_params(t: optuna.trial.FrozenTrial) -> dict[str, Any] | None:
    """Params for a trial: real params if set, else enqueue fixed_params.

    WAITING trials from ``study.enqueue_trial`` keep suggestions in
    ``system_attrs['fixed_params']`` until the trial is asked; ``t.params``
    is empty until then.
    """
    if t.params:
        return dict(t.params)
    fixed = t.system_attrs.get("fixed_params")
    if isinstance(fixed, dict) and fixed:
        return dict(fixed)
    return None


def drop_waiting_already_complete(
    study: optuna.Study,
    db_path: Path,
) -> list[int]:
    """Delete WAITING trials whose params already exist as COMPLETE on study.

    Cross-study sync only copies COMPLETE rows (see
    ``pull_unique_trials_from_peers``). Sibling WAITING queues are independent
    — without this, the same preset would be fully re-generated on every
    shard. Call after peer pull so imported COMPLETE cancels local WAITING.
    Returns deleted trial numbers.
    """
    complete_keys: set[tuple[tuple[str, Any], ...]] = set()
    for t in study.trials:
        if t.state != TrialState.COMPLETE:
            continue
        p = _trial_fixed_or_params(t)
        if p:
            complete_keys.add(_param_key(p))
    if not complete_keys:
        return []

    # Need storage trial_id for hard delete; map number → id via SQLite.
    import sqlite3

    waiting_numbers: list[int] = []
    for t in study.trials:
        if t.state != TrialState.WAITING:
            continue
        p = _trial_fixed_or_params(t)
        if not p:
            continue
        if _param_key(p) in complete_keys:
            waiting_numbers.append(int(t.number))
    if not waiting_numbers:
        return []

    try:
        with sqlite3.connect(str(db_path)) as con:
            study_id = con.execute(
                "SELECT study_id FROM studies WHERE study_name = ?",
                (study.study_name,),
            ).fetchone()
            if not study_id:
                return []
            sid = int(study_id[0])
            trial_ids: list[int] = []
            for num in waiting_numbers:
                row = con.execute(
                    "SELECT trial_id FROM trials "
                    "WHERE study_id = ? AND number = ? AND state = 'WAITING'",
                    (sid, num),
                ).fetchone()
                if row:
                    trial_ids.append(int(row[0]))
    except sqlite3.Error as e:
        print(
            f"warning: could not resolve WAITING ids for "
            f"{study.study_name}: {e}"
        )
        return []

    if not trial_ids:
        return []
    deleted = _hard_delete_trials_by_id(db_path, trial_ids)
    if deleted:
        return waiting_numbers
    return []


def _trial_raw_components(
    t: optuna.trial.FrozenTrial,
) -> tuple[float, float, float] | None:
    """Recover (stylo, 32 emb, 256 emb) from user_attrs or legacy values."""
    ua = t.user_attrs
    if all(k in ua for k in COMPONENT_ATTR_KEYS):
        return (
            float(ua[ATTR_STYLO]),
            float(ua[ATTR_EMB_32]),
            float(ua[ATTR_EMB_256]),
        )
    # Collapsed single-obj / prior multi-obj used obj1/obj2/obj3.
    if all(k in ua for k in ("obj1", "obj2", "obj3")):
        return float(ua["obj1"]), float(ua["obj2"]), float(ua["obj3"])
    # Very old multi-obj: values were (stylo, 32 emb, 256 emb).
    if t.values is not None and len(t.values) == 3:
        return float(t.values[0]), float(t.values[1]), float(t.values[2])
    return None


def _complete_trial_copy(
    t: optuna.trial.FrozenTrial,
    *,
    obj_index: int,
    imported_from: str | None = None,
) -> optuna.trial.FrozenTrial | None:
    """Build a COMPLETE single-objective FrozenTrial for ``obj_index``.

    Params + component user_attrs are preserved; Optuna ``value`` is that
    study's scalar only. Optional ``imported_from`` tags peer-synced rows.
    """
    if t.state != TrialState.COMPLETE:
        return None
    comps = _trial_raw_components(t)
    if comps is None:
        return None
    stylo, emb_32, emb_256 = comps
    # Cosine emb can be ≤0 (dissimilar windows); only require finite for import.
    if not all(np.isfinite(v) for v in (stylo, emb_32, emb_256)):
        return None
    value = objective_value_from_components(stylo, emb_32, emb_256, obj_index)
    if not np.isfinite(value):
        return None
    if not t.params or not t.distributions:
        return None
    sys_attrs = dict(t.system_attrs)
    sys_attrs.pop(SYS_IMPORTED_FROM, None)
    sys_attrs.pop("seeded", None)
    if imported_from is not None:
        sys_attrs[SYS_IMPORTED_FROM] = imported_from
    # Drop free min_p / legacy bit/nats param names (derived or renamed).
    _drop_params = {ATTR_MIN_P, "min_p_constant", "min_p_bits"}
    params = {k: v for k, v in t.params.items() if k not in _drop_params}
    distributions = {
        k: v for k, v in t.distributions.items() if k not in _drop_params
    }
    if not params or not distributions:
        return None
    user_attrs: dict[str, Any] = {
        ATTR_STYLO: stylo,
        ATTR_EMB_32: emb_32,
        ATTR_EMB_256: emb_256,
    }
    # Prefer stored logit_tolerance; recover from legacy bits/nats names.
    if ATTR_LOGIT_TOLERANCE in t.user_attrs:
        tol_ref = float(t.user_attrs[ATTR_LOGIT_TOLERANCE])
    elif PARAM_LOGIT_TOLERANCE in params:
        tol_ref = float(params[PARAM_LOGIT_TOLERANCE])
    elif "min_p_constant" in t.user_attrs:
        tol_ref = float(t.user_attrs["min_p_constant"])
    elif "min_p_constant" in params:
        tol_ref = float(params["min_p_constant"])
    elif "min_p_bits" in t.user_attrs:
        tol_ref = float(t.user_attrs["min_p_bits"]) * math.log(2.0)
    elif "min_p_bits" in params:
        tol_ref = float(params["min_p_bits"]) * math.log(2.0)
    else:
        tol_ref = LOGIT_TOLERANCE_DEFAULT
    user_attrs[ATTR_LOGIT_TOLERANCE] = tol_ref

    if ATTR_MIN_P in t.user_attrs:
        user_attrs[ATTR_MIN_P] = float(t.user_attrs[ATTR_MIN_P])
    elif "temperature" in params:
        user_attrs[ATTR_MIN_P] = min_p_from_temperature(
            float(params["temperature"]), tol_ref
        )
    elif ATTR_MIN_P in t.params:
        user_attrs[ATTR_MIN_P] = float(t.params[ATTR_MIN_P])
    # create_trial() stamps both datetimes to now() for finished trials, so
    # duration would always be 0. Copy the source trial's wall-clock span.
    frozen = create_trial(
        state=TrialState.COMPLETE,
        value=float(value),
        params=params,
        distributions=distributions,
        user_attrs=user_attrs,
        system_attrs=sys_attrs,
    )
    if t.datetime_start is not None:
        frozen.datetime_start = t.datetime_start
    if t.datetime_complete is not None:
        frozen.datetime_complete = t.datetime_complete
    return frozen


def ensure_single_objective_study(
    storage: str,
    study_name: str,
    sampler: optuna.samplers.BaseSampler,
    obj_index: int,
) -> optuna.Study:
    """Load or create a single-objective maximize study for ``obj_index``."""
    obj_name = OBJECTIVE_NAMES[obj_index]

    if not _study_exists(storage, study_name):
        study = optuna.create_study(
            study_name=study_name,
            storage=storage,
            sampler=sampler,
            direction="maximize",
            load_if_exists=False,
        )
        study.set_user_attr(ATTR_OBJECTIVE_INDEX, obj_index)
        study.set_user_attr(ATTR_BUDGET_SINCE_BEST, 0)
        print(
            f"Created single-obj study {_short_study_name(study_name)} "
            f"(maximize {obj_name})."
        )
        return study

    existing = optuna.load_study(
        study_name=study_name, storage=storage, sampler=sampler
    )
    if existing.user_attrs.get(ATTR_OBJECTIVE_INDEX) != obj_index:
        existing.set_user_attr(ATTR_OBJECTIVE_INDEX, obj_index)
    if ATTR_BUDGET_SINCE_BEST not in existing.user_attrs:
        existing.set_user_attr(ATTR_BUDGET_SINCE_BEST, 0)
    return existing


def ensure_exam_study_shards(
    storage: str,
    study_key: str,
    sampler_factory: Any,
) -> list[optuna.Study]:
    """Ensure 3 single-objective studies ``{study_key}:i`` for this exam hash."""
    names = shard_study_names(study_key)
    return [
        ensure_single_objective_study(
            storage, name, sampler_factory(), obj_index=i
        )
        for i, name in enumerate(names)
    ]


def pull_unique_trials_from_peers(
    target: optuna.Study,
    peers: list[optuna.Study],
) -> int:
    """Copy COMPLETE trials with unseen params from peer studies into target.

    Re-maps Optuna value to target's single objective from shared component
    scores. Uniqueness is by parameter fingerprint. Exam-hash isolation is
    the caller's responsibility (only pass sibling shards).

    Only COMPLETE is synced — WAITING on peers is ignored. After a pull, call
    ``drop_waiting_already_complete`` so local WAITING with the same params
    are cancelled instead of re-running generation.
    """
    obj_index = shard_objective_index(target.study_name)
    have: set[tuple[tuple[str, Any], ...]] = set()
    for t in target.trials:
        if t.state == TrialState.COMPLETE:
            p = _trial_fixed_or_params(t)
            if p:
                have.add(_param_key(p))

    added = 0
    for peer in sorted(peers, key=lambda s: s.study_name):
        if peer.study_name == target.study_name:
            continue
        for t in sorted(peer.trials, key=lambda x: x.number):
            if t.state != TrialState.COMPLETE:
                continue
            p = _trial_fixed_or_params(t)
            if not p:
                continue
            key = _param_key(p)
            if key in have:
                continue
            frozen = _complete_trial_copy(
                t,
                obj_index=obj_index,
                imported_from=peer.study_name,
            )
            if frozen is None:
                continue
            target.add_trial(frozen)
            have.add(key)
            added += 1
    return added


@dataclass(frozen=True)
class StudyBestInfo:
    """'Last seen best' from this study's single scalar score."""

    last_best_number: int  # -1 if no COMPLETE values yet
    last_best_value: float | None
    trials_since_best: int  # COMPLETE trial numbers after best (diagnostic)
    budget_since_best: int  # own optimize steps since last new best
    n_complete: int
    obj_index: int
    obj_name: str


def study_work_budget(study: optuna.Study) -> int:
    """Per-study optimize steps without a new best (study.user_attrs)."""
    raw = study.user_attrs.get(ATTR_BUDGET_SINCE_BEST, 0)
    try:
        return max(0, int(raw))
    except (TypeError, ValueError):
        return 0


def set_study_work_budget(study: optuna.Study, n: int) -> None:
    """Persist work-budget counter on the study (survives restarts)."""
    study.set_user_attr(ATTR_BUDGET_SINCE_BEST, int(max(0, n)))


def _best_value_improved(
    before: float | None, after: float | None
) -> bool:
    """True if maximize-objective best rose (or first best appeared)."""
    if after is None:
        return False
    if before is None:
        return True
    return float(after) > float(before)


def study_best_info(study: optuna.Study) -> StudyBestInfo:
    """Last best = study.best_trial (single-objective); since = later trials.

    Empty studies (no COMPLETE) use trials_since_best=0 so they stay under
    SINCE_BEST_CAP and get work — never the old 10**12 sentinel that made
    brand-new studies look "already done" and exit the optimizer immediately.

    ``budget_since_best`` is the individual work meter used for scheduling
    (imported peer trials do not inflate it).
    """
    obj_index = shard_objective_index(study.study_name)
    obj_name = OBJECTIVE_NAMES[obj_index]
    budget = study_work_budget(study)
    complete = [
        t
        for t in study.trials
        if t.state == TrialState.COMPLETE and t.values is not None
    ]
    n_complete = len(complete)
    if not complete:
        return StudyBestInfo(
            last_best_number=-1,
            last_best_value=None,
            trials_since_best=0,
            budget_since_best=budget,
            n_complete=0,
            obj_index=obj_index,
            obj_name=obj_name,
        )

    max_num = max(t.number for t in complete)
    try:
        best = study.best_trial
    except ValueError:
        # COMPLETE rows exist but no best (e.g. all non-finite) — treat as
        # "no best recorded yet", not as past the cap.
        return StudyBestInfo(
            last_best_number=-1,
            last_best_value=None,
            trials_since_best=0,
            budget_since_best=budget,
            n_complete=n_complete,
            obj_index=obj_index,
            obj_name=obj_name,
        )

    return StudyBestInfo(
        last_best_number=int(best.number),
        last_best_value=float(best.value) if best.value is not None else None,
        trials_since_best=int(max_num - best.number),
        budget_since_best=budget,
        n_complete=n_complete,
        obj_index=obj_index,
        obj_name=obj_name,
    )


def pick_study_lowest_budget(
    studies: list[optuna.Study],
    *,
    since_best_cap: int = SINCE_BEST_CAP,
) -> optuna.Study | None:
    """Pick the study with the lowest ``last_budget_for_new_best``.

    Only studies still under ``since_best_cap`` are eligible. Returns None
    when every study has already hit the cap (optimizer should exit).

    Opposite of the old "most trials since best" policy: prefer the least
    worked / most recently improved shard. Tie-break: fewest COMPLETE trials,
    then study name.
    """
    if not studies:
        return None
    eligible = [s for s in studies if study_work_budget(s) < since_best_cap]
    if not eligible:
        return None
    return min(
        eligible,
        key=lambda s: (
            study_work_budget(s),
            study_best_info(s).n_complete,
            s.study_name,
        ),
    )


def reload_exam_shards(
    storage: str,
    study_key: str,
    sampler_factory: Any,
) -> list[optuna.Study]:
    """Reload all three single-obj studies after SQLite hard-deletes."""
    out: list[optuna.Study] = []
    for name in shard_study_names(study_key):
        st = optuna.load_study(
            study_name=name, storage=storage, sampler=sampler_factory()
        )
        out.append(st)
    return out


def _print_study_status_line(study: optuna.Study) -> None:
    info = study_best_info(study)
    best_s = (
        f"{info.last_best_value:.4f}"
        if info.last_best_value is not None
        else "—"
    )
    print(
        f"  {_short_study_name(study.study_name)} "
        f"n={info.n_complete} best#{info.last_best_number}={best_s} "
        f"bud{info.budget_since_best}/{SINCE_BEST_CAP}"
    )


def _print_study_bests(studies: list[optuna.Study]) -> None:
    for study in studies:
        info = study_best_info(study)
        best_s = (
            f"{info.last_best_value:.4f}"
            if info.last_best_value is not None
            else "—"
        )
        print(
            f"best {_short_study_name(study.study_name)} "
            f"#{info.last_best_number}={best_s} "
            f"n={info.n_complete} bud{info.budget_since_best}/{SINCE_BEST_CAP}"
        )
        try:
            best = study.best_trial
        except ValueError:
            continue
        sty = best.user_attrs.get(ATTR_STYLO)
        e256 = best.user_attrs.get(ATTR_EMB_256)
        e32 = best.user_attrs.get(ATTR_EMB_32)

        def _f(v: Any) -> str:
            try:
                return f"{float(v):.3f}"
            except (TypeError, ValueError):
                return "—"

        p = best.params
        print(
            f"  sty:{_f(sty)} e256:{_f(e256)} e32:{_f(e32)} "
            f"k{p.get('top_k','—')} p{p.get('top_p','—')} T{p.get('temperature','—')}"
        )


_SIGINT_HITS = 0


def _force_exit_now(msg: str = "Stopped.") -> None:
    """Immediate process death — kill stylo process workers first (no orphans)."""
    try:
        sys.stderr.write(f"\n{msg}\n")
        sys.stderr.flush()
    except Exception:
        pass
    try:
        # SIGKILL process-pool children; shutdown(wait=False) alone leaves them.
        shutdown_metric_pool(wait=False)
    except Exception:
        pass
    # 130 = 128 + SIGINT. Never return: atexit would wait on spaCy workers.
    os._exit(130)


def _install_sigint_fast_exit() -> None:
    """Ctrl+C must not hang for minutes on stylo/embed pool joins.

    Previous behavior raised KeyboardInterrupt then tried SQLite + best
    summary while finalize() blocked on thousands of fut.result() calls —
    second Ctrl+C often never reached Python (stuck in C/BLIS or atexit).

    Now: **first** SIGINT hard-exits after cancelling the metric pool.
    """
    import signal

    def _handler(signum: int, frame: Any) -> None:
        global _SIGINT_HITS
        _SIGINT_HITS += 1
        # Always hard-exit. Soft KeyboardInterrupt paths hang under load.
        _force_exit_now(
            "Ctrl+C — force exit (skip stylo drain / atexit joins)."
            if _SIGINT_HITS >= 1
            else "Stopped."
        )

    try:
        signal.signal(signal.SIGINT, _handler)
        signal.signal(signal.SIGTERM, _handler)
    except Exception:
        pass


def _interrupt_cleanup_and_exit(
    db_path: Path,
    studies: list[optuna.Study],
    study_key: str,
    sampler_factory: Any,
) -> None:
    """Legacy path if KeyboardInterrupt arrives without the signal handler."""
    _force_exit_now("Stopped (KeyboardInterrupt).")


def main() -> None:
    load_dotenv(ROOT / ".env")
    _install_sigint_fast_exit()
    # BLAS/OpenMP pins also applied at import time (see top of file).
    # Previous Ctrl+C used os._exit without killing ProcessPool children.
    try:
        reap_stale_metric_workers()
    except Exception:
        pass

    exam_toml = ROOT / "exam.toml"
    exam_txt = ROOT / "exam.txt"
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    tokenizer_path = TOKENIZER_PATH
    for p in (exam_toml, exam_txt):
        if not p.is_file():
            raise SystemExit(f"missing required file: {p}")

    cfg, toml_hash = load_exam_toml(exam_toml)
    ensure_tokenizer(tokenizer_path, cfg.tokenizer_url)
    tok = load_tokenizer(tokenizer_path)
    exam = build_exam(cfg, exam_txt, tok, toml_hash=toml_hash)
    print(
        f"exam {cfg.model} pref{cfg.prefill_tokens} "
        f"ans{cfg.answer_tokens} n={cfg.n_samples} "
        f"prompt{len(exam.prompt_ids)} "
        f"key{exam.study_key[:8]} shards{N_STUDY_SHARDS}"
    )
    print(
        f"  mode=fullgen stylo={STYLO_CHUNK}/{STYLO_STRIDE} "
        f"db=presets_fin.db"
    )

    nai_token = os.getenv("NOVELAI_TOKEN")
    if not nai_token:
        raise SystemExit("NOVELAI_TOKEN missing in .env")

    embed_client = make_embed_client()

    # One SQLite file; study_key fingerprints exam.toml + exam.txt + layout.
    # Three single-obj studies: "{key}:0" / ":1" / ":2".
    db_path = (ROOT / "presets_fin.db").resolve()
    storage = f"sqlite:///{db_path.as_posix()}"
    study_key = exam.study_key

    def _new_sampler() -> BaseSampler:
        # Lazy load so import of this module stays light; hub may download once.
        return optunahub.load_module("samplers/auto_sampler").AutoSampler()

    studies = ensure_exam_study_shards(storage, study_key, _new_sampler)

    # Drop leftover FAIL (e.g. catch=Exception) and RUNNING (crashed workers).
    purged = purge_fail_running_trials(db_path)
    if purged:
        print(
            f"Purged {len(purged)} FAIL/RUNNING trial(s) at startup: "
            + ", ".join(f"#{n}({st})" for n, st in purged[:12])
            + ("…" if len(purged) > 12 else "")
        )
        studies = reload_exam_shards(storage, study_key, _new_sampler)

    # Metric pool: process+fork COW on Linux (shared spaCy pages) + batched pipe.
    print("stylo pool warm…", flush=True)
    info = metric_pool_info(N_SAMPLES)
    n_metric_workers = warmup_metric_workers(N_SAMPLES)
    print(
        f"stylo {STYLO_CHUNK}/{STYLO_STRIDE} "
        f"kind={info.get('pool_kind', '?')} "
        f"workers={n_metric_workers} ready "
        f"stylo≈{info.get('stylo_ram_est_gib', '?')}g "
        f"script≈{info.get('script_ram_est_gib', '?')}g/"
        f"≤{info.get('script_ram_budget_gib', 8)}g "
        f"cpu={info['cpu_count']} host_ram={info['ram_gb']}g "
        f"emb_flush={EMBED_FLUSH_SIZE} emb_par={_embed_batch_parallel()} "
        f"METRIC_WORKERS={info['override']}",
        flush=True,
    )
    human = ensure_human_cache(exam, embed_client, tok, db_path=db_path)

    for st in studies:
        _print_study_status_line(st)

    print(
        f"opt until bud≥{SINCE_BEST_CAP}/study  "
        f"objs={','.join(OBJECTIVE_SHORT)}  "
        f"db={db_path.name} key{study_key[:8]}  Ctrl+C stop"
    )

    try:
        while True:
            target = pick_study_lowest_budget(studies)
            if target is None:
                print(
                    f"\nAll {len(studies)} studies hit "
                    f"work-budget≥{SINCE_BEST_CAP}; ending optimizer."
                )
                for st in studies:
                    _print_study_status_line(st)
                break

            target_name = target.study_name
            obj_index = shard_objective_index(target_name)
            peers = [s for s in studies if s.study_name != target_name]
            best_before_pull = study_best_info(target).last_best_value
            n_pulled = pull_unique_trials_from_peers(target, peers)
            if n_pulled:
                print(
                    f"  pulled {n_pulled} unique trial(s) into "
                    f"{_short_study_name(target_name)} from peers"
                )
                target = optuna.load_study(
                    study_name=target_name,
                    storage=storage,
                    sampler=_new_sampler(),
                )
                best_after_pull = study_best_info(target).last_best_value
                if _best_value_improved(best_before_pull, best_after_pull):
                    set_study_work_budget(target, 0)
                    print(
                        f"  budget reset on {_short_study_name(target_name)} "
                        f"(peer import improved best → {best_after_pull:.4f})"
                    )

            # COMPLETE (local or just imported) cancels same-param WAITING so
            # we do not re-generate a preset another shard already scored.
            dropped = drop_waiting_already_complete(target, db_path)
            if dropped:
                print(
                    f"  dropped WAITING already COMPLETE on "
                    f"{_short_study_name(target_name)}: "
                    + ", ".join(f"#{n}" for n in dropped)
                )
                target = optuna.load_study(
                    study_name=target_name,
                    storage=storage,
                    sampler=_new_sampler(),
                )

            before = study_best_info(target)
            best_s = (
                f"{before.last_best_value:.4f}"
                if before.last_best_value is not None
                else "—"
            )
            print(
                f"→ {_short_study_name(target_name)} "
                f"best#{before.last_best_number}={best_s} "
                f"bud{before.budget_since_best}/{SINCE_BEST_CAP} "
                f"n{before.n_complete}"
            )

            objective = make_objective(
                exam, nai_token, embed_client, human, tok, obj_index
            )

            def _after_trial(
                _study: optuna.Study, trial: optuna.trial.FrozenTrial
            ) -> None:
                removed = purge_fail_running_trials(db_path)
                if removed:
                    print(
                        "  purged FAIL/RUNNING: "
                        + ", ".join(f"#{n}({st})" for n, st in removed)
                    )
                cur = optuna.load_study(
                    study_name=target_name,
                    storage=storage,
                    sampler=_new_sampler(),
                )
                is_best = False
                if trial.state == TrialState.COMPLETE and trial.value is not None:
                    try:
                        is_best = cur.best_trial.number == trial.number
                    except ValueError:
                        is_best = False
                if is_best:
                    set_study_work_budget(cur, 0)
                    print(
                        f"  ★ {_short_study_name(target_name)} "
                        f"#{trial.number}={float(trial.value):.4f} bud0"
                    )
                else:
                    new_b = study_work_budget(cur) + 1
                    set_study_work_budget(cur, new_b)
                    st = (
                        f"#{trial.number}"
                        if trial.state == TrialState.COMPLETE
                        else trial.state.name[:4]
                    )
                    print(
                        f"  bud {_short_study_name(target_name)} "
                        f"→{new_b}/{SINCE_BEST_CAP} {st}"
                    )

            target.optimize(
                objective,
                n_trials=1,
                catch=(Exception,),
                callbacks=[_after_trial],
            )

            studies = reload_exam_shards(storage, study_key, _new_sampler)
    except KeyboardInterrupt:
        print("\nStopped (Ctrl+C). Cleaning up…", flush=True)
        _interrupt_cleanup_and_exit(db_path, studies, study_key, _new_sampler)

    final_purged = purge_fail_running_trials(db_path)
    if final_purged:
        print(
            f"Purged {len(final_purged)} FAIL/RUNNING at shutdown: "
            + ", ".join(f"#{n}({st})" for n, st in final_purged)
        )
        studies = reload_exam_shards(storage, study_key, _new_sampler)
    _print_study_bests(studies)


if __name__ == "__main__":
    main()

