Module sentspace.lexical.utils

import os
import pickle

import numpy as np
import sentspace.utils
from sentspace.utils import io, text
from sentspace.utils.caching import cache_to_mem  # , cache_to_disk
from sentspace.utils.misc import merge_lists
from pathlib import Path


# --------- Lexical features

# list of acceptable feature terms to load_databases(...)
# @cache_to_mem
def get_feature_list():
    return [
        "NRC_Arousal",  # (Mohammad 2018)
        "NRC_Valence",  # (Mohammad 2018)
        "NRC_Dominance",  # (Mohammad 2018)
        "OSC",  # Orthographic-semantics consistency (Marelli & Amenta, 2018)
        "aoa",  # Age of Acquisition (Kuperman et al., 2012)
        "concreteness",  # (Brysbaert et al., 2014)
        "lexical_decision_RT",  # (Balota et al., 2007)
        "log_contextual_diversity",
        "log_lexical_frequency",
        "n_orthographic_neighbors",
        "num_morpheme",
        "prevalence",
        "surprisal-1",
        "surprisal-2",
        "surprisal-3",
        "surprisal-4",
        "total_degree_centrality",
        "imageability",  # Glasgow norms (Scott et al, 2019) https://link.springer.com/article/10.3758/s13428-018-1099-3#Sec1
        "body-object-interaction",  # (Pexman et al, 2019)  https://link.springer.com/article/10.3758/s13428-018-1171-z#Sec9
        "zipf",  # SUBTLEXus
        "socialness",  # (Diveica et al., 2022) https://link.springer.com/article/10.3758/s13428-022-01810-x#Sec15
        # Lancaster norms (Lynott et al, 2020) https://link.springer.com/article/10.3758/s13428-019-01316-z#Bib1
        "Auditory",
        "Gustatory",
        "Interoceptive",
        "Haptic",
        "Olfactory",
        "Visual",
        "Foot_leg",
        "Hand_arm",
        "Head",
        "Mouth",
        "Torso",
    ]


def get_feature_list_using_third_party_libraries():
    return ["polysemy", "num_morpheme_poly"]


def get_feature_list_requiring_calculation():
    return ["PMI"]


def get_user_contributed_features():
    return []


def get_all_features(sentence: "sentspace.Sentence.Sentence", databases):
    """
    Given list of words, return dict mapping feature to list of feature values
    """

    result = {}
    for feature in (
        get_feature_list()
        + get_feature_list_using_third_party_libraries()
        + get_user_contributed_features()
    ):
        # we don't want to compute num_morpheme using the dictionary DB by default.
        # we want to do it only if the polyglot library is unavailable.
        if feature == "num_morpheme":
            continue
        computed_feature = get_feature(sentence, feature, databases)
        # even though we are computing "num_morpheme_poly", we want to
        # record it as "num_morpheme", since the _poly suffix comes from the polyglot library
        # that provides the feature
        if feature == "num_morpheme_polyglot":
            try:
                import polyglot
            except ImportError:
                feature = "num_morpheme"
        result[feature] = computed_feature
    return result


