python/elidedb/context.py

"""Context retrieval: the ingest-side half.

WHY THIS EXISTS
---------------
`embeddings` holds one mean-pooled SigLIP vector per window. Mean pooling is
order-blind: reverse the frames and the vector is identical. So the index can
answer "is there a car and a person here" but never "is the person walking
*toward* the car" — appearance, not context. Reranking with a VLM fixes the
ranking but costs seconds per query, which is the wrong place to spend time in
a database.

The fix is the oldest trick a database has: precompute the expensive operator
into an index and make the query a lookup. Concretely,

    ingest (once, offline)      query (every time, hot)
    ------------------------    -----------------------
    VLM reads 3 frames of the   text -> SigLIP text vector
    window and writes a
    relational caption          ONE matmul against ctx vectors
        |
    SigLIP text tower           (no VLM in the loop, ever)
        |
    target vector  -----------> train a small temporal tower to
                                predict it from cheap frame vectors

The VLM only ever labels a subset. The tower generalises the label to every
window, including windows ingested later, so a new day of footage costs a
forward pass over frame vectors instead of a day of VLM time. This is
pseudo-labelling in the sense of "Distilling Vision-Language Models on
Millions of Videos" (arXiv 2401.06129), applied at ingest instead of at
pretraining scale.

Three tables come out of this module, all ordinary Parquet:

    frame_vectors     ts, stream, vector[d]              one row per FRAME
    context_captions  ts, t1, stream, caption, vector[d] teacher labels
    context           ts, t1, stream, vector[d]          student output
"""
from __future__ import annotations

import re
import time

import numpy as np
import pyarrow as pa
import pyarrow.compute as pc

from .embeddings import (DEFAULT_MODEL, MODELS, _embed_images,
                         _load_model, resolve_model)

# The teacher is asked for RELATIONS and CHANGE, not for a list of objects.
# An object list is exactly what SigLIP already encodes, so a caption that
# reads "a car, a road, a building" teaches the student nothing it does not
# already know. Everything interesting is in the verbs.
CAPTION_PROMPT = (
    "These frames are in time order from one short video clip. "
    "Reply with ONE sentence of at most 25 words describing what is "
    "happening: who or what is present, what each is doing, how they are "
    "positioned relative to each other, and how the scene moves or changes. "
    "Be concrete and literal. Do not say 'frame', 'image', or 'video'."
)

# The caption is the index. If it does not use the words a user would use, the
# lexical ranker can never fire and the caption-LSA space is built around the
# wrong distinctions. Measured on BridgeData2: the generic prompt above
# produced "a robot arm interacts with a wooden box", while the human label
# for the same clip was "put red object in the drawer" — different noun,
# different granularity, no overlap for retrieval to work with.
#
# So the prompt is a per-domain parameter. This one asks for the ACTION and
# the OBJECT MOVED, which is how manipulation data is described. It stays
# deliberately generic: it never names objects that appear in the labels
# (that would be leaking the eval set into the index), only the SHAPE of the
# description — what moved, and where it ended up.
# The caption IS the index: whatever verbs the prompt teaches are the only
# verbs lexical recall can ever match. The first version of this prompt said
# "which object the arm picks up or moves, and where it puts it" — and the
# resulting 2,348 captions contained picks x2371, puts x1673 and ZERO
# instances of close/open/wipe/push, so "close the drawer" was unfindable by
# construction. The prompt must be VERB-OPEN: describe the action in its own
# words, and always report state changes.
# S0 VIOLATION, FIXED. The previous MANIPULATION_PROMPT enumerated the
# task vocabulary - "picking up, putting, opening, closing, pushing,
# pouring, wiping, pressing" - and this file's own comment stated the
# consequence: "The caption IS the index: whatever verbs the prompt
# teaches are the only verbs lexical recall can ever match." A verb
# outside that list was unfindable by construction, and a non-
# manipulation corpus was unfindable entirely.
#
# It was also TUNED AGAINST THE EVAL LABELS. The removed comment
# recorded the generic prompt producing "a robot arm interacts with a
# wooden box" where "the human label for the same clip was 'put red
# object in the drawer'", and the prompt being rewritten to close that
# gap. bridge_ingest.py deliberately keeps the task strings out of the
# store; this prompt let them back in through the side door.
#
# The replacement names no verb, no object class and no domain. It asks
# for what MOVED and what CHANGED, which is answerable for a kitchen, a
# road, a warehouse or a surgical table, and lets the corpus supply its
# own words.
MANIPULATION_PROMPT = (
    "These frames are in time order from one short clip. "
    "Reply with ONE short sentence in plain English: what moved, what it "
    "did, and what was different at the end. Use whatever words fit; do "
    "not choose from a list. If anything changed state, say so. "
    "Do not say 'frame', 'image', or 'video'."
)

PROMPTS = {"scene": CAPTION_PROMPT, "manipulation": MANIPULATION_PROMPT}

# SigLIP's text tower truncates at 64 tokens, so a rambling caption is
# silently cut mid-clause and the tail is lost anyway. Trim deliberately
# instead: drop the VLM's framing preamble, keep whole sentences.
#
# The preamble pattern is deliberately narrow — a leading PREPOSITIONAL
# phrase only ("In the first frame,", "Across this sequence,"). An earlier
# looser version also matched "The video shows a street with cars," and
# amputated the subject, leaving captions that began "and various buildings".
# Requiring a leading preposition and at most two filler words makes that
# impossible.
_PREAMBLE = re.compile(
    r"^(?:in|across|throughout|over|during)\s+(?:the|this|these)\s+"
    r"(?:[\w-]+\s+){0,2}?(?:frames?|images?|pictures?|sequence|clip|video)\s*,\s*",
    re.I)

# Prompt-echo: asking for "the action verb in plain English" made the VLM
# write 'The action verb is "picking up" as the arm picks up...' — meta-
# language that pollutes the lexical index. Strip the frame, keep the deed.
_VERB_ECHO = re.compile(
    r"^the action(?:\s+verb)?\s+is\s+['\"]?[\w-]+(?:\s+[\w-]+){0,2}?"
    r"['\"]?[,.]?\s+(?:as|because|where|since|:)\s+", re.I)


