"""The `pe` channel: Meta Perception Encoder (published SOTA zero-shot
video-text retrieval), adopted by measurement — best single-encoder
precision@10 on the bench (48/140 vs shipping fusion 40), complementary
profile (put-into-drawer 10/10 where others fail). Naive equal fusion
measured WORSE than PE alone, so this channel routes through the learned
per-query-type weights like every other.

8 frame vectors per recording in `pe_vectors`; recording score =
top-5-mean of frame cosines (the winning pooling from the zero-shot
sweep). Query cost: one open_clip text forward (~15 ms, cached)."""
from __future__ import annotations

import numpy as np

_TEXT = {}
_IDX = {}


def _text_vec(text):
    import torch
    if text in _TEXT.get("cache", {}):
        return _TEXT["cache"][text]
    if "model" not in _TEXT:
        import open_clip

        from .device import pick, strip_vision, text_only
        dev, dtype = pick()
        m, _, _ = open_clip.create_model_and_transforms(
            "PE-Core-L-14-336", pretrained="meta")
        m = strip_vision(m.to(dev).eval(), "visual")
        if text_only():
            # serving keeps only the text tower; at reduced precision
            # it is a few hundred MB instead of the fp32 gigabytes
            m = m.to(dtype)
        _TEXT["model"] = m
        _TEXT["dtype"] = m.ln_final.weight.dtype \
            if hasattr(m, "ln_final") else None
        _TEXT["tok"] = open_clip.get_tokenizer("PE-Core-L-14-336")
        _TEXT["dev"] = dev
        _TEXT["cache"] = {}
    with torch.no_grad():
        t = _TEXT["model"].encode_text(
            _TEXT["tok"]([text]).to(_TEXT["dev"]))
    v = (t / t.norm(dim=-1, keepdim=True))[0].cpu().float().numpy()
    if len(_TEXT["cache"]) > 256:
        _TEXT["cache"].clear()
    _TEXT["cache"][text] = v
    return v


def pe_lookup(store, text):
    """(lookup(stream, t0, t1) -> top5-mean cosine | nan, candidates)."""
    from .embeddings import _vec_table
    ver = store.table("pe_vectors").state().version
    key = (str(store.dir), ver)
    if key not in _IDX:
        tbl, _ = _vec_table(store, "pe_vectors")
        ss = np.asarray(tbl.column("stream").to_pylist())
        sa = np.asarray([int(v) for v in tbl.column("ts").to_pylist()])
        sb = np.asarray([int(v) for v in tbl.column("t1").to_pylist()])
        # rows grouped by recording (stream, t0, t1): build run map
        recs = {}
        order = np.lexsort((sa, ss))
        for r in order:
            recs.setdefault((str(ss[r]), int(sa[r]), int(sb[r])),
                            []).append(int(r))
        idx = {}
        for (s, a, b), rows in recs.items():
            idx.setdefault(s, []).append((a, b, np.array(rows)))
        for s in idx:
            idx[s].sort(key=lambda x: x[0])
        if len(_IDX) > 8:
            _IDX.clear()
        _IDX[key] = idx
    idx = _IDX[key]
    _, vecs = _vec_table(store, "pe_vectors")
    sc = vecs @ _text_vec(text)

    def rec_score(rows):
        v = sc[rows]
        k = min(5, len(v))
        return float(np.sort(v)[-k:].mean())

    def lookup(s, a, b):
        lst = idx.get(s)
        if not lst:
            return float("nan")
        starts = [x[0] for x in lst]
        j = int(np.searchsorted(starts, a, side="right")) - 1
        if j >= 0 and b <= lst[j][1] + 1:
            return rec_score(lst[j][2])
        return float("nan")

    cands = []
    for s, lst in idx.items():
        for a, b, rows in lst:
            cands.append((s, a, b, rec_score(rows)))
    cands.sort(key=lambda x: -x[3])
    return lookup, cands[:64]