# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/index/01_df.ipynb.

# %% auto 0
__all__ = ['df_to_index', 'df_index_search']

# %% ../../nbs/index/01_df.ipynb 2
from ..imports import *
from ..utils import *
from ..mapper import *
from .utils import *


# %% ../../nbs/index/01_df.ipynb 4
def df_to_index(
    df,
    index,
    ems_col="embedding",  # Column name of embeddings in df.
    ems_key="embedding",  # Key name of embeddings in json file.
    verbose=False,  # Whether to print out information.
):
    "Add the embeddings in df to the index."

    ems = read_ems(df, ems_col=ems_col, ems_key=ems_key)
    if verbose:
        msg.info(f"Ems Shape: {ems.shape}")
    index.add(ems)
    if verbose:
        msg.info(f"Index Size: {index.ntotal}")
    return df


def df_index_search(
    df,
    ems_col="embedding",  # Column name of embeddings in df.
    ems_key="embedding",  # Key name of embeddings in json file.
    index_col="index",  # Column name of index in df.
    k=1,  # Number of nearest neighbors to return.
    verbose=False,  # Whether to print out information.
):
    "Find the nearest neighbors of the embeddings in df."

    if verbose:
        msg.info(f"Index Col: {df[index_col]}")
    index = faiss.read_index(str(df[index_col]))
    if verbose:
        msg.info(f"Index Size: {index.ntotal}")
    ems = read_ems(df, ems_col=ems_col, ems_key=ems_key)
    if verbose:
        msg.info(f"Ems Shape: {ems.shape}")
    d, i = index.search(ems, k)
    if verbose:
        msg.good(f"IDs: {i}, Distances: {d}")
    df["index_size"] = index.ntotal
    df["distances"] = d  # .tolist()[0]
    df["ids"] = i  # .tolist()[0]
    return df