def tidy_caption(text: str, max_words: int = 32) -> str:
    t = " ".join(text.strip().split())
    t = _PREAMBLE.sub("", t)
    t = _VERB_ECHO.sub("", t)
    parts = re.split(r"(?<=[.!?])\s+", t)
    # A generation cut off at max_tokens ends mid-clause. Keep only sentences
    # that actually terminate, unless that would leave nothing at all.
    whole = [p for p in parts if p.rstrip().endswith((".", "!", "?"))]
    parts = whole or parts[:1]
    out = []
    for p in parts:
        if out and len(" ".join(out + [p]).split()) > max_words:
            break
        out.append(p)
    t = " ".join(out).strip()
    return (t[:1].upper() + t[1:]) if t else " ".join(text.split())[:200]


# ---------------------------------------------------------------------------
# 1. Per-frame vectors — the sequence a temporal model needs
# ---------------------------------------------------------------------------
def embed_frames(store, frame_table="frames", model=None, width=512,
                 batch=32, incremental=True, verbose=True, stride=1,
                 streams=None, engine="siglip"):
    """One SigLIP vector per frame → `frame_vectors`.

    Deliberately NOT pooled. Pooling is the student's job, and pooling here
    would throw away the only signal that distinguishes context from
    appearance. Decode goes through the same byte-range path queries use, so
    this costs the frames it reads and nothing else.
    """
    from PIL import Image

    from .video import FrameSet
    # engine="fdnnv": the distilled streaming encoder — every frame, no
    # stride, state carried per stream. Same output space as the teacher, so
    # everything downstream (pooling, search, context) is unchanged.
    if engine == "fdnnv":
        return _embed_frames_fdnnv(store, frame_table, incremental, verbose,
                                   streams)
    model = resolve_model(model)
    frames = store.table(frame_table).scan()
    allst = sorted(set(frames.column("stream").to_pylist()))
    streams = [s for s in allst if s in streams] if streams else allst

    done = {}
    if incremental:
        try:
            prev = store.table("frame_vectors").scan()
            for s_, t_ in zip(prev.column("stream").to_pylist(),
                              prev.column("ts").to_pylist()):
                done[s_] = max(done.get(s_, -1), t_)
        except Exception:
            pass

    t_start = time.time()
    rows_ts, rows_stream, rows_vec = [], [], []
    for s in streams:
        sel = frames.filter(pc.equal(frames.column("stream"), s))
        if s in done:
            sel = sel.filter(pc.greater(sel.column("ts"), done[s]))
        if len(sel) == 0:
            continue
        # `stride` subsamples the frame index before decoding. A 5 Hz robot
        # camera does not need every frame embedded for a 4 s window to be
        # well described, and the cost here is linear in frames decoded, so
        # this is the dial between ingest time and temporal resolution.
        decoded = FrameSet(store, frame_table, sel).decode(width=width,
                                                           stride=stride)
        if verbose:
            print(f"  {s}: {len(decoded)} frames decoded", flush=True)
        for i in range(0, len(decoded), batch):
            chunk = decoded[i:i + batch]
            vecs = _embed_images([Image.fromarray(a) for _, a in chunk], model)
            for (ts, _), v in zip(chunk, vecs):
                rows_ts.append(int(ts))
                rows_stream.append(s)
                rows_vec.append(v)
    if not rows_ts:
        return {"frames": 0, "note": "nothing new (incremental)"}

    dim = len(rows_vec[0])
    tbl = pa.table({
        "ts": pa.array(rows_ts, pa.int64()),
        "stream": pa.array(rows_stream),
        "vector": pa.array([v.tolist() for v in rows_vec],
                           pa.list_(pa.float32(), dim)),
    })
    version = store.table("frame_vectors").append(
        tbl, kind="embeddings",
        meta={"model": model, "dim": dim, "decode_width": width,
              "source_table": frame_table})
    return {"frames": len(tbl), "dim": dim, "version": version,
            "seconds": round(time.time() - t_start, 1)}


# ---------------------------------------------------------------------------
# 2. Window plan — shared by teacher and student so labels line up exactly
# ---------------------------------------------------------------------------
def plan_windows(store, window_s=2.0, stride_s=0.5, table="frame_vectors",
                 min_frames=4):
    """Sliding windows over each stream's timeline.

    Stride < window on purpose: overlapping windows are how a *sliding* index
    avoids the boundary problem where an event straddles two tumbling windows
    and lands strongly in neither. Merging overlaps back into one answer is
    already handled downstream by the segment merger.
    """
    fv = store.table(table).scan()
    win = int(window_s * 1e9)
    stride = int(stride_s * 1e9)
    out = []
    for s in sorted(set(fv.column("stream").to_pylist())):
        rows = fv.filter(pc.equal(fv.column("stream"), s))
        ts = np.sort(rows.column("ts").to_numpy())
        if len(ts) == 0:
            continue
        t = int(ts[0])
        end = int(ts[-1])
        while t <= end - win // 2:
            lo, hi = np.searchsorted(ts, [t, t + win])
            if hi - lo >= min_frames:
                out.append((s, t, t + win - 1))
            t += stride
    return out


def window_sequences(store, windows, table="frame_vectors", max_len=32):
    """(stream, t0, t1) → (T, d) float32 stack of that window's frame vectors.

    Subsampled to `max_len` evenly. A 2 s window at 16 Hz is 32 frames; the
    cap keeps the tower's cost independent of frame rate, which is what makes
    the same model valid across a 10 Hz LiDAR-synced camera and a 60 Hz one.
    """
    fv = store.table(table).scan()
    by_stream = {}
    for s in sorted(set(fv.column("stream").to_pylist())):
        rows = fv.filter(pc.equal(fv.column("stream"), s))
        ts = rows.column("ts").to_numpy()
        order = np.argsort(ts)
        vecs = np.asarray(rows.column("vector").to_pylist(), dtype=np.float32)
        by_stream[s] = (ts[order], vecs[order])
    seqs = []
    for (s, t0, t1) in windows:
        ts, vecs = by_stream[s]
        lo, hi = np.searchsorted(ts, [t0, t1 + 1])
        idx = np.arange(lo, hi)
        if len(idx) > max_len:
            idx = idx[np.linspace(0, len(idx) - 1, max_len).round().astype(int)]
        seqs.append(vecs[idx])
    return seqs


