from transformers import pipeline
import pandas as pd

def TAPAS(question, table_main):
    """
    Processing the question using an expression and the main and geom table.

    Args:
        question (str): the question.
        table_main (df): main table
        table_geom (df): geom table
    Returns:
        answer (str): answer to the question
    """
    
    # set up a TAPAS pipeline for table-based question answering
    tqa = pipeline(task="table-question-answering", model="google/tapas-large-finetuned-wtq")

    # use the tqa pipeline to perform table-based question answering.
    i = tqa(table=table_main, query=question)['cells'][0]

    # Check if the output is the link to the TEMP DB:
    # Has to be done because the entrys for geometry, ... are an array :(
    if ';' in i:
        i = i.split(";")
        path = i[0]
        r = int(i[1])
        c = int(i[2])
        answer_table = pd.read_csv(path)
        answer = answer_table.iloc[r,c]
        return(answer)
    
    answer = str(i)
    return(answer)