"""
NovelAI sampling-preset optimizer (lean single file).

Loads exam.txt + exam.toml, downloads tokenizer.json from tokenizer_url when
missing or when the URL changes (tracked via tokenizer.json.meta), generates
11 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.

Per exam fingerprint (sha256 of exam.toml + exam.txt) 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:   per-sample text-metric vectors (see text_metrics.py) → geometric
           median of the sample vectors → scale-normalized L2 match vs human metrics
           Metrics start as soon as each sample's full text is ready (overlap gen).
  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: every time a
new full 32-tok or 256-tok window completes, one OpenRouter embed call is fired
in a background thread. By generation end, scores are mostly ready (only a
possible final covering window + in-flight 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

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 text_metrics import (
    FEATURE_NAMES,
    extract_vector,
    extract_vectors,
    metric_pool_info,
    submit_extract_vector,
    warmup_metric_workers,
)

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
N_SAMPLES = 11
BASE_DELAY = 1.0
MAX_DELAY = 128.0
# Bump when text_metrics / chunk embedding semantics change (human cache).
METRICS_VERSION = "v4-stride16-py"
# 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.
#
# Always free: top_k, top_p, temperature.
# min_p is never free: min_p = exp(-logit_tolerance / temperature).
# logit_tolerance (nats) from exam.toml when set; else Optuna suggests.
# 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
    # Nats in min_p = exp(-logit_tolerance / T). None → Optuna tunes it.
    logit_tolerance: float | None


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


def make_study_key(toml_hash: str, exam_hash: str) -> str:
    """Study identity: exam.toml + exam.txt + N_SAMPLES + chunk layout.

    N_SAMPLES and embed chunk (size, stride) are fingerprinted so runs with
    different sample counts or window overlap never share Optuna studies.
    Human baselines are one reference text; the cache file is keyed by this
    study_key so a layout change triggers a recompute.
    """
    return hashlib.sha256(
        f"{toml_hash}:{exam_hash}:n_samples={N_SAMPLES}"
        f":chunk32={CHUNK_32}/{CHUNK_32_STRIDE}"
        f":chunk256={CHUNK_256}/{CHUNK_256_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 logit_tolerance (nats). Absent → Optuna searches.
    # Legacy keys: min_p_constant (nats, alias), min_p_bits → nats = bits*ln2.
    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}"
            )
    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,
        logit_tolerance=logit_tolerance,
    )
    if cfg.prefill_tokens < 1 or cfg.answer_tokens < 1:
        raise SystemExit("prefill_tokens and answer_tokens must be >= 1")
    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")
    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:
    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}"
        )
    prefill = ids[: cfg.prefill_tokens]
    answer_ids = ids[cfg.prefill_tokens : need]
    answer_text = tok.decode(answer_ids)
    prompt = list(cfg.system_prefix) + 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)
    # prompt + generation must fit NovelAI context (same budget every resume).
    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,
    )


_METRIC_WARNED = False


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 _metric_vector(text: str) -> np.ndarray:
    """Fixed-order fiction metrics; non-finite → 0. Failures soft-zero once warned."""
    global _METRIC_WARNED
    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:
        if not _METRIC_WARNED:
            print(f"warning: metric extract failed ({type(e).__name__}: {e})")
            _METRIC_WARNED = True
        return np.zeros(len(FEATURE_NAMES), dtype=float)
    return _sanitize_metric_row(v)


def _metric_matrix(texts: list[str]) -> np.ndarray:
    """Batch metrics for several texts (obj1 multi-sample path)."""
    global _METRIC_WARNED
    cleaned = [(t or "").strip() for t in texts]
    if any(not t for t in cleaned):
        raise ValueError("empty text in metric batch")
    try:
        mat = np.asarray(extract_vectors(cleaned), dtype=float)
    except Exception as e:
        if not _METRIC_WARNED:
            print(f"warning: metric extract failed ({type(e).__name__}: {e})")
            _METRIC_WARNED = True
        # Soft-fallback: per-row zeros keep shape so callers can still fail later
        # with a clearer distance/score path if needed.
        return np.zeros((len(cleaned), len(FEATURE_NAMES)), dtype=float)
    if mat.ndim != 2 or mat.shape[0] != len(cleaned):
        raise RuntimeError(
            f"metric matrix shape {mat.shape}, expected ({len(cleaned)}, n_feat)"
        )
    return np.stack([_sanitize_metric_row(mat[i]) for i in range(mat.shape[0])], axis=0)


def _submit_sample_metrics(text: str) -> Future[np.ndarray]:
    """Fire obj1 extract as soon as one full sample text is ready."""
    return submit_extract_vector(text, n_texts_hint=N_SAMPLES)


def _resolve_sample_metrics(sample: GeneratedSample) -> np.ndarray:
    """Wait for early-started metrics future, or compute now as fallback."""
    global _METRIC_WARNED
    fut = sample.metrics_future
    if fut is not None:
        try:
            return _sanitize_metric_row(np.asarray(fut.result(), dtype=float))
        except Exception as e:
            if not _METRIC_WARNED:
                print(
                    f"warning: early metric future failed "
                    f"({type(e).__name__}: {e}); recomputing"
                )
                _METRIC_WARNED = True
    return _metric_vector(sample.text)


def _metrics_matrix_from_samples(samples: list[GeneratedSample]) -> np.ndarray:
    """
    Collect obj1 vectors for finished samples.

    Prefer futures started when each sample completed (overlap with remaining
    generation). Join all futures (work already parallel); only samples without
    a future are extracted now via the batch path.
    """
    n = len(samples)
    out = [None] * n  # type: list[np.ndarray | None]
    need_batch: list[int] = []

    # Kick off any missing extracts first so late samples still parallelize.
    pending: list[tuple[int, Future[np.ndarray]]] = []
    for i, s in enumerate(samples):
        if s.metrics_future is not None:
            pending.append((i, s.metrics_future))
        else:
            need_batch.append(i)

    if need_batch:
        texts = [(samples[i].text or "").strip() for i in need_batch]
        mat = _metric_matrix(texts)
        for row, i in enumerate(need_batch):
            out[i] = mat[row]

    for i, fut in pending:
        try:
            out[i] = _sanitize_metric_row(np.asarray(fut.result(), dtype=float))
        except Exception:
            out[i] = _metric_vector(samples[i].text)

    if any(v is None for v in out):
        raise RuntimeError("internal: incomplete metric matrix from samples")
    return np.stack(out, axis=0)  # type: ignore[arg-type]


def _metric_scale(human_metrics: np.ndarray) -> np.ndarray:
    """Per-feature scale so count-like dims do not dominate L2 / geom-median."""
    return np.maximum(np.abs(np.asarray(human_metrics, dtype=float)), 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")
    return OpenAI(base_url="https://openrouter.ai/api/v1", api_key=key)


def embed_texts(client: OpenAI, texts: list[str]) -> np.ndarray:
    """Embed texts in one API call; force max dimensions (4096). Empty → (0, d).

    Used for human baseline batching and any multi-text call. Live generation
    prefers one request per chunk (see LiveChunkEmbedder) so embed latency
    overlaps NovelAI streaming.
    """
    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]
    resp = client.embeddings.create(
        model=EMBED_MODEL,
        input=clean_texts,
        dimensions=EMBED_DIM,
        encoding_format="float",
    )
    raw = getattr(resp, "data", None)
    if raw is None:
        raise RuntimeError("embedding API returned data=None")
    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)}, 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: {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}), got {arr.shape}"
        )
    if not np.all(np.isfinite(arr)):
        raise RuntimeError("embedding contains non-finite values")
    return arr


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."""
    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 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(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:
    """Submit one embed request per completed sliding window as ids stream in.

    Full windows (CHUNK_32 / CHUNK_256 with stride 16) fire as soon as enough
    tokens exist. A final covering window (if needed) is submitted in
    finalize(). Soft-cut resume reuses the same instance so already-submitted
    starts are not re-embedded.
    """

    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 submitted (full windows only).
        self._next_start_32 = 0
        self._next_start_256 = 0
        # Starts already submitted (so finalize can skip duplicates).
        self._starts_32: set[int] = set()
        self._starts_256: set[int] = set()
        self._futs_32: list[Future[np.ndarray]] = []
        self._futs_256: list[Future[np.ndarray]] = []
        self._lock = threading.Lock()
        self._finalized = False
        # 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 and fire embeds for any newly completed full windows."""
        with self._lock:
            if self._finalized:
                return
            self.token_ids = list(token_ids)
            self._submit_full_unlocked()

    def _try_submit(self, text: str) -> Future[np.ndarray] | None:
        """Submit one embed job; None if the pool is already shut down."""
        try:
            return self._pool.submit(embed_one, self._client, text)
        except RuntimeError:
            # ThreadPoolExecutor shut down after trial completed enough samples.
            return None

    def _submit_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 = self._tok.decode(ids[start : start + CHUNK_32])
            fut = self._try_submit(text)
            if fut is None:
                return
            self._futs_32.append(fut)
            self._starts_32.add(start)
            self._next_start_32 += CHUNK_32_STRIDE
        while self._next_start_256 + CHUNK_256 <= n:
            start = self._next_start_256
            text = self._tok.decode(ids[start : start + CHUNK_256])
            fut = self._try_submit(text)
            if fut is None:
                return
            self._futs_256.append(fut)
            self._starts_256.add(start)
            self._next_start_256 += CHUNK_256_STRIDE

    def progress(self) -> tuple[int, int]:
        """(completed futures, submitted futures) across both chunk sizes."""
        with self._lock:
            futs = list(self._futs_32) + list(self._futs_256)
        done = sum(1 for f in futs if f.done())
        return done, len(futs)

    def _submit_span_unlocked(
        self,
        ids: list[int],
        start: int,
        end: int,
        *,
        starts: set[int],
        futs: list[Future[np.ndarray]],
    ) -> bool:
        """Submit one span if start not yet done. False if pool is shut down."""
        if start in starts:
            return True
        text = self._tok.decode(ids[start:end])
        fut = self._try_submit(text)
        if fut is None:
            return False
        futs.append(fut)
        starts.add(start)
        return True

    def finalize(self) -> tuple[np.ndarray, np.ndarray]:
        """Cover any remaining spans, wait all futures → (emb_32, emb_256)."""
        with self._lock:
            if self._finalized:
                raise RuntimeError("LiveChunkEmbedder.finalize() called twice")
            self._finalized = True
            self._submit_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._submit_span_unlocked(
                    ids,
                    start,
                    end,
                    starts=self._starts_32,
                    futs=self._futs_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._submit_span_unlocked(
                    ids,
                    start,
                    end,
                    starts=self._starts_256,
                    futs=self._futs_256,
                ):
                    raise RuntimeError(
                        "embed pool shut down before final window embed (256)"
                    )
            futs_32 = list(self._futs_32)
            futs_256 = list(self._futs_256)

        rows_32 = [f.result() for f in futs_32]
        rows_256 = [f.result() for f in futs_256]
        emb_32 = _stack_embs(rows_32)
        emb_256 = _stack_embs(rows_256)

        # Sanity: count must match offline token_id_chunks (order = span order).
        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})"
            )
        return emb_32, emb_256


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


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