# ---------------------------------------------------------------------------
# 3. The teacher — a VLM that actually reads the pixels, run ONCE per window
# ---------------------------------------------------------------------------
def caption_windows(store, windows, frames_per_window=3, model_id=None,
                    max_tokens=64, width=448, verbose=True, limit=None,
                    every=1, prompt="scene"):
    """VLM captions for `windows` → `context_captions` table.

    The VLM is shown several frames of the SAME window in order, so the
    caption can describe motion. A single-frame caption would be another
    appearance label and the student would learn nothing a mean-pool cannot
    already produce.
    """
    import tempfile
    from pathlib import Path

    from PIL import Image

    from .rerank import DEFAULT_VLM, _load
    from .video import FrameSet
    model_id = model_id or DEFAULT_VLM
    from mlx_vlm import generate
    from mlx_vlm.prompt_utils import apply_chat_template

    vlm, processor, cfg, _, _ = _load(model_id)
    rot = store.meta.get("display", {}).get("rotate", 0)
    frames = store.table("frames").scan()
    tmpdir = Path(tempfile.mkdtemp(prefix="elidedb_ctx_"))

    # `every` samples the window list uniformly instead of taking a prefix.
    # On a corpus too large to caption in full this matters: a prefix would
    # confine the caption vocabulary to whatever happens early in the
    # timeline, so query terms for anything later would be out of vocabulary
    # and the lexical ranker would abstain on the entire tail.
    if every > 1:
        windows = windows[::every]
    if limit:
        windows = windows[:limit]
    text = PROMPTS.get(prompt, prompt)          # preset name or raw prompt
    prompt = apply_chat_template(processor, cfg, text,
                                 num_images=frames_per_window)

    recs = {"ts": [], "t1": [], "stream": [], "caption": []}
    t_start = time.time()
    for n, (s, t0, t1) in enumerate(windows):
        sel = frames.filter(pc.and_(
            pc.equal(frames.column("stream"), s),
            pc.and_(pc.greater_equal(frames.column("ts"), t0),
                    pc.less_equal(frames.column("ts"), t1))))
        fs = FrameSet(store, "frames", sel)
        dec = fs.decode(width=width)
        if len(dec) < 1:
            continue
        picks = np.linspace(0, len(dec) - 1,
                            min(frames_per_window, len(dec))).round().astype(int)
        paths = []
        for j, p in enumerate(picks):
            im = Image.fromarray(dec[p][1])
            if rot:
                im = im.rotate(rot, expand=True)
            fp = tmpdir / f"w{n}_{j}.jpg"
            im.save(fp, "JPEG", quality=85)
            paths.append(str(fp))
        while len(paths) < frames_per_window:      # pad short windows
            paths.append(paths[-1])
        r = generate(vlm, processor, prompt, image=paths,
                     max_tokens=max_tokens, verbose=False)
        cap = tidy_caption(r.text if hasattr(r, "text") else str(r))
        recs["ts"].append(t0)
        recs["t1"].append(t1)
        recs["stream"].append(s)
        recs["caption"].append(cap)
        if verbose and n % 10 == 0:
            el = time.time() - t_start
            print(f"  [{n + 1}/{len(windows)}] {el:.0f}s  {s} +"
                  f"{(t0 - int(frames.column('ts')[0].as_py())) / 1e9:.1f}s :: "
                  f"{cap[:90]}", flush=True)
    if not recs["ts"]:
        return {"captions": 0}

    # Caption → frozen SigLIP TEXT tower. This is the whole point of using
    # SigLIP as the teacher's codec: the target lands in the SAME space a
    # user's query lands in, so the student is learning to be the image side
    # of a dual encoder — not to regress an arbitrary embedding.
    vecs = embed_texts(recs["caption"])
    dim = vecs.shape[1]
    tbl = pa.table({
        "ts": pa.array(recs["ts"], pa.int64()),
        "t1": pa.array(recs["t1"], pa.int64()),
        "stream": pa.array(recs["stream"]),
        "caption": pa.array(recs["caption"]),
        "vector": pa.array([v.tolist() for v in vecs],
                           pa.list_(pa.float32(), dim)),
    })
    version = store.table("context_captions").append(
        tbl, kind="embeddings",
        meta={"teacher": model_id, "text_model": DEFAULT_MODEL, "dim": dim,
              "frames_per_window": frames_per_window, "prompt": text[:200],
              "seconds": round(time.time() - t_start, 1)})
    return {"captions": len(tbl), "version": version,
            "seconds": round(time.time() - t_start, 1)}


def embed_texts(texts, model_id=DEFAULT_MODEL, batch=32):
    """Batched SigLIP text tower. Same normalisation as image vectors so the
    two are directly comparable by dot product. Delegates to the
    portable single-text path when mlx is not on this machine."""
    from .embeddings import _backend, embed_text
    if _backend() != "mlx":
        return np.stack([embed_text(t if t.strip() else "a scene",
                                    model_id) for t in texts])
    import mlx.core as mx
    model, processor = _load_model(model_id)
    out = []
    for i in range(0, len(texts), batch):
        chunk = [t if t.strip() else "a scene" for t in texts[i:i + batch]]
        ti = processor(text=chunk, padding="max_length", max_length=64,
                       truncation=True, return_tensors="np")
        v = np.array(model.get_text_features(mx.array(ti["input_ids"])),
                     dtype=np.float32)
        out.append(v / np.linalg.norm(v, axis=1, keepdims=True))
    return np.concatenate(out, axis=0)


