# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/pdf/00_pdf_utils.ipynb.

# %% auto 0
__all__ = ['cid_to_char', 'process_text', 'bold_text', 'print_doc', 'create_pdf_df', 'create_idx_q_df', 'doc_proc_text',
           'df_pdf_docs_', 'df_pdf_docs', 'text_cat', 'docs_cat', 'df_docs_cat', 'docs_ems', 'df_docs_ems',
           'df_to_faiss', 'df_search_faiss']

# %% ../../nbs/pdf/00_pdf_utils.ipynb 2
from dreamai.imports import *
from ..imports import *
from ..chains import *

# %% ../../nbs/pdf/00_pdf_utils.ipynb 3
def cid_to_char(cidx: str):
    try:
        return chr(int(re.findall(r"\(cid\:(\d+)\)", cidx)[0]) + 29)
    except:
        return cidx


def process_text(text: str):
    text = text.strip()
    text = cid_to_char(text)
    text = re.sub("\xa0", " ", text)
    text = re.sub(r"\uf0b7", " ", text)
    text = re.sub(r"\(cid:\d{0,3}\)", " ", text)
    text = re.sub(r"•", "", text)
    text = re.sub(r"●", "", text)
    text = re.sub(r"▪", "", text)
    text = re.sub(r"", "", text)
    text = re.sub(r"➢", "", text)
    text = re.sub(r"\u2b9a", "", text)
    text = re.sub(r"\u201c", "", text)
    text = re.sub(r"\u201d", "", text)
    text = re.sub(r"\u2013", " ", text)
    text = re.sub(r"\u2019", "'", text)
    text = re.sub(r"\u2018", "'", text)
    text = re.sub(r"\u00f4", " ", text)
    text = re.sub(r"\u00f6", "o", text)
    text = re.sub(r"\u00e9", "e", text)
    text = re.sub(r"\u00e8", "e", text)
    text = re.sub(r"\u00e7", " ", text)
    text = re.sub(r"\u00a7", "", text)
    text = re.sub(r"\u00e3", "a", text)
    text = re.sub(r"\uf0a7", "", text)
    text = re.sub(r"\uf076", "", text)
    text = re.sub(r"\u00ad", "", text)
    text = re.sub(r"\u00ab", "", text)
    text = re.sub(r"\u00bb", "", text)
    text = re.sub(r"\uf02d", "", text)
    text = re.sub(r"\uf0fc", "", text)
    text = re.sub(r"\uf06e", "", text)
    text = re.sub(r"\uf07a", "", text)
    text = re.sub(r"\ufb01", "fi", text)
    text = re.sub(r"\ufb00", "ff", text)
    text = re.sub(r"\uf0d8", "", text)
    text = re.sub(r"\u00b7", "", text)
    text = re.sub("\t", " ", text)
    text = re.sub(" +", " ", text)
    return text.strip()


def bold_text(text):
    return "\033[1m" + text + "\033[0m"


def print_doc(doc):
    print(f"{bold_text('Page_Content:')} {doc.page_content}\n")
    print(f"{bold_text('Metadata:')} {doc.metadata}\n")


def create_pdf_df(pdf_folder):
    pdfs = get_files(pdf_folder, extensions=[".pdf"])
    return pd.DataFrame({"pdf_path": pdfs})


def create_idx_q_df(query, index_folder, index_name):
    index_names = [
        f.stem for f in Path(index_folder).glob("*.faiss") if f.stem.startswith(index_name)
    ]
    return pd.DataFrame(
        {"index_folder": str(index_folder), "index_name": index_names, "query": query}
    )


def doc_proc_text(doc):
    doc.page_content = process_text(doc.page_content)
    return doc


def df_pdf_docs_(
    df,
    path_col="pdf_path",
    splitter=None,
    chunk_size=200,
    chunk_overlap=60,
    separators=["\n\n", "\n", "(?<=\. )", " ", ""],
    add_start_index=True,
    proc=True,
):
    if splitter is None:
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=separators,
            add_start_index=add_start_index,
        )
    docs = PyPDFLoader(str(df[path_col])).load_and_split(text_splitter=splitter)
    if proc:
        docs = [doc_proc_text(doc) for doc in docs]
    df["docs"] = docs
    return df


def df_pdf_docs(
    df,
    path_col="pdf_path",
    docs_col="doc",
    splitter=None,
    chunk_size=200,
    chunk_overlap=60,
    separators=["\n\n", "\n", "(?<=\. )", " ", ""],
    add_start_index=True,
    proc=True,
    **kwargs,
):
    if splitter is None:
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=separators,
            add_start_index=add_start_index,
        )
    df = df.apply(df_pdf_docs_, axis=1, path_col=path_col, splitter=splitter, proc=proc)
    return pd.DataFrame({docs_col: flatten_list(df["docs"])})


def text_cat(
    text,
    cats_model,
    thresh=0.6,
    classes=["Work Experience", "Education", "Certifications", "Other"],
    other_class="Other",
):
    if other_class not in classes:
        classes.append(other_class)
    ot_id = classes.index(other_class)
    probs, ids = cats_model.predict_proba(text).detach().cpu().max(1)
    ids[probs < thresh] = ot_id
    pred_classes = np.array(classes)[ids]
    return pred_classes.tolist()


def docs_cat(docs, cats_model):
    cats = text_cat([doc.page_content for doc in docs], cats_model)
    for doc, cat in zip(docs, cats):
        doc.metadata["category"] = cat
    return docs


def df_docs_cat(df, cats_model, docs_col="doc"):
    df[docs_col] = docs_cat(df[docs_col], cats_model)
    return df


def docs_ems(docs, ems_model):
    return ems_model.encode([doc.page_content for doc in docs]).tolist()


def df_docs_ems(df, ems_model, ems_folder, docs_col="doc"):
    os.makedirs(ems_folder, exist_ok=True)
    ems = docs_ems(df[docs_col], ems_model)
    for doc, em in zip(df[docs_col], ems):
        source = Path(doc.metadata["source"])
        ems_path = (Path(ems_folder) / source.stem).with_suffix(".json")
        if ems_path.exists():
            ems_path = find_alternate_path(ems_path, first_idx=1, verbose=False)
        with open(ems_path, "w") as f:
            ems_dict = {"embedding": em, "metadata": doc.metadata}
            json.dump(ems_dict, f)
    return df


def df_to_faiss(
    df,
    ems_model,
    index_folder="/media/hamza/data2/faiss_data/saved_indexes/",
    index_name="index",
    docs_col="doc",
):
    db = FAISS.from_documents(df[docs_col], ems_model)
    index_path = (Path(index_folder) / index_name).with_suffix(".faiss")
    if index_path.exists():
        index_path = find_alternate_path(index_path, first_idx=1, verbose=False)
    index_name = index_path.stem
    db.save_local(index_folder, index_name)
    return df


def df_search_faiss(df, ems_model, filter=None, k=2):
    index_folder = str(df["index_folder"])
    index_name = df["index_name"]
    query = df["query"]
    db = FAISS.load_local(index_folder, embeddings=ems_model, index_name=index_name)
    q_sims = db.similarity_search_with_score(query, filter=filter, k=k)
    # print(f'\n\nQSIMS: {q_sims}\n\n')
    df["results"] = q_sims
    return df
