from apis.layoutlm import LayoutLM
import pandas as pd
import gradio as gr
import os

layoutlm = None

def auth(username, password):
    u = os.environ.get('USERNAME')
    p = os.environ.get('PASSWORD')
    return (username == u and password == p)

def inference(img) -> pd.DataFrame:
    return layoutlm.inference(img)

def ask(img, question) -> str:
    return layoutlm.answer_the_question(img, question, is_debug=True)

if __name__ == '__main__':

    try:
        layoutlm = LayoutLM()
        layoutlm.set_model(layoutlm.default_model)

        with gr.Blocks() as demo:

            with gr.Row():
                inp_img = gr.Image(type='pil')
            
                with gr.Column():
                    out = gr.Dataframe(
                        headers=['Data', 'Value'],
                        datatype=['str', 'str'],
                        row_count=8,
                        col_count=(2, 'fixed'), 
                        interactive=False
                    )

                    txt_custom_question = gr.Textbox(label='Your question')
                    btn_ask = gr.Button('Ask me')
                    txt_out_answer = gr.Textbox(label='Answer', interactive=False)

                    # event
                    inp_img.change(inference, inp_img, out)
                    btn_ask.click(ask, [inp_img, txt_custom_question], txt_out_answer)

        demo.launch(auth=auth)
    except Exception as e:
        print(str(e))