import gradio as gr
from Levenshtein import ratio
import json

data_path = "./data.json"

with open(data_path, "r") as f:
    documents_data = json.load(f)

def search_similar_texts(query, selected_vols, top_n=5):
    results = []
    
    for doc in documents_data:
        # 選択された巻のみを検索対象とする
        if not selected_vols or str(doc["vol"]) in selected_vols:
            score = ratio(query, doc["text"])
            results.append({
                "vol": doc["vol"],
                "page": doc["page"],
                "score": score,
                "text": doc["text"]
            })

    results.sort(key=lambda x: x["score"], reverse=True)
    top_results = results[:top_n]  # top_nで指定された件数だけを取得

    '''
    # top_n件の結果における巻ごとの件数をカウント
    vol_counts = {}
    total_results = len(top_results)
    
    for doc in top_results:
        vol_str = f"巻{doc['vol']}"
        vol_counts[vol_str] = vol_counts.get(vol_str, 0) + 1

    # 巻ごとの割合を計算
    vol_percentages = {
        vol: f"{(count/total_results*100):.1f}%" 
        for vol, count in vol_counts.items()
    } if total_results > 0 else {}

    return [top_results, vol_percentages]  # リストとして返す
    '''

    return [top_results] # , vol_percentages

# Gradioインターフェースの作成
demo = gr.Interface(
    fn=search_similar_texts,
    inputs=[
        gr.Textbox(label="検索テキスト", placeholder="検索したいテキストを入力してください"),
        gr.Dropdown(
            choices=[str(i) for i in range(1, 55)],
            label="巻",
            multiselect=True,
            value=[],
        ),
        gr.Slider(minimum=1, maximum=10, value=5, step=1, label="表示件数")
    ],
    outputs=[
        gr.JSON(),
        # gr.JSON(label="巻ごとの割合")
    ],
    title="校異源氏物語 類似テキスト検索",
    description="テキストを入力すると、校異源氏物語の類似する箇所を検索します。",
    allow_flagging="never",
)

# インターフェースの起動
demo.launch()