# 必要なライブラリのインポート
import numpy as np
import pandas as pd
import wikipedia
from sentence_transformers import SentenceTransformer
import faiss

# Wikipedia APIの初期化
wiki_wiki = wikipedia.set_lang("jp")
# SentenceTransformerモデルの初期化
model = SentenceTransformer('cl-tohoku/bert-base-japanese')

# CSVファイルからキーワードを読み込む
df = pd.read_csv('keyword_data.csv')
df.head()

# Wikipediaの文章を取得し、ベクトルデータに変換してリストに保存
document_vectors = []
for keyword in df['keyword']:
    try:
      summary = wikipedia.summary(keyword)
      vector = model.encode(summary)
      document_vectors.append(vector)
    except:
      print(keyword)

# ベクトルデータをnumpy arrayに変換
vectors_np = np.array(document_vectors).astype('float32')

# FAISSインデックスを作成
dimension = len(vectors_np[0])
index = faiss.IndexFlatL2(dimension)
index.add(vectors_np)
# 実施の運用では下記のように作成したインデックスをファイルに保存しておくことが可能
# faiss.write_index(index, 'faiss_距離による類似度.index')

# 検索
search_word = 'すき焼き'
search_summary = wikipedia.summary(search_word)
search_vector = np.array([model.encode(search_summary)]).astype('float32')
D, I = index.search(search_vector, 50)  # 上位50件を表示

# FAISSインデックスを作成(コサイン類似度)
# 10個のクラスタリングに分けて、３個のクラスタ分を検索対象とする
dimension = len(vectors_np[0])
index_cos = faiss.IndexIVFFlat(quantizer, dimension, 10, faiss.METRIC_INNER_PRODUCT)
index_cos.nprobe = 3

# クラスタリングのためのトレーニング
train_vectors = vectors_np[:int(vectors_np.shape[0]/2)].copy()
faiss.normalize_L2(train_vectors)
index_cos.train(train_vectors)

# 全データのインデックスを作成
all_vectors = vectors_np.copy()
faiss.normalize_L2(all_vectors)
index_cos.set_direct_map_type(faiss.DirectMap.Hashtable)
index_cos.add_with_ids(all_vectors, df['no'])
# faiss.write_index(index_cos, 'faiss_コサイン類似度.index')

# 検索
search_word = 'すき焼き'
search_summary = wikipedia.summary(search_word)
search_vector = np.array([model.encode(search_summary)]).astype('float32')
faiss.normalize_L2(search_vector)
D, I = index_cos.search(search_vector, 50)

# 類似度が高い順にキーワードを表示
for i, keyword_index in enumerate(I[0]):
    similarity_score = D[0][i]  # 類似度スコア
    similar_keyword = df.loc[keyword_index, 'keyword']
    print(similar_keyword + ' score:' + str(similarity_score))