# -*- coding: utf-8 -*-
from gensim import corpora
from gensim.models.keyedvectors import KeyedVectors

from src.preprocessings.cleaning import clean
from src.preprocessings.normalization import normalize
from src.preprocessings.tokenizer import MeCabTokenizer


class WordMover(object):
    def __init__(self, vec_path, vocab_path, dic_path):
        self._word_vectors = KeyedVectors.load_word2vec_format(vec_path, binary=False)
        self._dic = corpora.Dictionary.load(vocab_path)
        self._tokenizer = MeCabTokenizer(user_dic_path=dic_path)

    def get_distance(self, doc1, doc2):
        document1 = self._preprocess(doc1)
        document2 = self._preprocess(doc2)
        distance = self._word_vectors.wmdistance(document1, document2)
        return distance

    def _preprocess(self, sent):
        text = clean(sent)
        tokenized_words = self._tokenizer.filter_by_pos(text, pos=('名詞', '動詞'))
        normalized_words = [normalize(word) for word in tokenized_words]
        document_ids = self._words2ids(normalized_words)
        document_ids = [str(id) for id in document_ids]
        return document_ids

    def _words2ids(self, doc):
        ids = [self._dic.token2id[word] for word in doc if word in self._dic.token2id]
        return ids

# -*- coding: utf-8 -*-
from gensim import corpora
from gensim.models.keyedvectors import KeyedVectors

from src.preprocessings.cleaning import clean
from src.preprocessings.normalization import normalize
from src.preprocessings.tokenizer import MeCabTokenizer


class WordMover(object):
    def __init__(self, vec_path, vocab_path, dic_path):
        self._word_vectors = KeyedVectors.load_word2vec_format(vec_path, binary=False)
        self._dic = corpora.Dictionary.load(vocab_path)
        self._tokenizer = MeCabTokenizer(user_dic_path=dic_path)

    def get_distance(self, doc1, doc2):
        document1 = self._preprocess(doc1)
        document2 = self._preprocess(doc2)
        distance = self._word_vectors.wmdistance(document1, document2)
        return distance

    def _preprocess(self, sent):
        text = clean(sent)
        tokenized_words = self._tokenizer.filter_by_pos(text, pos=('名詞', '動詞'))
        normalized_words = [normalize(word) for word in tokenized_words]
        document_ids = self._words2ids(normalized_words)
        document_ids = [str(id) for id in document_ids]
        return document_ids

    def _words2ids(self, doc):
        ids = [self._dic.token2id[word] for word in doc if word in self._dic.token2id]
        return ids

# -*- coding: utf-8 -*-