import gradio as gr
import json

# Import your modules here
from Agents.togetherAIAgent import generate_article_from_query
from Agents.wikiAgent import get_wiki_data
from Agents.rankerAgent import rankerAgent
from Query_Modification.QueryModification import query_Modifier, getKeywords
from Ranking.RRF.RRF_implementation import reciprocal_rank_fusion_three, reciprocal_rank_fusion_six
from Retrieval.tf_idf import tf_idf_pipeline
from Retrieval.bm25 import bm25_pipeline
from Retrieval.vision import vision_pipeline
from Retrieval.openSource import open_source_pipeline
from Baseline.boolean import boolean_pipeline
from AnswerGeneration.getAnswer import generate_answer_withContext, generate_answer_zeroShot

# Load miniWikiCollection
miniWikiCollection = json.load(open('Datasets/mini_wiki_collection.json', 'r'))
miniWikiCollectionDict = {wiki['wikipedia_id']: " ".join(wiki['text']) for wiki in miniWikiCollection}

def process_query(query):
    # Query modification
    modified_query = query_Modifier(query)

    # Context Generation
    article = generate_article_from_query(query)

    # Keyword Extraction and getting context from Wiki
    keywords = getKeywords(query)
    wiki_data = get_wiki_data(keywords)

    # Retrieve rankings
    boolean_ranking = boolean_pipeline(query)
    tf_idf_ranking = tf_idf_pipeline(query)
    bm25_ranking = bm25_pipeline(query)
    vision_ranking = vision_pipeline(query)
    open_source_ranking = open_source_pipeline(query)

    # Modified queries
    boolean_ranking_modified = boolean_pipeline(modified_query)
    tf_idf_ranking_modified = tf_idf_pipeline(modified_query)
    bm25_ranking_modified = bm25_pipeline(modified_query)
    vision_ranking_modified = vision_pipeline(modified_query)
    open_source_ranking_modified = open_source_pipeline(modified_query)

    # RRF rankings
    tf_idf_bm25_open_RRF_Ranking = reciprocal_rank_fusion_three(tf_idf_ranking, bm25_ranking, open_source_ranking)
    tf_idf_bm25_open_RRF_Ranking_modified = reciprocal_rank_fusion_three(tf_idf_ranking_modified, bm25_ranking_modified, open_source_ranking_modified)
    tf_idf_bm25_open_RRF_Ranking_combined = reciprocal_rank_fusion_six(
        tf_idf_ranking, bm25_ranking, open_source_ranking,
        tf_idf_ranking_modified, bm25_ranking_modified, open_source_ranking_modified
    )

    # Retrieve contexts
    boolean_context = miniWikiCollectionDict[boolean_ranking[0]]
    tf_idf_context = miniWikiCollectionDict[tf_idf_ranking[0]]
    bm25_context = miniWikiCollectionDict[str(bm25_ranking[0])]
    vision_context = miniWikiCollectionDict[vision_ranking[0]]
    open_source_context = miniWikiCollectionDict[open_source_ranking[0]]

    tf_idf_bm25_open_RRF_Ranking_context = miniWikiCollectionDict[tf_idf_bm25_open_RRF_Ranking[0]]

    # Generating answers
    agent1_context = wiki_data[0]
    agent2_context = article

    agent1_answer = generate_answer_withContext(query, agent1_context)
    agent2_answer = generate_answer_withContext(query, agent2_context)
    boolean_answer = generate_answer_withContext(query, boolean_context)
    tf_idf_answer = generate_answer_withContext(query, tf_idf_context)
    bm25_answer = generate_answer_withContext(query, bm25_context)
    vision_answer = generate_answer_withContext(query, vision_context)
    open_source_answer = generate_answer_withContext(query, open_source_context)

    tf_idf_bm25_open_RRF_Ranking_answer = generate_answer_withContext(query, tf_idf_bm25_open_RRF_Ranking_context)

    zeroShot = generate_answer_zeroShot(query)

    # Ranking the best answer
    rankerAgentInput = {
        "query": query,
        "agent1": agent1_answer,
        "agent2": agent2_answer,
        "boolean": boolean_answer,
        "tf_idf": tf_idf_answer,
        "bm25": bm25_answer,
        "vision": vision_answer,
        "open_source": open_source_answer,
        "tf_idf_bm25_open_RRF_Ranking": tf_idf_bm25_open_RRF_Ranking_answer,
        "zeroShot": zeroShot,
    }

    best_model, best_answer = rankerAgent(rankerAgentInput)

    return best_model, best_answer

# Gradio interface
interface = gr.Interface(
    fn=process_query,
    inputs=gr.Textbox(label="Enter your query"),
    outputs=[
        gr.Textbox(label="Best Model"),
        gr.Textbox(label="Best Answer"),
    ],
    title="Query Answering System",
    description="Enter a query to get the best model and the best answer using multiple retrieval models and ranking techniques.",
    allow_flagging="never"
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()