# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks_build/03_spe2vec.ipynb (unless otherwise specified).

__all__ = ['Corpus', 'learn_spe2vec', 'load_spe2vec', 'SPE2Vec']

# Cell
import os
import gensim


class Corpus(object):
    '''
    *filename*: A file that stores SMILES line-by-line.

    *tokenizer*: SPE tokenizer

    *dropout*: SPE dropout, default = 0
    '''
    def __init__(self, infile, tokenizer, isdir=False, dropout=0):
        self.infile = infile
        self.tokenizer = tokenizer
        self.dropout = dropout
        self.isdir = isdir

    def __iter__(self):
        if self.isdir:
            for fname in os.listdir(self.infile):
                for smi in open(os.path.join(self.infile, fname)):
                    yield self.tokenizer.tokenize(smi, dropout=self.dropout).split(' ')
        else:
            for smi in open(self.infile):
                yield self.tokenizer.tokenize(smi, dropout=self.dropout).split(' ')

def learn_spe2vec(corpus, outfile=None,
                  vector_size=100, window=10, min_count=10, n_jobs = 1, method = 'skip-gram',
                  **kwargs):
    '''
    Train a spe2vec model.

    *corpus*: an instance of `Class Corpus()`

    *outfile*: str, name of the spe2vec model file.

    *vector_size*: dimensions of embedding.

    *window*: number of tokens considered as context

    *min_count*: number of occurrences a token should have to be considered in training

    *n_jobs*: number of cpu cores used for training.

    *method*: modeling method, choose from ['cbow', 'skip-gram']

    More training parameter can be found https://radimrehurek.com/gensim/models/word2vec.html#gensim.models.word2vec.Word2Vec
    '''

    if method.lower() == 'skip-gram':
        sg = 1
    elif method.lower() == 'cbow':
        sg = 0
    else:
        raise ValueError("Invalid option,  choose from ['cbow', 'skip-gram']")

    model = gensim.models.Word2Vec(corpus, size=vector_size, window=window, min_count=min_count, workers=n_jobs, sg=sg,
                              **kwargs)

    if outfile:
        model.save(outfile)

    return model

def load_spe2vec(model_path):
    return gensim.models.Word2Vec.load(model_path)

# Cell

import numpy as np

class SPE2Vec(object):
    def __init__(self, model_path, tokenizer):
        self.model = gensim.models.Word2Vec.load(model_path)
        self.tokenizer = tokenizer
        self.token_keys = set(self.model.wv.vocab.keys())

        #get the vector for unknown tokens. simply averge the vectors of all known tokens.

        vectors = []
        for word in self.model.wv.vocab:
            vectors.append(self.model.wv[word])
        self.unknown = np.mean(vectors, axis=0)

    def tokenize(self, smi, dropout=0):
        '''
        tokenize SMILES into substructure tokens.
        '''
        return self.tokenizer.tokenize(smi, dropout)

    def smiles2vec(self, smi, dropout=0, mode = 'average'):
        '''
        Generate a vector for a SMILES. The vector is construc in four modes: ['average', 'sum', 'avg_pool', 'sum_pool']

        `average`: average the embedding of all tokens

        `sum`: sum the embedding of all tokens

        `avg_pool`: concatenation of average, max pooling and min pooling

        `sum_pool`: concatenation of sum, max pooling and min pooling

        The Unknown token will be skipped
        '''

        if mode not in ['average', 'sum', 'avg_pool', 'sum_pool']:
            raise ValueError("Invalid option,  choose from ['average', 'sum', 'avg_pool', 'sum_pool']")

        tokens = self.tokenizer.tokenize(smi, dropout).split(' ')

        if mode == 'average':

            return np.mean([self.model.wv[tok] for tok in tokens if tok in self.token_keys], axis=0)

        if mode == 'sum':
            return np.sum([self.model.wv[tok] for tok in tokens if tok in self.token_keys], axis=0)

        if mode == 'avg_pool':
            tok_mean = np.mean([self.model.wv[tok] for tok in tokens if tok in self.token_keys], axis=0)
            tok_max = np.amax([self.model.wv[tok] for tok in tokens if tok in self.token_keys], axis=0)
            tok_min = np.amin([self.model.wv[tok] for tok in tokens if tok in self.token_keys], axis=0)
            tok_concate = np.concatenate((tok_mean, tok_max, tok_min))
            return tok_concate

        if mode == 'sum_pool':
            tok_sum = np.sum([self.model.wv[tok] for tok in tokens if tok in self.token_keys], axis=0)
            tok_max = np.amax([self.model.wv[tok] for tok in tokens if tok in self.token_keys], axis=0)
            tok_min = np.amin([self.model.wv[tok] for tok in tokens if tok in self.token_keys], axis=0)
            tok_concate = np.concatenate((tok_sum, tok_max, tok_min))
            return tok_concate

    def spe2vec(self, smi, dropout=0, skip_unknown=False):
        '''
        Generate a list of vectors (np.array). Each vector is spe vector of each token.

        The unknown token will be represented by the mean of all token vectors from the model.
        '''

        token_keys = set(self.model.wv.vocab.keys())
        tokens = self.tokenizer.tokenize(smi, dropout).split(' ')

        if skip_unknown:
            vec = [self.model.wv[tok] for tok in tokens if tok in self.token_keys]
        else:
            vec = [self.model.wv[tok] if tok in self.token_keys else self.unknown for tok in tokens]

        return vec