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

__all__ = ['JSONizerWithBool', 'extractembeds']

# Cell
from copy import deepcopy
import json
import torch.nn as nn
import pandas as pd
from fastcore.dispatch import *
from .utils import *

# Cell
class JSONizerWithBool(json.JSONEncoder):
    def default(self, obj):
        return super().encode(bool(obj)) \
            if isinstance(obj, np.bool_) \
            else super().default(obj)

# Cell
def _catdict2embedsdictstruct(catdict):
    embedsdict = {}
    for cat, classes in catdict.items():
        embedsdict[cat] = {}
        embedsdict[cat]["classes"] = classes
    return embedsdict

# Cell
@typedispatch
def extractembeds(model: nn.Module, df: pd.DataFrame, *, transfercats, allcats, path=None, kind="bson"):
    catdict = getcatdict(df, transfercats)
    return extractembeds(model, catdict, transfercats=transfercats, allcats=allcats, path=path, kind=kind)


@typedispatch
def extractembeds(model: nn.Module, catdict: dict, *, transfercats, allcats, path=None, kind="bson"):
    '''
    Extracts embedding weights from `model`, which can be further transferred to other models.

    model: Any pytorch model, containing the embedding layers.
    catdict: A dictionary with category as key, and classes as value.
    transfercats: Names of categories to be transferred.
    allcats: Names of all categories corresponding to the embedding layers in model.
    path: Path for the json to be stored.
    '''
    embedsdict = _catdict2embedsdictstruct(catdict)
    model_dict = list(model.state_dict().items())
    for i, cat in enumerate(transfercats):
        classes = catdict[cat]
        catidx = allcats.index(cat)
        assert (model_dict[catidx][1].shape[0] == len(classes)), \
            (f"embeddings dimension {model_dict[catidx][1].shape[0]} !="
            f"num of classes {len(classes)} for vairable {cat}. Embeddings should have"
            f"same number of classes. Something might have gone wrong.")
        embedsdict[cat]["embeddings"] = model_dict[catidx][1].numpy().tolist()
    if (path != None):
        with open(path, 'w') as fp:
            json.dump(embedsdict, fp, cls = JSONizerWithBool) if kind == "json" else store_bson(path, embedsdict)
    return embedsdict

# Cell
#nbdev_comment _all_ = ['extractembeds']