from __future__ import annotations

from pathlib import Path
import numpy as np


class EmbeddingTable:
    """
    bf16 memory map embedding lookup.
    Lookup(ids) returns float32 (len(ids), dim).
    Exact conversion: (memmap[ids].astype(np.uint32) << 16).view(np.float32)
    """
    def __init__(self, path: str | Path, rows: int, dim: int) -> None:
        self.path = Path(path)
        self.rows = int(rows)
        self.dim = int(dim)
        # raw little-endian uint16
        self.mmap = np.memmap(self.path, dtype="<u2", mode="r", shape=(self.rows, self.dim))

    def lookup(self, ids: np.ndarray | list[int]) -> np.ndarray:
        ids_arr = np.asarray(ids, dtype=np.int64)
        bf16_rows = self.mmap[ids_arr]  # (..., dim) uint16
        u32 = bf16_rows.astype(np.uint32) << 16
        fp32 = u32.view(np.float32)
        return np.ascontiguousarray(fp32)