@dataclass(frozen=True)
class HumanCache:
    metrics: np.ndarray  # (n_features,) fiction text-metric vector
    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",
)


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 embeddings collected during streaming.

    metrics_future: obj1 vector job started as soon as this sample's full
    text is ready (overlaps later samples still streaming). None if not queued.
    """

    text: str
    emb_32: np.ndarray  # (n_chunks_32, EMBED_DIM)
    emb_256: np.ndarray  # (n_chunks_256, EMBED_DIM)
    metrics_future: Future[np.ndarray] | None = None


def compute_human_cache(
    answer_text: str, embed_client: OpenAI, tokenizer: Tokenizer
) -> HumanCache:
    metrics = _metric_vector(answer_text)
    emb_32 = embed_token_chunks(embed_client, answer_text, tokenizer, CHUNK_32)
    emb_256 = embed_token_chunks(embed_client, answer_text, tokenizer, CHUNK_256)
    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 text metrics + embeddings already 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")

    # stylo: per-sample metrics (usually already running since each sample finished)
    # → geometric median → match human vector.
    ai_metrics = _metrics_matrix_from_samples(samples)
    scale = _metric_scale(human.metrics)
    metrics_dist_per = [
        _scaled_metric_distance(ai_metrics[i], human.metrics, scale)
        for i in range(N_SAMPLES)
    ]
    s_stylo = [_per_sample_match_score(d) for d in metrics_dist_per]
    # Geom-median in human-scaled feature space so mixed units are comparable.
    med_metrics = geometric_median(ai_metrics / scale)
    metrics_dist = _finite(
        float(np.linalg.norm(med_metrics - human.metrics / scale))
    )
    stylo = _per_sample_match_score(metrics_dist)

    # 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 workers per trial. I/O-bound (OpenRouter RTT); not local CPU.
# Effectively unlimited so chunk jobs across live streams never queue on the pool.
EMBED_POOL_WORKERS = 9999


class NaiPool:
    """
    Fill N_SAMPLES completions for one preset.

    Slots in flight: finished + streaming (+ idle parked). At most
    MAX_CONCURRENT_STREAMS NovelAI stream POSTs at once. Prefer resuming
    idle partials over new blanks. 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). At delay 128 after a
    fail: wipe that partial and retry blank. Any successful stream token
    only resets the *delay value* back to 1s — it never clears the active
    countdown (so backoff is not cancelled mid-flight).

    While tokens stream, LiveChunkEmbedder fires one OpenRouter embed request
    per completed 32/256-token window so scoring is nearly ready at gen end.

    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
        self.prompt_ids = list(prompt_ids)
        self.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] = []
        self._streaming = 0
        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] = {}
        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()

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

    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 self.tokenizer.decode(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(self) -> tuple[int, int]:
        """Sum (done, submitted) embed futures across live + idle embedders."""
        done = 0
        total = 0
        for emb in self._live_embedders.values():
            d, t = emb.progress()
            done += d
            total += t
        for emb in self._idle:
            d, t = emb.progress()
            done += d
            total += t
        return done, total

    def _progress_msg(self, *, note: str = "") -> str:
        left = max(0.0, self._deadline - time.perf_counter())
        tgt = self.answer_tokens
        # Single combined counter: finished samples (full) + live + idle partials
        # over N_SAMPLES * answer_tokens, e.g. tok=121/59750
        cur = (
            len(self._finished) * tgt
            + sum(self._live_tokens.values())
            + sum(len(p.token_ids) for p in self._idle)
        )
        total = N_SAMPLES * tgt
        emb_done, emb_total = self._embed_progress()
        elapsed = time.perf_counter() - self._t0
        http_s = (
            f"{self._http_attempts}/{MAX_HTTP_ATTEMPTS}"
            if MAX_HTTP_ATTEMPTS is not None
            else str(self._http_attempts)
        )
        msg = (
            f"{self._progress_prefix}"
            f"done={len(self._finished)}/{N_SAMPLES} "
            f"idle={len(self._idle)} "
            f"tok={cur}/{total} "
            f"emb={emb_done}/{emb_total} "
            f"delay={self._delay:.0f}s cd={left:.0f}s "
            f"http={http_s} "
            f"t={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)."""
        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
            msg = self._progress_msg(note=note)
        # progress may have been detached after we built msg
        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:
        """Silence late worker UI and refuse new work after run() returns."""
        with self._cv:
            self._stop = True
            self._progress = None
            self._cv.notify_all()

    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()
                    # Remaining sample capacity and concurrent stream cap.
                    sample_slots = (
                        N_SAMPLES - len(self._finished) - self._streaming
                    )
                    stream_slots = MAX_CONCURRENT_STREAMS - self._streaming
                    slots_open = min(sample_slots, stream_slots)
                    if slots_open <= 0:
                        self._cv.wait(timeout=0.1)
                    elif not self._idle and (
                        len(self._finished) + self._streaming + len(self._idle)
                    ) >= N_SAMPLES:
                        # Capacity full (finished+streaming+idle); wait for a finish.
                        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
                        emb = (
                            self._idle.pop(0)
                            if self._idle
                            else 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 — release the slot.
                        with self._cv:
                            self._streaming = max(0, self._streaming - 1)
                            self._live_tokens.pop(wid, None)
                            self._live_embedders.pop(wid, None)
                            self._error = RuntimeError(
                                f"failed to start stream worker: {e}"
                            )
                            self._cv.notify_all()
        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 _arm_fail(
        self,
        embedder: LiveChunkEmbedder | None,
        partial_ids: list[int],
        *,
        park: bool,
    ) -> None:
        """Wait current delay, then double (cap 128). Wipe partial at 128."""
        # Caller must hold _cv.
        if self._stop:
            return
        wait = float(self._delay)
        # Never shorten an existing longer backoff (concurrent fail races).
        self._deadline = max(self._deadline, time.perf_counter() + wait)
        if wait >= MAX_DELAY:
            if park:
                # Blank retry after 128s wait — drop old embeds with the wipe.
                self._idle.append(self._new_embedder())
            self._delay = BASE_DELAY
        else:
            next_delay = min(wait * 2.0, MAX_DELAY)
            if park:
                emb = embedder if embedder is not None else self._new_embedder()
                emb.set_tokens(partial_ids)
                self._idle.append(emb)
            self._delay = next_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 _worker(self, embedder: LiveChunkEmbedder, wid: int) -> None:
        try:
            if self._stop:
                with self._cv:
                    self._streaming = max(0, self._streaming - 1)
                    self._live_tokens.pop(wid, None)
                    self._live_embedders.pop(wid, None)
                    self._cv.notify_all()
                return
            clean_partial, _dropped = _sanitize_token_ids(embedder.token_ids)
            embedder.set_tokens(clean_partial)
            self._set_live(wid, len(clean_partial))
            outcome, text, tokens = self._one_stream(embedder, wid)
            sample: GeneratedSample | None = None
            if outcome == "done" and (text or "").strip():
                # Full text ready: start obj1 metrics immediately so they overlap
                # remaining streams / embed finalization of other samples.
                metrics_fut: Future[np.ndarray] | None = None
                try:
                    metrics_fut = _submit_sample_metrics(text)
                    self._ui(force=True, note="obj1 metrics queued")
                except Exception as met_err:
                    # Scoring can still compute metrics later; do not drop sample.
                    self._ui(
                        force=True,
                        note=f"obj1 queue fail: {type(met_err).__name__}",
                    )
                    metrics_fut = None
                # Wait for in-flight embeds + short final windows before parking
                # the finished sample. Most full-window embeds already completed
                # during the stream.
                try:
                    emb_32, emb_256 = embedder.finalize()
                    sample = GeneratedSample(
                        text=text,
                        emb_32=emb_32,
                        emb_256=emb_256,
                        metrics_future=metrics_fut,
                    )
                except Exception as emb_err:
                    # Embed blips (null data, rate-limit payload quirks) must not
                    # kill the whole trial — drop this sample and soft-retry.
                    if metrics_fut is not None:
                        metrics_fut.cancel()
                    self._ui(
                        force=True,
                        note=f"embed fail: {type(emb_err).__name__}: {emb_err}",
                    )
                    sample = None
            with self._cv:
                self._streaming = max(0, self._streaming - 1)
                self._live_tokens.pop(wid, None)
                self._live_embedders.pop(wid, None)
                if self._stop:
                    # run() already left — do not mutate pool results / UI
                    self._cv.notify_all()
                    return
                if outcome == "done":
                    if sample is None:
                        self._arm_fail(None, [], park=True)
                    elif len(self._finished) < N_SAMPLES:
                        self._finished.append(sample)
                        self._delay = BASE_DELAY
                    # else: already have enough samples (race); drop extra
                elif outcome == "partial":
                    self._arm_fail(embedder, tokens, park=True)
                else:
                    # empty drop — no idle slot; just backoff
                    self._arm_fail(None, [], park=False)
                self._cv.notify_all()
            self._ui(force=True)
        except BaseException as e:
            with self._cv:
                self._streaming = max(0, self._streaming - 1)
                self._live_tokens.pop(wid, None)
                self._live_embedders.pop(wid, None)
                if not self._stop:
                    # First hard error wins; later workers stay quiet.
                    if self._error is None:
                        self._error = e
                self._cv.notify_all()

    def _one_stream(
        self, embedder: LiveChunkEmbedder, wid: int
    ) -> tuple[str, str, list[int]]:
        """Single POST stream → (done|partial|empty, text, token_ids).

        On every token growth, LiveChunkEmbedder submits embeds for any newly
        completed full 32/256-token windows (one HTTP request per window).
        """
        out_ids, _dropped = _sanitize_token_ids(embedder.token_ids)
        embedder.set_tokens(out_ids)
        remaining = self.answer_tokens - len(out_ids)
        if remaining <= 0:
            return "done", self._decode(out_ids), out_ids

        body: dict[str, Any] = {
            "model": self.model,
            "prompt": self.prompt_ids + out_ids,
            "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 partial token ids: wipe and retry blank.
                    if out_ids and (
                        "token beyond max" in err_body
                        or "beyond max token" in err_body
                    ):
                        return "empty", "", []
                    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, self.answer_tokens - 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) >= self.answer_tokens:
                            self._set_live(wid, len(out_ids))
                            return _finish("done")
                    elif text_delta:
                        try:
                            approx = self._ids_from_text(text_delta)
                        except Exception:
                            approx = []
                        room = self.answer_tokens - 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) >= self.answer_tokens:
                                self._set_live(wid, len(out_ids))
                                return _finish("done")
                        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) >= self.answer_tokens:
            return _finish("done")
        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 completions; embeddings + obj1 metrics overlap streams.

    Per completed sample: full text → queue obj1 extract immediately, then
    finalize remaining embeds. Metrics run while other samples still generate.
    """
    with ThreadPoolExecutor(
        max_workers=EMBED_POOL_WORKERS, thread_name_prefix="embed"
    ) as embed_pool:
        pool = 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 pool.run()


# ---------------------------------------------------------------------------
# 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",
)

HUMAN_CACHE_DIR = ROOT / "human_cache"


def human_cache_path(exam: ExamBundle) -> Path:
    """One compressed npz per study_key (toml + exam + N_SAMPLES fingerprint)."""
    return HUMAN_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",
                "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["embed_dim"]) != EMBED_DIM
                or int(z["n_samples"]) != 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 == 1
        and metrics_h.shape[0] == n_feat
        and np.all(np.isfinite(metrics_h))
        and _valid_chunk_emb(emb_32)
        and _valid_chunk_emb(emb_256)
    ):
        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),
        embed_dim=np.asarray(EMBED_DIM, dtype=np.int32),
        n_samples=np.asarray(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 human_cache/<study_key>.npz — never in presets.db / Optuna
    user_attrs (embeddings are huge and bloat SQLite).
    """
    n_feat = len(FEATURE_NAMES)
    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"Loaded human baselines from {path.relative_to(ROOT) if path.is_relative_to(ROOT) else path} "
            f"(study_key={exam.study_key[:12]}… "
            f"emb32={cached.emb_32.shape[0]}×{EMBED_DIM} "
            f"emb256={cached.emb_256.shape[0]}×{EMBED_DIM})"
        )
        return cached
    if path.is_file():
        print(f"Human cache stale/invalid ({path.name}) — recomputing…")
    else:
        print(f"No human cache at {path.name} — computing…")

    n32 = len(token_chunks(exam.answer_text, tokenizer, CHUNK_32))
    n256 = len(token_chunks(exam.answer_text, tokenizer, CHUNK_256))
    print(
        f"Computing human baselines "
        f"(metrics[{n_feat}], {n32}×{CHUNK_32}-tok emb, "
        f"{n256}×{CHUNK_256}-tok emb)…"
    )
    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.relative_to(ROOT) if path.is_relative_to(ROOT) else path} "
        f"({size_mb:.1f} MiB)  metrics_dim={human.metrics.shape[0]}  "
        f"emb32={human.emb_32.shape[0]}×{EMBED_DIM}  "
        f"emb256={human.emb_256.shape[0]}×{EMBED_DIM}"
    )
    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``."""
    obj_name = OBJECTIVE_NAMES[obj_index]

    def objective(trial: optuna.Trial) -> float:
        tk = SEARCH_SPACE["top_k"]
        tp = SEARCH_SPACE["top_p"]
        temp = SEARCH_SPACE["temperature"]
        temperature = trial.suggest_float(
            "temperature", temp.low, temp.high, log=temp.log
        )
        # min_p = exp(-logit_tolerance / T). Fixed from exam.toml, else Optuna.
        fixed_tol = exam.config.logit_tolerance
        if fixed_tol is not None:
            logit_tolerance = float(fixed_tol)
        else:
            w = SEARCH_SPACE[PARAM_LOGIT_TOLERANCE]
            logit_tolerance = trial.suggest_float(
                PARAM_LOGIT_TOLERANCE, w.low, w.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)
        preset = Preset(
            top_k=trial.suggest_int("top_k", tk.low, tk.high, log=tk.log),
            min_p=min_p,
            top_p=trial.suggest_float("top_p", tp.low, tp.high),
            temperature=temperature,
        )
        prefix = (
            f"[trial {trial.number} {obj_name}] "
            f"k={preset.top_k} min_p={preset.min_p:.3f} "
            f"logit_tol={logit_tolerance:.3f} "
            f"top_p={preset.top_p:.3f} T={preset.temperature:.2f} | "
        )
        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 streamed live; stylo metrics usually already running since
            # each sample finished — scoring mostly joins futures + aggregates.
            n_ready = sum(
                1
                for s in samples
                if s.metrics_future is not None and s.metrics_future.done()
            )
            prog.update(
                f"{prefix}scoring… stylo={n_ready}/{len(samples)} ready"
            )
            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])
            v0, v1, v2 = result.values
            # Component display order is stylo / 256 emb / 32 emb (not 32 then 256).
            # score_std tuple is (stylo, 32, 256) — reorder + label to match.
            prog.end(
                f"{prefix}"
                f"{obj_name}={value:.4f} "
                f"(all {v0:.4f}/{v1:.4f}/{v2:.4f}) "
                f"[{ATTR_STYLO}={result.stylo:.4f} "
                f"{ATTR_EMB_256}={result.emb_256:.4f} "
                f"{ATTR_EMB_32}={result.emb_32:.4f}] | "
                f"std {ATTR_STYLO}={d_s:.4f} "
                f"{ATTR_EMB_256}={d_256:.4f} "
                f"{ATTR_EMB_32}={d_32:.4f}"
            )
            del samples
            gc.collect()
            return value
        except optuna.TrialPruned:
            prog.end(f"{prefix}pruned")
            raise
        except KeyboardInterrupt:
            prog.end(f"{prefix}interrupted")
            raise
        except Exception as e:
            prog.end(f"{prefix}error: {type(e).__name__}: {e}")
            tb = "".join(
                traceback.format_exception(type(e), e, e.__traceback__, limit=6)
            )
            sys.stderr.write(tb)
            sys.stderr.flush()
            gc.collect()
            raise optuna.TrialPruned() from e

    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: first 12 of exam key + shard suffix if present."""
    if SHARD_SEP in name:
        base, _, suf = name.rpartition(SHARD_SEP)
        idx = suf if suf.isdigit() else suf
        label = (
            OBJECTIVE_NAMES[int(idx)]
            if idx.isdigit() and int(idx) < len(OBJECTIVE_NAMES)
            else idx
        )
        return f"{base[:12]}…:{idx}({label})"
    return f"{name[:12]}…"


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
    if not all(np.isfinite(v) and v > 0.0 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])
    return create_trial(
        state=TrialState.COMPLETE,
        value=float(value),
        params=params,
        distributions=distributions,
        user_attrs=user_attrs,
        system_attrs=sys_attrs,
    )


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"trials={info.n_complete} "
        f"best=#{info.last_best_number} value={best_s} "
        f"budget={info.budget_since_best}/{SINCE_BEST_CAP} "
        f"(trial_since={info.trials_since_best})"
    )


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} value={best_s} "
            f"trials={info.n_complete} "
            f"budget={info.budget_since_best}/{SINCE_BEST_CAP}"
        )
        try:
            best = study.best_trial
        except ValueError:
            continue
        print(
            f"  {ATTR_STYLO}={best.user_attrs.get(ATTR_STYLO)} "
            f"{ATTR_EMB_256}={best.user_attrs.get(ATTR_EMB_256)} "
            f"{ATTR_EMB_32}={best.user_attrs.get(ATTR_EMB_32)} "
            f"{ATTR_MIN_P}={best.user_attrs.get(ATTR_MIN_P)} "
            f"{ATTR_LOGIT_TOLERANCE}="
            f"{best.user_attrs.get(ATTR_LOGIT_TOLERANCE)} "
            f"params={best.params}"
        )


