# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/index/00_utils.ipynb.

# %% auto 0
__all__ = ['read_ems', 'create_index', 'index_heap', 'random_ems']

# %% ../../nbs/index/00_utils.ipynb 2
from ..imports import *
from ..utils import *
from ..mapper import *


# %% ../../nbs/index/00_utils.ipynb 4
def read_ems(df, ems_col="embedding", ems_key="embedding"):
    ems = df[ems_col]
    if path_or_str(ems):
        with open(df[ems_col], "r") as f:
            ems = np.expand_dims(json.load(f)[ems_key], 0)
    else:
        ems = np.expand_dims(np.array(df[ems_col]), 0)
    return ems


def create_index(dim=768):
    return faiss.IndexFlatL2(dim)


def index_heap(
    df, k=1, dist_col="distances", ids_col="ids", size_col="index_size", verbose=False
):
    rh = faiss.ResultHeap(nq=1, k=k)

    ds = df[dist_col].tolist()
    # if verbose:
    # msg.info(f"First 10 Distances: {ds[:10]}")
    ids = df[ids_col].tolist()
    # if verbose:
    # msg.info(f"First 10 IDs: {ids[:10]}")
    sz = df[size_col].tolist()
    # if verbose:
    # msg.info(f"First 10 Index Sizes: {sz[:10]}")
    for i in range(len(ds)):
        offset = sum(sz[:i])
        if verbose:
            msg.info(f"Adding Result: {ds[i]}, {ids[i] + offset}")
        rh.add_result(ds[i], ids[i] + offset)
        if verbose:
            msg.good(f"Added Result: {ds[i]}, {ids[i] + offset}")
    rh.finalize()
    return {"distances": rh.D, "ids": rh.I}


def random_ems(num_ems=100, ems_dim=768, ems_folder="embeddings"):
    "Generate random embeddings and save them to a folder"

    ems = [np.random.random((1, ems_dim))[0].tolist() for i in range(num_ems)]
    os.makedirs(ems_folder, exist_ok=True)
    for i, em in enumerate(ems):
        em_path = str(Path(ems_folder) / f"em_{i}.json")
        json.dump({"embedding": em}, open(em_path, "w"))
