# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/pdf/02_df.ipynb.

# %% auto 0
__all__ = ['create_pdf_df', 'create_ner_df', 'extract_df_text', 'write_df_ner', 'df_segs', 'write_df_segs', 'write_df_ems']

# %% ../../nbs/pdf/02_df.ipynb 2
from ..imports import *
from ..utils import *
from .core import *
from .extract import *
from .ner import *


# %% ../../nbs/pdf/02_df.ipynb 4
def create_pdf_df(data_path, segs_folder, ems_folder):
    files = [str(x) for x in resolve_data_path(data_path) if Path(x).suffix == ".pdf"]
    df = pd.DataFrame({"path": files})
    df["segs_folder"] = str(segs_folder)
    df["ems_folder"] = str(ems_folder)
    return df


def create_ner_df(data_path, segs_folder):
    data_path = Path(data_path)
    if data_path.is_file():
        files = [str(data_path)]
    else:
        files = get_files(data_path, extensions=[".pdf"], make_str=True)
    df = pd.DataFrame({"path": files})
    df["segs_folder"] = str(segs_folder)
    return df


def extract_df_text(df, col="path", n_lines=3):
    df["text"] = extract_text_list(df[col], n_lines=n_lines)
    return df


def write_df_ner(df, tner, jner, work_ner_dict={"company": "", "date": ""}):
    segs_path = json_file(df["path"], df["segs_folder"])
    if os.path.exists(segs_path):
        with open(segs_path, "r") as f:
            segs = json.load(f)
        if not is_list(segs["Work Experience"][0]):
            segs = {k: [v] for k, v in segs.items()}
        work_ners = [
            work_ner(txt, tner, jner, ner_dict=work_ner_dict)
            for txt in segs["Work Experience"][0]
        ]
        work_ners = [ner for ner in work_ners if len(ner) > 0]
        segs["Work Experience"].append(work_ners)
        edu_ners = [edu_ner(txt, tner) for txt in segs["Education"][0]]
        edu_ners = [ner for ner in edu_ners if len(ner) > 0]
        segs["Education"].append(edu_ners)
        ner_path = extend_path_name(segs_path, "_ner")
        with open(ner_path, "w") as f:
            json.dump(segs, f, indent=4)
    return df


def df_segs(
    df,
    segs_model,
    thresh=0.6,
    classes=["Work Experience", "Education", "Certifications", "Other"],
    other_class="Other",
    write=False,
):
    if len(df["text"]) == 0:
        return {}, [], np.array([])
    segs, preds, probs = text_to_segments(
        df["text"], segs_model, thresh, classes=classes, other_class=other_class
    )
    if write:
        segs_path = json_file(df["path"], df["segs_folder"])
        with open(segs_path, "w") as f:
            json.dump(segs, f, indent=4)
    df["segs"] = segs
    df["preds"] = preds
    df["probs"] = probs
    return df


def write_df_segs(
    df,
    segs_model,
    thresh=0.6,
    write=False,
    classes=["Work Experience", "Education", "Certifications", "Other"],
    other_class="Other",
):
    if len(df["text"]) == 0:
        return {}, [], np.array([])
    segs, preds, probs = text_to_segments(
        df["text"], segs_model, thresh, classes=classes, other_class=other_class
    )
    if write:
        segs_path = json_file(df["path"], df["segs_folder"])
        with open(segs_path, "w") as f:
            json.dump(segs, f, indent=4)
    return segs, preds, probs


def write_df_ems(
    df,
    ems_model,
    classes=["Work Experience", "Education", "Certifications", "Other"],
    other_class="Other",
):
    classes = add_other_class(classes, other_class)
    seg_ids = {c: i + 1 for i, c in enumerate(classes)}
    fn = Path(df["path"]).stem
    ems_folder = Path(df["ems_folder"])
    os.makedirs(ems_folder, exist_ok=True)
    ems_path = str(ems_folder / fn)
    txt = df["text"]
    pred_classes = df.get("classes", [""] * len(txt))
    if len(txt) == 0 or len(pred_classes) == 0:
        return df
    ems = ems_model.encode(txt).tolist()
    for i, data in enumerate(zip(pred_classes, ems)):
        pc, em = data
        seg_id = seg_ids.get(pc, 0)
        if seg_id == 0:
            seg = ""
        else:
            seg = f"_{seg_id}"
        jf = ems_path + f"{seg}_{i+1}.json"
        em_dict = {"id": fn, "embedding": em}
        with open(jf, "w") as f:
            json.dump(em_dict, f)
    return df