# ---------------------------------------------------------------------------
# 4. CaptionSpace — the output space, chosen by measurement
# ---------------------------------------------------------------------------
# Three candidate spaces were benchmarked against a VLM judge on windows the
# tower never trained on (mean yes/no logprob margin over each method's top-5,
# six relational queries, higher is better):
#
#     appearance only  (SigLIP image-text)                +0.276
#     caption EMBEDDING (SigLIP text tower, oracle)       +0.218   <- worse
#     VLM rerank at query time (2.1 s/query)              +0.314
#     caption TEXT, LSA-48                                +0.339   <- winner
#     caption TEXT, raw TF-IDF, fused with appearance     +0.345
#
# The embedding route loses because SigLIP's text tower is trained to sit
# near IMAGES, not near other text; comparing a query embedding to a caption
# embedding uses a geometry the model was never optimised for. Matching the
# caption as TEXT sidesteps that entirely.
#
# LSA-48 is chosen over raw TF-IDF despite scoring 0.006 lower: it is dense
# and fixed-width, so (a) it is a target a small tower can actually regress,
# which is what lets unlabelled windows get a context vector at all, and
# (b) it is one more fixed_size_list column, so every index already in the
# store consumes it unchanged.


class CaptionSpace:
    """TF-IDF + LSA over the caption corpus. Queries and captions share it.

    Not persisted as a pickle. The captions themselves are the durable
    artifact — they live in `context_captions` as ordinary Parquet — and the
    lexical index is refit from them on load (a few ms for this corpus) and
    cached. That keeps the index unconditionally consistent with the data and
    free of any sklearn version pinning. At corpus sizes where refitting
    stops being free, persist the vocabulary and the SVD basis; the interface
    does not change.
    """

    def __init__(self, vec, svd):
        self.vec, self.svd = vec, svd
        self.dim = svd.n_components

    @staticmethod
    def fit(texts, dim=48):
        from sklearn.decomposition import TruncatedSVD
        from sklearn.feature_extraction.text import TfidfVectorizer
        vec = TfidfVectorizer(stop_words="english", ngram_range=(1, 2),
                              sublinear_tf=True).fit(texts)
        X = vec.transform(texts)
        dim = int(min(dim, X.shape[1] - 1, len(texts) - 1))
        svd = TruncatedSVD(n_components=dim, random_state=0).fit(X)
        return CaptionSpace(vec, svd)

    def transform(self, texts):
        v = self.svd.transform(self.vec.transform(list(texts)))
        v = np.asarray(v, dtype=np.float32)
        return v / (np.linalg.norm(v, axis=1, keepdims=True) + 1e-8)


_SPACE_CACHE: dict = {}


def caption_space(store, dim=48):
    t = store.table("context_captions")
    key = (str(store.dir), t.state().version, dim)
    if key not in _SPACE_CACHE:
        _SPACE_CACHE.clear()
        _SPACE_CACHE[key] = CaptionSpace.fit(
            t.scan().column("caption").to_pylist(), dim=dim)
    return _SPACE_CACHE[key]


# ---------------------------------------------------------------------------
# 5. The student — train the tower, then materialise context vectors
# ---------------------------------------------------------------------------
def _labelled(store, windows):
    """Join the caption table onto the window plan by (stream, ts)."""
    caps = store.table("context_captions").scan()
    key = {(s, t): i for i, (s, t) in enumerate(
        zip(caps.column("stream").to_pylist(), caps.column("ts").to_pylist()))}
    vecs = np.asarray(caps.column("vector").to_pylist(), dtype=np.float32)
    texts = caps.column("caption").to_pylist()
    keep, tv, tt = [], [], []
    for w in windows:
        i = key.get((w[0], w[1]))
        if i is not None:
            keep.append(w)
            tv.append(vecs[i])
            tt.append(texts[i])
    return keep, np.asarray(tv, dtype=np.float32), tt


def train_context(store, window_s=2.0, stride_s=0.5, val_frac=0.3,
                  d_in=48, d_out=48, epochs=400, cfg=None, verbose=True,
                  seed=0):
    """Fit the input PCA + tower to predict caption-LSA coordinates.

    The validation split is by TIME, never at random: windows slide with 75%
    overlap, so a random split puts near-duplicate windows on both sides and
    reports a score that measures nothing.
    """
    from .ctxtower import ContextCodec, retrieval_r1, save_tower, train_tower

    windows = plan_windows(store, window_s, stride_s)
    windows, _, cap_text = _labelled(store, windows)
    if len(windows) < 16:
        raise RuntimeError(f"only {len(windows)} captioned windows — run "
                           "caption_windows() first")

    fv = store.table("frame_vectors").scan()
    frame_mat = np.asarray(fv.column("vector").to_pylist(), dtype=np.float32)
    codec = ContextCodec.fit(frame_mat, d_in=d_in)
    space = caption_space(store, dim=d_out)

    seqs = [codec.encode(s) for s in window_sequences(store, windows)]
    targets = space.transform(cap_text)

    t0s = np.array([w[1] for w in windows])
    cut = np.quantile(t0s, 1.0 - val_frac)
    val_idx = np.where(t0s >= cut)[0]
    if verbose:
        print(f"  {len(windows)} labelled windows | train "
              f"{len(windows) - len(val_idx)} | val {len(val_idx)} "
              f"(time split at +{(cut - t0s.min()) / 1e9:.1f}s) | "
              f"target dim {space.dim}", flush=True)

    cfg = {**dict(d_in=codec.P_in.shape[0], d_out=space.dim), **(cfg or {})}
    model, info = train_tower(seqs, targets, windows, val_idx, cfg=cfg,
                              epochs=epochs, verbose=verbose, seed=seed)

    import mlx.core as mx
    val = np.zeros(len(seqs), bool)
    val[val_idx] = True
    Xva = _pad_stack([s for s, m in zip(seqs, val) if m])
    pred = np.array(model(mx.array(Xva)))
    pred /= np.linalg.norm(pred, axis=1, keepdims=True) + 1e-8
    tgt_va = targets[val]

    # Baseline: predict the TRAIN-set mean target for every window. That is
    # the best a model can do while knowing nothing about the specific window,
    # so any lift over it is genuine per-window information and not the corpus
    # prior leaking through.
    prior = targets[~val].mean(axis=0)
    prior /= np.linalg.norm(prior) + 1e-8
    prior = np.repeat(prior[None, :], len(tgt_va), axis=0)

    metrics = {
        "labelled_windows": len(windows),
        "train": int((~val).sum()), "val": int(val.sum()),
        "target_dim": space.dim,
        "val_R@1_tower": retrieval_r1(pred, tgt_va),
        "val_cos_tower": float((pred * tgt_va).sum(1).mean()),
        "val_cos_prior": float((prior * tgt_va).sum(1).mean()),
        "params": int(sum(v.size for v in _flat(model).values())),
        **{k: v for k, v in info.items() if k != "history"},
    }
    save_tower(model, codec,
               {"window_s": window_s, "stride_s": stride_s, "d_out": space.dim,
                "metrics": metrics, "history": info["history"]},
               store.dir / "models" / "context")
    return model, codec, metrics


