import pandas as pd
import datasets
import faiss
import numpy as np
import gradio as gr
import torch
from datasets import Dataset
from transformers import FeatureExtractionPipeline, pipeline

def load_encoder_pipeline(encoder_path: str) -> FeatureExtractionPipeline:
    """訓練済みの教師なしSimCSEのエンコーダを読み込む"""
    encoder_pipeline = pipeline("feature-extraction", model=encoder_path)
    return encoder_pipeline

def load_dataset(dataset_dir: str) -> Dataset:
    """文埋め込み適用済みのデータセットを読み込み、Faissのインデックスを構築"""
    # ディスクに保存されたデータセットを読み込む
    dataset = datasets.load_from_disk(dataset_dir)

    # データセットの"embeddings"フィールドの値からFaissのインデックスを構築する
    emb_dim = len(dataset[0]["embeddings"])
    index = faiss.IndexFlatIP(emb_dim)
    dataset.add_faiss_index("embeddings", custom_index=index)
    return dataset

def embed_text(
    text: str, encoder_pipeline: FeatureExtractionPipeline
) -> np.ndarray:
    """教師なしSimCSEのエンコーダを用いてテキストの埋め込みを計算"""
    with torch.inference_mode():
        # encoder_pipelineが返すTensorのsizeは(1, トークン数, 埋め込みの次元数)
        encoded_text = encoder_pipeline(text, return_tensors="pt")[0][0]

    # ベクトルをNumPyのarrayに変換
    emb = encoded_text.cpu().numpy().astype(np.float32)
    # ベクトルのノルムが1になるように正規化
    emb = emb / np.linalg.norm(emb)
    return emb

def search_similar_texts(
    query_text: str,
    dataset: Dataset,
    encoder_pipeline: FeatureExtractionPipeline,
    k: int = 5,
#) -> list[dict[str, float | str]]:
) -> list[dict[str, str]]:
    """モデルとデータセットを用いてクエリの類似文検索を実行"""
    # クエリに対して類似テキストをk件取得する
    scores, retrieved_examples = dataset.get_nearest_examples(
        "embeddings", embed_text(query_text, encoder_pipeline), k=k
    )
    titles = retrieved_examples["title"]
    texts = retrieved_examples["text"]

    # 検索された類似テキストをdictのlistにして返す
    results = [
        {"score": score, "title": title, "text": text}
        for score, title, text in zip(scores, titles, texts)
    ]
    return results 

# 訓練済みの教師なしSimCSEのモデルを読み込む
encoder_pipeline = load_encoder_pipeline("outputs_unsup_simcse/encoder")

# 文埋め込み適用済みのデータセットを読み込む
dataset = load_dataset("outputs_unsup_simcse/embedded_paragraphs")

def simsearch(query_text:str, k:int =10) -> pd.DataFrame:
    ret = search_similar_texts(query_text, dataset, encoder_pipeline, k)
    return  pd.DataFrame(ret)

with gr.Blocks() as app:
    gr.Markdown("# Wikipediaのパラグラフを用いた類似文検索")
    # クエリの入力欄を表示し、入力された値を受け取る
    query_text = gr.Textbox(label="入力文", value="肺胞では、膜と毛細血管の壁を通して、呼吸による二酸化炭素と酸素の交換(ガス交換)が行われています．" )
    # 検索結果数のスライダーを表示し、設定された値を受け取る
    k = gr.Slider(1, 50, step=1, value=10, label="検索結果数", info="検索結果の数を指定する")
    # 検索を実行するボタンを表示し、押下されたらTrueを受け取る
    search_btn = gr.Button("類似検索")
    output = gr.Dataframe(label="類似検索結果")
    # クエリに対して類似文検索を実行し、検索結果を受け取る
    search_btn.click(fn=simsearch, inputs=[query_text, k], outputs=output, api_name="greet")

app.launch()