import gensim
from sklearn.metrics.pairwise import cosine_similarity


class CosineSimilarityCalculator:
    model_path = 'ja/ja.bin'

    def __init__(self):
        """
        CosineSimilarityCalculatorクラスを初期化し、
        事前トレーニング済みのWord2Vecモデルをロード

        Parameters:
        - なし

        Returns:
        - なし
        """
        self.model = gensim.models.Word2Vec.load(CosineSimilarityCalculator.model_path)

    def _convert_to_2d_array(self, vector):
        """
        埋め込みベクトルを2次元配列に変換

        Parameters:
        - vector (numpy.ndarray): 変換する1次元配列のベクトル

        Returns:
        - vector_2d (numpy.ndarray): 変換後の2次元配列のベクトル
        """
        return vector.reshape(1, -1)

    def _calculate_cosine_similarity(self, embedding1, embedding2):
        """
        ｃos類似度を計算

        Parameters:
        - embedding1 (numpy.ndarray): 1つ目の単語ベクトル（2次元配列）
        - embedding2 (numpy.ndarray): 2つ目の単語ベクトル（2次元配列）

        Returns:
        - similarity (numpy.ndarray): cos類似度
        """
        return cosine_similarity(embedding1, embedding2)

    def calculate_similarity(self, search_word, article_keyword_list):
        """
        指定された検索ワードと記事のキーワードリストの間のcos類似度を計算
        モデルにない単語の場合はエラーメッセージを出力しブレイクする

        Parameters:
        - search_word (str): 検索ワード
        - article_keyword_list (list): 記事のキーワードリスト
        Returns:
        - similarities (dict): 記事キーワードとそれぞれの検索ワードのcos類似度を含むdictを作成
                               モデルにない単語の場合はNoneを返す
        """
        # 検索ワードの埋め込みベクトルを取得
        if search_word in self.model.wv:
            search_embedding = self.model.wv[search_word]
        else:
            return None

        similarities = {}
        # 記事キーワードの埋め込みベクトルを取得し、cos類似度を計算
        for keyword in article_keyword_list:
            if keyword in self.model.wv:
                keyword_embedding = self.model.wv[keyword]
                search_embedding_2d = self._convert_to_2d_array(search_embedding)
                keyword_embedding_2d = self._convert_to_2d_array(keyword_embedding)
                similarity = self._calculate_cosine_similarity(search_embedding_2d, keyword_embedding_2d)
                similarities[keyword] = similarity[0][0]
            else:
                similarities[keyword] = None
        return similarities