def _flat(model):
    from mlx.utils import tree_flatten
    return {k: np.array(v) for k, v in tree_flatten(model.parameters())}


def _pad_stack(seqs):
    T = max(s.shape[0] for s in seqs)
    X = np.zeros((len(seqs), T, seqs[0].shape[1]), dtype=np.float32)
    for i, s in enumerate(seqs):
        X[i, :len(s)] = s
        if len(s) < T:
            X[i, len(s):] = s[-1]
    return X


def build_context(store, window_s=None, stride_s=None, batch=64, verbose=True):
    """Materialise the `context` table: one context vector per window.

    Where a window has a teacher caption, its EXACT caption-LSA vector is
    stored. Where it does not, the tower's prediction is stored and the row is
    flagged `estimated`. This is the ordinary database distinction between a
    materialised value and an estimated one, and it is the point of having a
    student at all: the VLM labels what you can afford, the tower covers the
    rest, and the query does not care which it got.
    """
    import time

    import mlx.core as mx

    from .ctxtower import load_tower
    model, codec, meta = load_tower(store.dir / "models" / "context")
    window_s = window_s or meta.get("window_s", 2.0)
    stride_s = stride_s or meta.get("stride_s", 0.5)
    space = caption_space(store, dim=meta.get("d_out", 48))

    windows = plan_windows(store, window_s, stride_s)
    raw = window_sequences(store, windows)
    seqs = [codec.encode(s) for s in raw]
    # The mean-pooled appearance vector for the SAME window, stored alongside.
    # Fusing appearance with context otherwise needs a join between two tables
    # built on different window plans; keeping both columns in one row makes
    # the fused query two matmuls over one Parquet scan and no join at all.
    appear = np.stack([s.mean(axis=0) for s in raw])
    appear /= np.linalg.norm(appear, axis=1, keepdims=True) + 1e-8

    t_start = time.time()
    out = []
    for i in range(0, len(seqs), batch):
        out.append(np.array(model(mx.array(_pad_stack(seqs[i:i + batch])))))
    z = np.concatenate(out, axis=0)
    z /= np.linalg.norm(z, axis=1, keepdims=True) + 1e-8
    infer_s = time.time() - t_start

    caps = store.table("context_captions").scan()
    known = {(a, b): c for a, b, c in zip(caps.column("stream").to_pylist(),
                                          caps.column("ts").to_pylist(),
                                          caps.column("caption").to_pylist())}
    have = [(i, known[(w[0], w[1])]) for i, w in enumerate(windows)
            if (w[0], w[1]) in known]
    estimated = np.ones(len(windows), bool)
    if have:
        exact = space.transform([c for _, c in have])
        for (i, _), v in zip(have, exact):
            z[i] = v
            estimated[i] = False

    dim, adim = z.shape[1], appear.shape[1]
    tbl = pa.table({
        "ts": pa.array([w[1] for w in windows], pa.int64()),
        "t1": pa.array([w[2] for w in windows], pa.int64()),
        "stream": pa.array([w[0] for w in windows]),
        "vector": pa.array([v.tolist() for v in z],
                           pa.list_(pa.float32(), dim)),
        "appearance": pa.array([v.tolist() for v in appear],
                               pa.list_(pa.float32(), adim)),
        "estimated": pa.array(estimated.tolist(), pa.bool_()),
    })
    st = store.table("context").state()
    meta_out = {"dim": dim, "window_s": window_s, "stride_s": stride_s,
                "estimated_rows": int(estimated.sum()),
                "exact_rows": int((~estimated).sum())}
    if st.files:                                 # replace: one atomic commit
        import uuid as _uuid

        from .log import FileEntry
        from .store import write_parquet
        fname = f"part-{_uuid.uuid4().hex[:12]}.parquet"
        p = store.dir / "tables" / "context" / fname
        tbl = tbl.take(pc.sort_indices(tbl.column("ts")))
        write_parquet(tbl, p)
        tsv = tbl.column("ts").to_numpy()
        version = store.table("context").log.commit(
            op="replace", kind="embeddings", schema=str(tbl.schema),
            add=[FileEntry(fname, len(tbl), p.stat().st_size,
                           int(tsv.min()), int(tsv.max()))],
            remove=[f.path for f in st.files], meta=meta_out)
    else:
        version = store.table("context").append(tbl, kind="embeddings",
                                                meta=meta_out)
    if verbose:
        print(f"  {len(tbl)} context vectors ({int((~estimated).sum())} exact, "
              f"{int(estimated.sum())} estimated) | tower inference "
              f"{infer_s * 1000:.0f} ms "
              f"({infer_s / len(tbl) * 1e6:.0f} us/window)")
    return {"windows": len(tbl), "dim": dim, "version": version,
            "inference_s": round(infer_s, 3),
            "us_per_window": round(infer_s / len(tbl) * 1e6, 1), **meta_out}


# ---------------------------------------------------------------------------
# 6. Post-training cellular turnover — capacity the data cannot support is
#    both latency and overfitting, so apoptosis pays twice.
# ---------------------------------------------------------------------------
def _prepare(store, window_s, stride_s, codec, val_frac, d_out):
    windows, _, texts = _labelled(store, plan_windows(store, window_s,
                                                      stride_s))
    seqs = [codec.encode(s) for s in window_sequences(store, windows)]
    targets = caption_space(store, dim=d_out).transform(texts)
    t0s = np.array([w[1] for w in windows])
    val = t0s >= np.quantile(t0s, 1.0 - val_frac)
    return windows, seqs, targets, val


