"""Where blocking gets its embedding matrices from.

Two interchangeable sources, selected by EMB["source"] (env M3_EMB_SOURCE):

  "internal"       the shard cache written by embed.py / encoder.py  (default)
  "hybrid_bridge"  the pre-computed multilingual-e5-large matrices already on
                   disk in /scratch/user4/dai/hybrid_bridge

The hybrid_bridge matrices were produced with, verified on 2026-09-26 by
re-encoding 12 records and comparing against the stored vectors:

    text   = "name: {business_name} | addr: {business_address} | country: {country}"
    S1     = "query: "   + text        (mean cos vs stored 0.99902)
    S2/S3  = "passage: " + text        (mean cos vs stored 0.99782)
    model  = intfloat/multilingual-e5-large (MIT), 1024-d, fp16, L2-normalised
    max_len= 96   (measured: p99 = 58 tokens, so nothing is truncated)

Row order is the exact row order of the corresponding parquet in
dataset/dataset/hf/{split}_p/, which is itself the row order of the original
TSVs. load() asserts the row counts, so a silent misalignment cannot pass.

MEASURED TRADE-OFF (20k real train queries, full 10.32M pool, country filter on,
dense channel only) -- see README "Embedding source":

    e5-large  (hybrid_bridge)  @20 = 0.9559   @30 = 0.9619   @50 = 0.9681
    e5-small  (internal)       @20 = 0.9652

The 1024-d e5-large matrices are 21 GB and score BELOW the 384-d e5-small ones
at the same k. Re-encoding the S1 side to make the prefixes symmetric was tested
and is worse still (passage:/passage: @20 = 0.9319), so the asymmetry is not the
cause. Use "hybrid_bridge" only with a larger emb_k, or to union the two dense
channels; it is not a drop-in improvement.
"""
from __future__ import annotations

import os
import numpy as np

import config as C


class ConcatMatrix:
    """Read-only row-concatenation of several memory-mapped (n_i, dim) matrices.

    Avoids materialising a 21 GB copy just to present S2+S3 as one matrix.
    Supports the two access patterns blocking uses: contiguous slices and
    integer-array fancy indexing.
    """

    def __init__(self, paths):
        self.parts = [np.load(p, mmap_mode="r") for p in paths]
        dims = {p.shape[1] for p in self.parts}
        if len(dims) != 1:
            raise SystemExit(f"dim mismatch across parts: {dims}")
        self.offs = np.cumsum([0] + [len(p) for p in self.parts])
        self.shape = (int(self.offs[-1]), self.parts[0].shape[1])
        self.dtype = self.parts[0].dtype

    def __len__(self):
        return self.shape[0]

    def _slice(self, a, b):
        out = np.empty((b - a, self.shape[1]), dtype=self.dtype)
        for i, part in enumerate(self.parts):
            lo, hi = int(self.offs[i]), int(self.offs[i + 1])
            s, e = max(a, lo), min(b, hi)
            if s < e:
                out[s - a:e - a] = part[s - lo:e - lo]
        return out

    def __getitem__(self, key):
        if isinstance(key, slice):
            a, b, _ = key.indices(self.shape[0])
            return self._slice(a, b)
        if isinstance(key, (int, np.integer)):
            return self._slice(int(key), int(key) + 1)[0]
        idx = np.asarray(key)
        out = np.empty((len(idx), self.shape[1]), dtype=self.dtype)
        for i, part in enumerate(self.parts):
            lo, hi = int(self.offs[i]), int(self.offs[i + 1])
            m = (idx >= lo) & (idx < hi)
            if m.any():
                sub = idx[m] - lo
                order = np.argsort(sub)            # memmap gather wants sorted reads
                rows = part[sub[order]]
                dest = np.where(m)[0][order]
                out[dest] = rows
        return out


def _check(mat, expect_rows, what):
    if mat.shape[0] != expect_rows:
        raise SystemExit(f"{what}: {mat.shape[0]:,} rows but parquet has "
                         f"{expect_rows:,} -- embeddings are not row-aligned")


def load_emb2(split, n_s1, n_s23):
    """Second dense channel: the pre-computed e5-large matrices, as (pool, query).

    No encoding is needed -- the S1 vectors are already on disk in parquet row
    order, so they are indexed directly by row.
    """
    H = C.HYBRID
    if split == "train":
        q = np.load(H["train_s1"], mmap_mode="r")
        c = ConcatMatrix([H["train_s2"], H["train_s3"]])
    elif split == "test":
        q = np.load(H["test_s1"], mmap_mode="r")
        c = ConcatMatrix([H["test_s23"]])
    else:
        raise SystemExit(f"unknown split {split!r}")
    _check(q, n_s1, f"{split} S1 emb2")
    _check(c, n_s23, f"{split} S2+S3 emb2")
    return c, q


def load(split, s1_ids, s23_ids):
    """-> (pool_matrix, query_matrix, meta) for split in {"train","test"}.

    s1_ids / s23_ids are the id lists in the order the pipeline built them and
    are used purely to validate row counts (and, cheaply, the first/last id).
    """
    src = C.EMB.get("source", "internal")
    if src == "internal":
        from encoder import load_matrix, read_ids
        pool_dir, q_dir = C.emb_dir(f"{split}_pool"), C.emb_dir(f"{split}_s1")
        c_emb, cmeta = load_matrix(pool_dir)
        q_emb, _ = load_matrix(q_dir)
        if list(s23_ids) != read_ids(os.path.join(pool_dir, "ids.txt")):
            raise SystemExit("pool embedding order != S2+S3 order")
        if list(s1_ids) != read_ids(os.path.join(q_dir, "ids.txt")):
            raise SystemExit("S1 embedding order != S1 order")
        return c_emb, q_emb, cmeta

    if src != "hybrid_bridge":
        raise SystemExit(f"unknown EMB['source'] {src!r}")

    H = C.HYBRID
    if split == "train":
        q_emb = np.load(H["train_s1"], mmap_mode="r")
        c_emb = ConcatMatrix([H["train_s2"], H["train_s3"]])
    elif split == "test":
        q_emb = np.load(H["test_s1"], mmap_mode="r")
        c_emb = ConcatMatrix([H["test_s23"]])
    else:
        raise SystemExit(f"unknown split {split!r}")

    _check(q_emb, len(s1_ids), f"{split} S1 embeddings")
    _check(c_emb, len(s23_ids), f"{split} S2+S3 embeddings")
    meta = {"model": H["model"], "dim": H["dim"], "n": c_emb.shape[0],
            "source": "hybrid_bridge", "template": H["template"],
            "query_prefix": H["query_prefix"], "doc_prefix": H["doc_prefix"]}
    return c_emb, q_emb, meta