import streamlit as st
import pandas as pd
import torch
from transformers import pipeline
#from transformers import TapasTokenizer, TapexTokenizer, BartForConditionalGeneration
from transformers import AutoTokenizer, AutoModelForTableQuestionAnswering
import datetime

#df = pd.read_excel('discrepantes.xlsx', index_col='Unnamed: 0')
df = pd.read_excel('discrepantes.xlsx')
df.fillna(0, inplace=True)
table_data = df.astype(str)
print(table_data.head())

def response(user_question, table_data):
    a = datetime.datetime.now()

    model_name = "google/tapas-base-finetuned-wtq"
    model = AutoModelForTableQuestionAnswering.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # The query should be passed as a list
    encoding = tokenizer(table=table_data, queries=[user_question], padding=True, return_tensors="pt", truncation=True)

    # Instead of using generate, we pass the encoding through the model to get the logits
    outputs = model(**encoding)

    # Extract the answer coordinates
    predicted_answer_coordinates = outputs.logits.argmax(-1)

    # Decode the answer from the table using the coordinates
    answer = tokenizer.convert_logits_to_predictions(
        encoding.data,
        predicted_answer_coordinates
    )

    # Process the answer into a readable format
    answer_text = answer[0][0][0] if len(answer[0]) > 0 else "Não foi possível encontrar uma resposta"

    query_result = {
        "Resposta": answer_text
    }

    b = datetime.datetime.now()
    print(b - a)

    return query_result

# Streamlit interface
st.markdown("""
<div style='display: flex; align-items: center;'>
    <div style='width: 40px; height: 40px; background-color: green; border-radius: 50%; margin-right: 5px;'></div>
    <div style='width: 40px; height: 40px; background-color: red; border-radius: 50%; margin-right: 5px;'></div>
    <div style='width: 40px; height: 40px; background-color: yellow; border-radius: 50%; margin-right: 5px;'></div>
    <span style='font-size: 40px; font-weight: bold;'>Chatbot do Tesouro RS</span>
</div>
""", unsafe_allow_html=True)

# Chat history
if 'history' not in st.session_state:
    st.session_state['history'] = []

# Input box for user question
user_question = st.text_input("Escreva sua questão aqui:", "")

if user_question:
    # Add person emoji when typing question
    st.session_state['history'].append(('👤', user_question))
    st.markdown(f"**👤 {user_question}**")
    
    # Generate the response
    bot_response = response(user_question, table_data)
    
    # Add robot emoji when generating response and align to the right
    st.session_state['history'].append(('🤖', bot_response))
    st.markdown(f"<div style='text-align: right'>**🤖 {bot_response}**</div>", unsafe_allow_html=True)

# Clear history button
if st.button("Limpar"):
    st.session_state['history'] = []

# Display chat history
for sender, message in st.session_state['history']:
    if sender == '👤':
        st.markdown(f"**👤 {message}**")
    elif sender == '🤖':
        st.markdown(f"<div style='text-align: right'>**🤖 {message}**</div>", unsafe_allow_html=True)