def prune_context(store, val_frac=0.3, ppo_iters=40, sparsity_coef=0.15,
                  finetune_epochs=300, rebirth_fraction=0.5,
                  select_tolerance=0.03, verbose=True, seed=0):
    """apoptosis → re-settle → neurogenesis → re-settle, then COMPACT so the
    surviving channels are the only ones that cost anything."""
    import time

    import mlx.core as mx
    import mlx.nn as nn

    from .ctxprune import compact, get_channel_mask, run_pruning_cycle
    from .ctxtower import (load_tower, overlap_mask, retrieval_r1,
                           save_tower, siglip_loss)

    model, codec, meta = load_tower(store.dir / "models" / "context")
    windows, seqs, targets, val = _prepare(
        store, meta["window_s"], meta["stride_s"], codec, val_frac,
        meta.get("d_out", 48))

    Xtr = mx.array(_pad_stack([s for s, m in zip(seqs, ~val) if m]))
    Xva = mx.array(_pad_stack([s for s, m in zip(seqs, val) if m]))
    Ytr, Yva = mx.array(targets[~val]), mx.array(targets[val])
    ig_tr = mx.array(overlap_mask([w for w, m in zip(windows, ~val) if m]))
    ig_va = mx.array(overlap_mask([w for w, m in zip(windows, val) if m]))
    log_t = mx.array(np.float32(meta["metrics"]["log_t"]))
    bias = mx.array(np.float32(meta["metrics"]["bias"]))

    def _norm(a):
        return a * mx.rsqrt(mx.sum(a * a, axis=-1, keepdims=True) + 1e-8)

    def _loss(m, x, y, ig):
        v = _norm(m(x))
        u = _norm(y)
        return siglip_loss(v, u, ignore=ig, log_t=log_t, bias=bias) \
            + 0.3 * mx.mean(1.0 - mx.sum(v * u, axis=-1))

    def loss_of(m):
        """The number PPO is scored against: VALIDATION retrieval loss. Using
        train loss here would reward a policy for keeping memorisers."""
        m.set_training(False)
        return float(_loss(m, Xva, Yva, ig_va).item())

    lg = nn.value_and_grad(model, lambda m: _loss(m, Xtr, Ytr, ig_tr))

    def r1(m):
        m.set_training(False)
        return retrieval_r1(np.array(_norm(m(Xva))), np.array(_norm(Yva)))

    def latency(m, reps=20):
        m.set_training(False)
        mx.eval(m(Xva))
        t = time.perf_counter()
        for _ in range(reps):
            mx.eval(m(Xva))
        return (time.perf_counter() - t) / reps / Xva.shape[0] * 1e6

    before = {"channels": int(get_channel_mask(model).sum()),
              "val_loss": loss_of(model), "val_R@1": r1(model),
              "us_per_window": latency(model),
              "params": int(sum(v.size for v in _flat(model).values()))}

    rec = run_pruning_cycle(model, np.array(Xtr), loss_of, lg,
                            ppo_iters=ppo_iters, sparsity_coef=sparsity_coef,
                            finetune_epochs=finetune_epochs,
                            rebirth_fraction=rebirth_fraction,
                            select_tolerance=select_tolerance, seed=seed,
                            verbose=verbose)

    # Compaction must be a no-op numerically — it only deletes channels the
    # mask already zeroed. Measure both sides and say so if they disagree,
    # rather than quietly shipping a tower that differs from the one the
    # search selected.
    restored = loss_of(model)
    model, kept = compact(model)
    after = {"channels": int(len(kept)), "val_loss": loss_of(model),
             "loss_before_compaction": restored, "val_R@1": r1(model),
             "us_per_window": latency(model),
             "params": int(sum(v.size for v in _flat(model).values()))}
    rec["before"], rec["after"] = before, after
    drift = abs(after["val_loss"] - restored)
    if drift > 1e-3:
        print(f"  WARNING: compaction changed val loss by {drift:.4f} "
              f"({restored:.4f} -> {after['val_loss']:.4f}) — expected ~0")
    if verbose:
        print(f"\n  channels {before['channels']} -> {after['channels']} | "
              f"params {before['params']:,} -> {after['params']:,} | "
              f"{before['us_per_window']:.1f} -> {after['us_per_window']:.1f} "
              f"us/window | val loss {before['val_loss']:.4f} -> "
              f"{after['val_loss']:.4f}")
    save_tower(model, codec,
               {**{k: v for k, v in meta.items() if k != "history"},
                "pruned": {"before": before, "after": after,
                           "ppo_history": rec.get("ppo_history"),
                           "keep_probs": rec.get("keep_probs"),
                           "reverse_attention": rec.get("reverse_attention"),
                           "selected": rec.get("selected"),
                           "stages": [{k: v for k, v in s.items()
                                       if k != "mask"} for s in rec["stages"]]}},
               store.dir / "models" / "context")
    return model, rec


# ---------------------------------------------------------------------------
# 7. The query path — two matmuls, no VLM, no decode
# ---------------------------------------------------------------------------
_CTX_CACHE: dict = {}


def _ctx_matrix(store):
    t = store.table("context")
    key = (str(store.dir), t.state().version)
    if key not in _CTX_CACHE:
        _CTX_CACHE[key] = np.stack([
            np.asarray(v, dtype=np.float32)
            for v in t.scan().column("vector").to_pylist()])
    return _CTX_CACHE[key]


def _lexical(store):
    """(vectorizer, matrix, has_caption) aligned to the `context` table rows.

    Windows the VLM never captioned have no text, so they are marked as
    abstentions rather than as empty documents — an empty document would score
    0 on every query and be ranked last by the lexical ranker, which is a veto
    dressed up as evidence.
    """
    t = store.table("context")
    key = (str(store.dir), t.state().version, "lex")
    if key in _CTX_CACHE:
        return _CTX_CACHE[key]
    ctx = t.scan()
    rows = list(zip(ctx.column("stream").to_pylist(),
                    ctx.column("ts").to_pylist()))
    caps = store.table("context_captions").scan()
    known = dict(zip(zip(caps.column("stream").to_pylist(),
                         caps.column("ts").to_pylist()),
                     caps.column("caption").to_pylist()))
    texts = [known.get(r, "") for r in rows]
    has = np.array([bool(x) for x in texts])
    space = caption_space(store)
    X = space.vec.transform(texts)
    nrm = np.sqrt(np.asarray(X.multiply(X).sum(1))).ravel() + 1e-8
    _CTX_CACHE[key] = (space, X, nrm, has)
    return _CTX_CACHE[key]


