from haystack.nodes.retriever import EmbeddingRetriever
from haystack.nodes import TableReader, FARMReader, RouteDocuments, JoinAnswers
from haystack import Pipeline

text_reader_types = {
    "minilm": "deepset/minilm-uncased-squad2", 
    "distilroberta": "deepset/tinyroberta-squad2", 
    "electra-base": "deepset/electra-base-squad2", 
    "bert-base": "deepset/bert-base-cased-squad2", 
    "deberta-large": "deepset/deberta-v3-large-squad2", 
    "gpt3": "implement openai answer generator"
}
table_reader_types = {
    "tapas": "deepset/tapas-large-nq-hn-reader",
    "text": "implement changing tables to text"
}


def create_retriever(document_store):
    retriever = EmbeddingRetriever(document_store=document_store, embedding_model="deepset/all-mpnet-base-v2-table")
    document_store.update_embeddings(retriever=retriever)
    return document_store, retriever

def create_readers_and_pipeline(retriever, text_reader_type = "deepset/roberta-base-squad2", table_reader_type="deepset/tapas-large-nq-hn-reader", use_table=True, use_text=True):
    both = (use_table and use_text)
    if use_text or both:
        print("Initializing Text reader..")
        text_reader = FARMReader(text_reader_type)
    if use_table or both:
        print("Initializing table reader..")
        table_reader = TableReader(table_reader_type)
    if both:
        route_documents = RouteDocuments()
        join_answers = JoinAnswers()
    
    text_table_qa_pipeline = Pipeline()
    text_table_qa_pipeline.add_node(component=retriever, name="EmbeddingRetriever", inputs=["Query"])
    if use_table and not use_text:
        text_table_qa_pipeline.add_node(component=table_reader, name="TableReader", inputs=["EmbeddingRetriever"])
    elif use_text and not use_table:
        text_table_qa_pipeline.add_node(component=text_reader, name="TextReader", inputs=["EmbeddingRetriever"])
    elif both:
        text_table_qa_pipeline.add_node(component=route_documents, name="RouteDocuments", inputs=["EmbeddingRetriever"])
        text_table_qa_pipeline.add_node(component=text_reader, name="TextReader", inputs=["RouteDocuments.output_1"])
        text_table_qa_pipeline.add_node(component=table_reader, name="TableReader", inputs=["RouteDocuments.output_2"])
        text_table_qa_pipeline.add_node(component=join_answers, name="JoinAnswers", inputs=["TextReader", "TableReader"])

    return text_table_qa_pipeline