from apis.layoutlm import LayoutLM
from apis.qc3.host_list import HostList
from PIL import Image
import pandas as pd
import gradio as gr
import os

layoutlm = None
hl = HostList(is_debug=True)

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 filter_green_out(img: Image):
    image_data = img.load()
    height,width = img.size
    for loop1 in range(height):
        for loop2 in range(width):
            (r,g,b) = image_data[loop1,loop2]
            if g < 70 and r < 70 and b < 70:
                (r,g,b) = (0,0,0)
            else:
                (r,g,b) = (255,255,255)
            image_data[loop1,loop2] = r,g,b
    
    img.save('./temp.jpg')
    return img

def ask(img: Image, question, top_k, max_answer_len, chk_is_remove_green) -> str:
    if chk_is_remove_green:
        img = filter_green_out(img.copy())

    return layoutlm.answer_the_question_without_filter(
        img, 
        question, 
        top_k=top_k, 
        max_answer_len=max_answer_len, 
        is_debug=True)

if __name__ == '__main__':

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

        with gr.Blocks() as demo:

            with gr.Tab('List'):
                with gr.Row():
                    with gr.Column():
                        list_inp_img = gr.Image(type="pil")
                        gr.Examples(
                            [['./examples/host-list1.JPG'], ['./examples/host-list2.JPG', './examples/host-list3.JPG']], 
                            list_inp_img
                        )

                    with gr.Column():
                        list_out_txt = gr.Textbox(label='Answer', interactive=False)

                        list_btn_ask = gr.Button('Ask me')
                        list_btn_ask.click(hl.process_image, [
                            list_inp_img
                        ], list_out_txt)

            with gr.Tab('Layout'):
                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')
                        sld_max_answer = gr.Slider(1, 10, value=1, step=1, label="Max answer", info="Top-K between 1 and 10")
                        sld_max_answer_len = gr.Slider(1, 200, value=15, step=1, label="Max answer length", info="Length between 15 and 200")
                        chk_is_remove_green = gr.Checkbox(label="Remove green", info="Do you need clean context?")
                        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, 
                            sld_max_answer, 
                            sld_max_answer_len,
                            chk_is_remove_green
                        ], txt_out_answer)

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