DEFAULT_WEIGHTS = {"appearance": 1.0, "context": 1.0, "lexical": 1.0}


def search(store, text, k=10, merge=True, t0=None, t1=None, streams=None,
           weights=None, rrf_k=60.0, neg_weight=0.5, min_score=None,
           percentile=None, rerank=False, rerank_top=12, rerank_alpha=0.7,
           explain_top=0):
    """Hybrid contextual search: three rankers fused by reciprocal rank.

        appearance  SigLIP image-text cosine over the window's frames.
                    Knows what OBJECTS are present. Order-blind.
        context     caption-LSA cosine. Knows what is HAPPENING, because the
                    caption was written by a VLM that watched three frames.
        lexical     TF-IDF over the caption text. Exact term evidence — the
                    ranker that actually knows what "red" means.

    Fusing by RRF rather than by a weighted score sum is the fix for
    "crossing red car" returning any clip of someone crossing: RRF rewards
    agreement across rankers, so a candidate that satisfies one strong signal
    alone can no longer win. See elidedb.fusion.

    `rerank=True` adds a final VLM pass over the top `rerank_top` — the
    expensive operator, last, on an already-pruned set.
    """
    from .embeddings import _parse_query, _score_windows, embed_text
    from .fusion import explain_fusion, rrf
    # Fail with a sentence, not a KeyError from three frames down. An empty
    # table has no schema, so the first column access explodes with
    # 'Field "t1" does not exist' — true, useless, and it names the wrong
    # problem.
    if not store.table("context").state().files:
        raise RuntimeError(
            f"store '{store.name}' has no context index. Build it with "
            "store.index_context() (frames -> per-frame vectors -> VLM "
            "captions -> context table), or use store.search_text() for "
            "appearance-only search.")
    pos, neg = _parse_query(text)
    ctx_tbl = store.table("context").scan()
    all_t0 = ctx_tbl.column("ts").to_numpy()
    all_t1 = ctx_tbl.column("t1").to_numpy()
    all_s = ctx_tbl.column("stream").to_numpy(zero_copy_only=False)

    # hybrid retrieval: time and stream predicates are pushed INTO candidate
    # selection, not applied to a global top-k afterwards
    pred = np.ones(len(all_t0), bool)
    if t0 is not None:
        pred &= all_t1 >= t0
    if t1 is not None:
        pred &= all_t0 <= t1
    if streams:
        pred &= np.isin(all_s, list(streams))
    idx = np.where(pred)[0]
    if len(idx) == 0:
        return [], {"index": "context", "total": len(all_t0), "scanned": 0,
                    "segments": 0}

    APP = _appearance_matrix(store)
    CTX = _ctx_matrix(store)
    space, X, nrm, has_cap = _lexical(store)

    pos_app = np.stack([embed_text(p) for p in pos])
    neg_app = np.stack([embed_text(n) for n in neg]) if neg else None
    app = _score_windows(APP, idx, pos_app, neg_app, neg_weight)

    q_join = " ".join(pos)
    ctx_sc = CTX[idx] @ space.transform([q_join])[0]

    qv = space.vec.transform([q_join])
    lex = np.asarray((X[idx] @ qv.T).todense()).ravel() / nrm[idx]
    lex = np.where(has_cap[idx], lex, np.nan)      # abstain, do not veto

    rankings = {"appearance": app, "context": ctx_sc, "lexical": lex}
    w = {**DEFAULT_WEIGHTS, **(weights or {})}
    scores = rrf(rankings, w, rrf_k)

    stats = {"index": "context", "total": len(all_t0), "scanned": len(idx),
             "method": "rrf", "rrf_k": rrf_k, "weights": w,
             "predicate_candidates": int(pred.sum()),
             "captioned_candidates": int(has_cap[idx].sum()),
             "positive_terms": pos, "negative_terms": neg}
    if explain_top:
        stats["why"] = explain_fusion(rankings, w, rrf_k, top=explain_top)

    keep = np.ones(len(idx), bool)
    if percentile is not None:
        keep &= scores >= np.percentile(scores, percentile)
    if min_score is not None:
        keep &= scores >= min_score
    idx, scores = idx[keep], scores[keep]
    stats["after_floor"] = int(len(idx))

    hits = _segments(idx, scores, all_s, all_t0, all_t1, k, merge, stats)
    if rerank and hits:
        from .rerank import rerank_hits
        hits, info = rerank_hits(store, hits, text, top_n=rerank_top,
                                 alpha=rerank_alpha)
        stats["rerank"] = info
    return hits, stats


def _segments(idx, scores, all_s, all_t0, all_t1, k, merge, stats):
    """Merge qualifying windows into maximal runs per stream.

    Fixed windows are an INDEXING granularity, not an answer granularity: a
    20 s event should come back as one 20 s hit, and a query matching only 2 s
    of it should come back as that 2 s.
    """
    streams_sel, w_t0, w_t1 = all_s[idx], all_t0[idx], all_t1[idx]
    if len(idx) == 0:
        stats["segments"] = 0
        return []
    if not merge:
        order = np.argsort(scores)[::-1][:k]
        return [{"stream": str(streams_sel[i]), "t0": int(w_t0[i]),
                 "t1": int(w_t1[i]), "score": float(scores[i]),
                 "windows": 1} for i in order]
    med, top = float(np.median(scores)), float(scores.max())
    thr = med + 0.55 * (top - med)
    stats["threshold"] = round(thr, 6)
    qual = np.where(scores >= thr)[0]
    order = np.lexsort((w_t0[qual], streams_sel[qual]))
    qual = qual[order]
    gap = int(np.median(w_t1[qual] - w_t0[qual])) + 1 if len(qual) else 0
    segs = []
    for i in qual:
        s_, a, b, sc = (str(streams_sel[i]), int(w_t0[i]), int(w_t1[i]),
                        float(scores[i]))
        last = segs[-1] if segs else None
        if last and last["stream"] == s_ and a - last["t1"] <= gap:
            last["t1"] = max(last["t1"], b)
            last["score"] = max(last["score"], sc)
            last["mean"] = (last["mean"] * last["windows"] + sc) / (last["windows"] + 1)
            last["windows"] += 1
        else:
            segs.append({"stream": s_, "t0": a, "t1": b, "score": sc,
                         "mean": sc, "windows": 1})
    segs.sort(key=lambda g: -g["score"])
    stats["qualifying_windows"] = len(qual)
    stats["segments"] = len(segs)
    return segs[:k]