def main() -> None:
    load_dotenv(ROOT / ".env")
    # Avoid BLAS/OpenMP oversubscription (many threads × spaCy × embed pool → RAM spikes).
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")
    os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")

    exam_toml = ROOT / "exam.toml"
    exam_txt = ROOT / "exam.txt"
    tokenizer_path = ROOT / "tokenizer.json"
    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)
    if cfg.logit_tolerance is not None:
        min_p_mode = (
            f"min_p=exp(-logit_tolerance/T) "
            f"logit_tolerance={cfg.logit_tolerance:.6g} (fixed in exam.toml)"
        )
    else:
        w = SEARCH_SPACE[PARAM_LOGIT_TOLERANCE]
        min_p_mode = (
            f"min_p=exp(-logit_tolerance/T) "
            f"logit_tolerance=Optuna[{w.low:g},{w.high:g}] "
            f"(exam.toml omits {PARAM_LOGIT_TOLERANCE})"
        )
    print(
        f"exam ready: model={cfg.model} prefill={cfg.prefill_tokens} "
        f"answer={cfg.answer_tokens} prompt_tokens={len(exam.prompt_ids)} "
        f"study_key={exam.study_key[:12]}… "
        f"(toml={exam.toml_hash[:8]} exam={exam.exam_hash[:8]}) "
        f"shards={N_STUDY_SHARDS} | {min_p_mode}"
    )

    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.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)

    # Preload thread-local spaCy models for stylo multi-sample extract.
    info = metric_pool_info(N_SAMPLES)
    n_metric_workers = warmup_metric_workers(N_SAMPLES)
    print(
        f"stylo metrics: workers={n_metric_workers}/{info['n_texts']} samples  "
        f"cpu={info['cpu_count']}  ram_gb={info['ram_gb']}  "
        f"override={info['override']}  embed_pool={EMBED_POOL_WORKERS}  "
        f"(1 spaCy/worker; min workers=2)"
    )
    human = ensure_human_cache(exam, embed_client, tok, db_path=db_path)

    for st in studies:
        _print_study_status_line(st)

    print(
        f"Optimizing until each study hits work-budget≥{SINCE_BEST_CAP} "
        f"without a new best (Ctrl+C to stop early) — 3 single-obj studies "
        f"({', '.join(OBJECTIVE_NAMES)}). "
        f"storage={storage} study_key={study_key[:16]}… "
        f"each step: pull unique peer trials (same exam) → run 1 trial on "
        f"study with lowest {ATTR_BUDGET_SINCE_BEST} "
        f"(FAIL/RUNNING auto-purged)"
    )

    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"→ work {_short_study_name(target_name)} "
                f"(best=#{before.last_best_number} value={best_s} "
                f"budget={before.budget_since_best}/{SINCE_BEST_CAP} "
                f"trials={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"  new best on {_short_study_name(target_name)} "
                        f"trial #{trial.number} "
                        f"value={float(trial.value):.4f} "
                        f"(budget reset)"
                    )
                else:
                    new_b = study_work_budget(cur) + 1
                    set_study_work_budget(cur, new_b)
                    print(
                        f"  budget {_short_study_name(target_name)} "
                        f"→ {new_b}/{SINCE_BEST_CAP}"
                        + (
                            f" (trial #{trial.number} not best)"
                            if trial.state == TrialState.COMPLETE
                            else f" (trial {trial.state.name})"
                        )
                    )

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

            studies = reload_exam_shards(storage, study_key, _new_sampler)
    except KeyboardInterrupt:
        print("\nStopped.")

    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()

