import gradio as gr
from langchain_community.vectorstores.faiss import FAISS
from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings

def load_index():
    index_path = "./storage"
    embedding_model = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large")
    return FAISS.load_local(
        folder_path=index_path, 
        embeddings=embedding_model,
        allow_dangerous_deserialization=True
    )

default_k = 4
default_fetch_k = 500
default_threshold = 0.5

def search_documents(input_text, courseOfStudy, index, k = default_k, fetch_k = default_fetch_k, score_threshold = default_threshold):
    if not input_text:
        return [[], []]  # Return empty outputs for both components
    
    metadata = {}
    if courseOfStudy:
        metadata["学校種別"] = courseOfStudy
    
    try:
        docs_and_scores = index.similarity_search_with_score(
            input_text, 
            filter=metadata,
            k=k,
            fetch_k=fetch_k,
            score_threshold=score_threshold
            )
    except Exception as e:
        print(f"Error during search: {e}")
        return [[], []]  # Return empty outputs for both components if an error occurs
    
    rows = [
        [
            f"<a href='https://w3id.org/jp-cos/{doc.metadata['id']}' target='_blank'>{doc.metadata['id']}</a>",      
            doc.metadata["学校種別"],
            doc.metadata["教科等"],
            round(float(score), 3),
            doc.page_content,
        ]
        for doc, score in docs_and_scores
    ]

    json_data = [
        {
            "dcterms:identifier": doc.metadata['id'],
            "jp-cos:courseOfStudy": doc.metadata["学校種別"],
            "jp-cos:subjectArea": doc.metadata["教科等"],
            "score": round(float(score), 3),
            "jp-cos:sectionText": doc.page_content,
        }
        for doc, score in docs_and_scores
    ]

    return [
        rows,
        json_data
    ]

def setup_gradio_interface():
    text_input = gr.Textbox(lines=2, placeholder="テキストを入力してください。", label="テキスト")
    schools = ['小学校', '特別支援学校小学部', '特別支援学校中学部', '幼稚園', '特別支援学校幼稚部', '中学校', '高等学校']
    metadata_selector = gr.CheckboxGroup(schools, label="学校種別")
    output_table = gr.Dataframe(
        headers=["学習指導要領コード",  "学校種別", "教科等", "スコア", "テキスト"], 
        datatype=["html", "str", "str", "number", "str"],
        label="検索結果"
    )

    k = gr.Number(value=default_k, label="検索結果の数", minimum=1, maximum=200)
    fetch_k = gr.Number(value=default_fetch_k, label="検索対象の数", minimum=1, maximum=2000)
    score_threshold = gr.Slider(value=default_threshold, label="スコアの閾値", minimum=0.0, maximum=1.0, step=0.01)
    
    index = load_index()

    examples=[
        ["空気鉄砲や水鉄砲、ペットボトルロケットなどのしくみを調べ、空気はおし縮められ体積が小さくなるにつれて反発する力が大きくなるが、水はおし縮められないことに気づく。", "小学校", 4, 500, 0.5],
        ["頼朝の妻・北条政子は、頼朝が築いた「ご恩と奉公」の絆で御家人たちを団結させ、復権をねらう朝廷に打ち勝った。鎌倉幕府の力をさらに強固なものにした政子の功績を知る。", "小学校", 4, 500, 0.5],
        ["胞子をつくるキノコ。胞子が菌糸を伸ばし新しいキノコを作る", "中学校", 4, 500, 0.5],
        ["今回のテーマは「鋭角の三角比の定義」。三角比は、高校で初めて扱う分野です。サイン、コサイン、タンジェント・・・これから何回かにわたり、この言葉を使って図形などの問題を解いていきますので、定義をばっちり覚えておきましょう。ミニコーナーでは、物語を書いたことで有名な、ある数学者を紹介します。\n出演：アイク　ぬわら、酒井　蒼澄、藤本　ばんび、湯浅　弘一", "高等学校", 4, 500, 0.5],
    ]
    
    interface = gr.Interface(
        fn=lambda text, courseOfStudy, k, fetch_k, score_threshold: search_documents(text, courseOfStudy, index, k=k, fetch_k=fetch_k, score_threshold=score_threshold),
        inputs=[text_input, metadata_selector, k, fetch_k, score_threshold], 
        outputs=[output_table, gr.JSON(label="JSON")], 
        title="Japanese Course of Study Predictor",
        description="入力したテキストに類似するテキストを持つ学習指導要領コードを検索します。",
        article=(
            "<p style='text-align: center'><a href='https://jp-cos.github.io/' target='_blank'>"
            "学習指導要領LOD</a> | <a href='https://www.nhk.or.jp/school/' target='_blank'>"
            "NHK for School</a></p>"
        ),
        examples=examples,
        allow_flagging="never"
    )
    return interface

demo = setup_gradio_interface()
demo.launch()