def get_feature(sentence: "sentspace.Sentence.Sentence", feature, databases={}):
    """
    get specific `feature` for the tokens in `sentence`; fall back to using `lemmas` if necessary
    """

    def get_feature_(token, lemma, feature):
        """given a `word` and a feature to extract, returns the value of that
            feature for the `word` using available databases

        Args:
            word (str): the token (word) to extract a feature for
            feature (str): name identifier of the feature acc to predefined convention
            databases (dict, in-scope): dictionary of feature --> (word --> feature_value) dictionaries.
                                        Defaults to {}.

        Returns:
            Any: feature value
        """
        # if the feature is from a database we have on disk
        # database can be a dictionary or an object that implements
        # get(key, default)
        if feature in get_feature_list():
            feature_dict = databases[feature]
            try:
                return feature_dict[token]
            except KeyError as e:
                try:
                    return feature_dict[lemma]
                except KeyError as e_:
                    return np.nan

        # Other databases we use from libraries we load such as NLTK-Wordnet and Polyglot
        elif feature in get_feature_list_using_third_party_libraries():
            if feature == "polysemy":
                from nltk.corpus import wordnet

                # first try it with the token itself
                if synsets := wordnet.synsets(token):
                    return len(synsets)  # TODO does a word's synset include itself?
                # if token is OOV, try again with the lemma
                elif synsets := wordnet.synsets(lemma):
                    return len(synsets)
                # otherwise the len of its synset is 1 (itself)
                return 1
            elif feature == "num_morpheme_poly":
                try:
                    from polyglot.text import Word

                    # try first to obtain # morphemes of the token
                    if morphed := Word(token, language="en").morphemes:
                        return len(morphed)
                    # otherwise, try using the lemmatized form
                    elif morphed := Word(lemma, language="en").morphemes:
                        return len(morphed)
                    # if both token and lemma OOV, then return nan? or 1 (i.e. full word is the morpheme?)
                    # but that only means we failed to analyze its morphology, not necessarily that is
                    # *is* a single morpheme
                    return 1  # np.nan
                except ImportError as e:
                    # fall back to simply using a dictionary-based feature
                    # TODO make a note of this somewhere
                    io.log(e.msg, type="WARN")
                    return get_feature_(token, lemma, "num_morpheme")

        else:
            raise ValueError(f"unable to compute unknown feature `{feature}`")

    features_list = []

    for token, lemma in zip(sentence.tokens, sentence.lemmas):
        features_list += [get_feature_(token, lemma, feature)]

    return features_list


def return_percentile_df(bench_df, usr_df):
    # Initialize df
    perc_df = pd.DataFrame(columns=usr_df.columns)
    # For each sentence get the percentile scores for each feature
    for index, row in usr_df.iterrows():
        temp_df = {}
        # Iterate through the features
        for col in usr_df.columns:
            if col == "Sentence no.":
                temp_df[col] = row[col]
                continue
            # print(percentileofscore(bench_df[col],row[col]))
            temp_df[col] = percentileofscore(bench_df[col], row[col])
            # pdb.set_trace()
        # Append percentile feature row
        perc_df = perc_df.append(temp_df, ignore_index=True)

    perc_df.drop(columns=["Sentence no."])
    return perc_df


@cache_to_mem
def load_databases(
    features="all",
    path="~/.cache/sentspace/",
    ignore_case=True,
):
    """
    Load dicts mapping word to feature value
    If one feature, provide in list format
    """
    path = str(Path(path).expanduser().resolve()) + "/"
    io.log("loading databases with all features")
    databases = {}
    if features == "all":
        features = get_feature_list()
    for feature in features:
        if not os.path.exists(path + feature + ".pkl"):
            sentspace.utils.s3.load_feature(key=feature + ".pkl", root_dir=path)
        with open(path + feature + ".pkl", "rb") as f:
            d = pickle.load(f)
            # if ignore_case:  # add lowercase version to feature database
            #     for key, val in d.copy().items():
            #         d[str(key).lower()] = val
            databases[feature] = d

    sanity_check_databases(databases)
    io.log("---done--- loading databases with all features")
    return databases


def sanity_check_databases(databases):
    """
    perform sanity checks upon loading various datasets
    to ensure validity of the loaded data
    """

    assert databases["NRC_Arousal"]["happy"] == 0.735
    assert databases["NRC_Valence"]["happy"] == 1
    assert databases["OSC"]["happy"] == 0.951549893181384
    assert abs(databases["aoa"]["a"] - 2.893384) < 1e-4
    assert databases["concreteness"]["roadsweeper"] == 4.85
    # assert abs(databases['imag']['abbey'] - 5.344) < 1e-4
    assert databases["total_degree_centrality"]["a"] == 30
    assert databases["lexical_decision_RT"]["a"] == 798.917
    assert abs(databases["log_contextual_diversity"]["a"] - 3.9234) < 1e-4
    assert abs(databases["log_lexical_frequency"]["a"] - 6.0175) < 1e-4
    assert databases["n_orthographic_neighbors"]["a"] == 950.59
    assert databases["num_morpheme"]["abbreviated"] == 4
    assert abs(databases["prevalence"]["a"] - 1.917) < 1e-3
    assert databases["surprisal-3"]["beekeeping"] == 10.258