def _appearance_matrix(store):
    t = store.table("context")
    key = (str(store.dir), t.state().version, "app")
    if key not in _CTX_CACHE:
        _CTX_CACHE[key] = np.stack([
            np.asarray(v, dtype=np.float32)
            for v in t.scan().column("appearance").to_pylist()])
    return _CTX_CACHE[key]


def explain(store, t0, t1, stream=None):
    """What the database believes is happening in a window — the teacher's own
    words. Lets a result be checked rather than trusted."""
    caps = store.table("context_captions").scan()
    out = []
    if len(caps) == 0 or "stream" not in caps.column_names:
        return out          # store has no captions (yet) — nothing to explain
    for s, a, b, c in zip(caps.column("stream").to_pylist(),
                          caps.column("ts").to_pylist(),
                          caps.column("t1").to_pylist(),
                          caps.column("caption").to_pylist()):
        if stream and s != stream:
            continue
        if b >= t0 and a <= t1:
            out.append({"stream": s, "t0": a, "t1": b, "caption": c})
    return out


def _embed_frames_fdnnv(store, frame_table, incremental, verbose, streams):
    """Every frame through the FDNN-V streaming encoder -> frame_vectors."""
    import pyarrow as pa

    from .fdnnvideo import embed_stream, load_encoder
    mdir = store.dir / "models" / "fdnnv"
    if not (mdir / "encoder.json").exists():
        # a NEW store has no encoder yet — adopt one from a sibling store
        # and COPY it in, so the store stays self-contained and the exact
        # weights that wrote its vectors are pinned with its data
        import shutil
        donors = sorted(store.dir.parent.glob("*/models/fdnnv/encoder.json"),
                        key=lambda p: p.stat().st_mtime, reverse=True)
        if not donors:
            raise RuntimeError(
                "no FDNN-V encoder found in this store or any sibling — "
                "train one first (scripts/fdnnv_train.py)")
        shutil.copytree(donors[0].parent, mdir)
        if verbose:
            print(f"  adopted encoder from {donors[0].parent}", flush=True)
    model, meta = load_encoder(mdir)
    # FRESH-DATA FIDELITY GATE: the fast student was distilled on one
    # style of footage; on an arbitrary upload its fidelity to the teacher
    # is unknown. 32 of THIS store's frames go through both encoders; if
    # mean cosine < 0.90 the store gets the TEACHER (slower ingest, right
    # space) instead of a fast-but-wrong index. Automatic — a fresh
    # customer never has to know this exists.
    try:
        from PIL import Image
        from .embeddings import DEFAULT_MODEL, _embed_images
        from .video import FrameSet
        frames_all = store.table(frame_table).scan()
        pick = np.linspace(0, len(frames_all) - 1,
                           min(32, len(frames_all))).round().astype(int)
        rows = frames_all.take(pick)
        ts_s, sv, _, _ = embed_stream(store, model, rows)
        dec = FrameSet(store, frame_table, rows).decode(width=448)
        imgs = [Image.fromarray(d[1]) for d in sorted(dec)]
        tv = _embed_images(imgs, DEFAULT_MODEL)
        n = min(len(sv), len(tv))
        svn = sv[:n] / (np.linalg.norm(sv[:n], axis=1,
                                       keepdims=True) + 1e-8)
        fid = float((svn * tv[:n]).sum(1).mean())
        if verbose:
            print(f"  student fidelity on this corpus: {fid:.3f}",
                  flush=True)
        if fid < 0.90:
            print(f"  fidelity {fid:.3f} < 0.90 — falling back to the "
                  f"TEACHER encoder for this store", flush=True)
            return embed_frames(store, frame_table, model="fast",
                                incremental=incremental, verbose=verbose,
                                streams=streams, engine="siglip")
    except Exception as e:
        if verbose:
            print(f"  fidelity gate skipped ({type(e).__name__})",
                  flush=True)
    frames = store.table(frame_table).scan()
    allst = sorted(set(frames.column("stream").to_pylist()))
    use = [s for s in allst if s in streams] if streams else allst
    done = {}
    if incremental:
        try:
            prev = store.table("frame_vectors").scan()
            for s_, t_ in zip(prev.column("stream").to_pylist(),
                              prev.column("ts").to_pylist()):
                done[s_] = max(done.get(s_, -1), t_)
        except Exception:
            pass
    t_start = time.time()
    rows_ts, rows_stream, rows_vec = [], [], []
    for s in use:
        sel = frames.filter(pc.equal(frames.column("stream"), s))
        if s in done:
            sel = sel.filter(pc.greater(sel.column("ts"), done[s]))
        if len(sel) == 0:
            continue
        sel = sel.take(pc.sort_indices(sel.column("ts")))
        ts, vecs, dec_s, emb_s = embed_stream(store, model, sel)
        if verbose:
            print(f"  {s}: {len(ts):,} frames (decode {dec_s:.1f}s, "
                  f"embed {emb_s:.1f}s)", flush=True)
        rows_ts.extend(int(t) for t in ts)
        rows_stream.extend([s] * len(ts))
        rows_vec.extend(vecs)
    if not rows_ts:
        return {"frames": 0, "note": "nothing new (incremental)"}
    dim = len(rows_vec[0])
    tbl = pa.table({
        "ts": pa.array(rows_ts, pa.int64()),
        "stream": pa.array(rows_stream),
        "vector": pa.array([v.tolist() for v in rows_vec],
                           pa.list_(pa.float32(), dim)),
    })
    version = store.table("frame_vectors").append(
        tbl, kind="embeddings",
        meta={"model": "fdnnv", "teacher": meta.get("teacher"),
              "dim": dim, "every_frame": True,
              "source_table": frame_table})
    return {"frames": len(tbl), "dim": dim, "version": version,
            "engine": "fdnnv",
            "seconds": round(time.time() - t_start, 1)}