# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00_utils.ipynb (unless otherwise specified).

__all__ = ['getcatdict', 'store_bson', 'load_bson', 'generate_files_embedprojector_for_comparision',
           '_load_extractembeds_from_dir', '_check_embdict_validity']

# Cell
import bson
import json
from pathlib import Path

# Cell
def getcatdict(df, catcols=None, add_na=False):
    if catcols == None:
        columns = list(df.columns)
        catcols = list(map(lambda arg: columns[arg[0]], filter(lambda arg: arg[1] == 'object', enumerate(df.dtypes))))
    catdict = {}
    for cat in catcols:
        catdict[cat] = ["nan"] + list(df[cat].unique()) if add_na else list(df[cat].unique())
    return catdict

# Cell
def store_bson(path, data):
    bdata = bson.dumps(data)
    with open(path, "wb") as fp:
        fp.write(bdata)

def load_bson(path):
    with open(path, "rb") as fp:
        bdata = fp.read()
    return bson.loads(bdata)

# Cell
def _load_extractembeds_from_dir(path):
    path = Path(path)
    json_file = {}
    if path.is_dir():
        for json_f in path.glob("*.json"):
            f = open(json_f, 'rb')
            file = json.load(f)
            for i in file.keys():
                json_file[i] = file[i]
        for bson_f in path.glob("*.bson"):
            file = load_bson(bson_f)
            for i in file.keys():
                json_file[i] = file[i]
    return json_file

# Cell
def _check_embdict_validity(embdict):
    for catcol, info in embdict.items():
        if set(list(info.keys())) == set(('classes', 'embeddings')) :
            class_len = len(info['classes'])
            embedding_len =  len(info['embeddings'])
            assert class_len == embedding_len, f"Class Size of {class_len} and Embedding size of {embedding_len} does not match for {catcol}"

# Cell
def generate_files_embedprojector_for_comparision(model1_path, model2_path, dir_path):
    model1_embedsdict = load_bson(model1_path)
    model2_embedsdict = load_bson(model2_path)
    for ((catcol, info1), (_, info2)) in zip(model1_embedsdict.items(), model2_embedsdict.items()):
        model1_data = [[clas + '_Model1'] for clas in info1['classes']]
        model2_data = [[clas + '_Model2'] for clas in info2['classes']]
        with open(dir_path + '/' + catcol + '_meta.tsv', 'w') as tsvfile:
            writer = csv.writer(tsvfile, delimiter='\n')
            writer.writerows(model1_data)
            writer.writerows(model2_data)
        with open(dir_path + '/' + catcol + '_embeds.tsv', 'w') as tsvfile:
            writer = csv.writer(tsvfile, delimiter='\t')
            writer.writerows(info1['embeddings'])
            writer.writerows(info2['embeddings'])

# Cell
#nbdev_comment _all_ = ['_load_extractembeds_from_dir', '_check_embdict_validity']