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

__all__ = ['Prod2Vec']

# Cell
import logging
import gensim
import numpy as np

# Cell
class Prod2Vec(object):
    """
    Implementation of the Prod2Vec skipgram model from
    Grbovic Mihajlo, Vladan Radosavljevic, Nemanja Djuric, Narayan Bhamidipati, Jaikit Savla, Varun Bhagwan, and Doug Sharp.
    "E-commerce in your inbox: Product recommendations at scale."
    In Proceedings of the 21th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining,
    pp. 1809-1818. ACM, 2015.
    """

    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
    logger = logging.getLogger()

    def __init__(self, min_count=2, negative=5, size=100, window=5, decay_alpha=0.9):
        """
        :param min_count: (optional) the minimum item frequency. Items less frequent that min_count will be pruned
        :param negative: (optional) the minimum negative samples
        :param size: (optional) the size of the embeddings
        :param window: (optional) the size of the context window
        :param decay_alpha: (optional) the exponential decay factor used to discount the similarity scores for items
                back in the user profile. Lower values mean higher discounting of past user interactions. Allows values in [0-1].
        """
        super(Prod2Vec, self).__init__()
        self.min_count = min_count
        self.negative = negative
        self.size = size
        self.window = window
        self.decay_alpha = decay_alpha

    def __str__(self):
        return 'Prod2Vec(min_count={min_count}, ' \
               'size={size}, ' \
               'window={window}, ' \
               'decay_alpha={decay_alpha})'.format(**self.__dict__)

    def fit(self, train_data):
        self.model = gensim.models.Word2Vec(train_data,
                                            min_count=self.min_count,
                                            negative=self.negative,
                                            window=self.window,
                                            hs=1,
                                            size=self.size,
                                            sg=1,
                                            workers=-1)
        self.model.train(train_data, total_examples = self.model.corpus_count,
                         epochs=10, report_delay=1)
        # As we do not plan to train the model any further, we are calling
        # init_sims(), which will make the model much more memory-efficient
        self.model.init_sims(replace=True)

    def aggregate_vectors(self, products):
        product_vec = []
        for i in products:
            try:
                product_vec.append(self.model[i])
            except KeyError:
                continue

        return np.mean(product_vec, axis=0)

    def recommend(self, user_profile, topk=5):
        """
        Given the user profile return a list of recommendation

        Args:
            user_profile: list of item ids visited/interacted by the user
            topk: (optional) top-k recommendations
        """
        rec = []
        try:
            vec = self.aggregate_vectors(user_profile)
            # extract most similar products for the input vector
            rec = self.model.wv.similar_by_vector(vec, topn= topk+1)[1:]
        except KeyError:
            rec = []

        return rec