"""Per-query retrieval latency with a fixed thread count (HF CPU Basic profile: 2 vCPU).

Run on an otherwise idle machine after scripts/run_all.py:  uv run python scripts/latency.py
"""
import os
import sys
from pathlib import Path

import yaml

ROOT = Path(__file__).resolve().parents[1]
CFG = yaml.safe_load((ROOT / "configs" / "default.yaml").read_text())
THREADS = str(CFG["latency"]["threads"])
for var in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"):
    os.environ[var] = THREADS  # must be set before numpy/torch/faiss load

import time  # noqa: E402

import faiss  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
import torch  # noqa: E402

from gatestats import data, retrieve  # noqa: E402


def main():
    lat = CFG["latency"]
    torch.set_num_threads(lat["threads"])
    faiss.omp_set_num_threads(lat["threads"])
    name, cache = lat["dataset"], ROOT / "cache" / lat["dataset"]
    ds = data.load(name)

    fns = {}
    for system in lat["systems"]:
        if system == "bm25":
            index = retrieve.bm25_index(ds.doc_texts)
            fns[system] = lambda q, index=index: retrieve.bm25_run(index, [q], 10)
            continue
        key = CFG["hnsw"]["model"] if system.startswith("hnsw") else system
        mcfg = CFG["models"][key]
        model = retrieve.load_model(mcfg, CFG["max_seq_length"])
        if system.startswith("hnsw"):
            index = faiss.read_index(str(cache / "hnsw_s0.faiss"))
            index.hnsw.efSearch = int(system.split("_ef")[1].split("_")[0])
        else:
            index = retrieve.exact_index(np.load(cache / f"emb_{key}_docs.npy"))
        fns[system] = lambda q, m=model, p=mcfg["query_prefix"], ix=index: ix.search(
            retrieve.embed(m, [q], p, progress=False), 10)

    rows, rng = [], np.random.default_rng([CFG["seed"], 5])
    for run in range(lat["runs"]):
        for system, fn in fns.items():
            order = rng.permutation(len(ds.query_ids))
            for i in order[: lat["warmup"]]:
                fn(ds.query_texts[i])
            for i in order:
                t0 = time.perf_counter()
                fn(ds.query_texts[i])
                rows.append((system, run, ds.query_ids[i], (time.perf_counter() - t0) * 1000))
        print(f"run {run + 1}/{lat['runs']} done", file=sys.stderr, flush=True)
    out = pd.DataFrame(rows, columns=["system", "run", "qid", "ms"])
    out.to_parquet(ROOT / "results" / "latency.parquet", index=False)
    print(out.groupby("system")["ms"].describe(percentiles=[0.5, 0.95, 0.99]).round(2))


if __name__ == "__main__":
    main()