import os
import pickle

import faiss
import numpy as np


class CodeEmbedding:
    """Class to manage  Embeddings for a Doczy.ai execution"""

    def __init__(
        self,
        level,
        index_name="faiss_index.bin",
        embedding_name="embeddings.npy",
        choices_name="choices.pkl",
    ):
        self.level = level
        self.index = faiss.read_index(os.path.join("embeddings", level, index_name))
        self.embeddings = np.load(os.path.join("embeddings", level, embedding_name))
        with open(os.path.join("embeddings", level, choices_name), "rb") as f:
            self.choices = pickle.load(f)