import gradio as gr

def single_inference_iit(head_img, head_id, tail_img, tail_id, question_txt, question_id):
    return question_txt

def single_inference_tti(head_txt, head_id, tail_txt, tail_id, question_img, question_id):
    return head_txt

def blended_inference_iti(head_img, head_id, tail_txt, tail_id, question_img, question_id):
    return tail_txt

def single_tab_iit():
    with gr.Column():
        gr.Markdown(""" $(I_h, I_t) : (T_q, ?)$
                    """)
        with gr.Row():
            with gr.Column():
                head_image = gr.Image(type='pil', label="Head Image")
                head_ent = gr.Textbox(lines=1, label="Head Entity")
            with gr.Column():
                tail_image = gr.Image(type='pil', label="Tail Image")
                tail_ent = gr.Textbox(lines=1, label="Tail Entity")
            with gr.Column():
                question_text = gr.Textbox(lines=1, label="Question Name")
                question_ent = gr.Textbox(lines=1, label="Question Entity")
                
    submit_btn = gr.Button("Submit")
    output_text = gr.Textbox(label="Output")
    
    submit_btn.click(fn=single_inference_iit, 
        inputs=[head_image, head_ent, tail_image, tail_ent, question_text, question_ent], 
        outputs=[output_text])
        
    examples=[['examples/qinghai_lake.png', 'Q201294', 'examples/inland_lake.png', 'Q31805992', "campaign", 'Q18812548']]
    ex = gr.Examples(
        examples=examples, 
        fn=single_inference_iit, 
        inputs=[head_image, head_ent, tail_image, tail_ent, question_text, question_ent], 
        outputs=[output_text],
        cache_examples=True, 
        run_on_click=True
    )
   
def single_tab_tti():    
    with gr.Column():
        gr.Markdown(""" $(T_h, T_t) : (I_q, ?)$
                    """)
        with gr.Row():
            with gr.Column():
                head_text = gr.Textbox(lines=1, label="Head Name")
                head_ent = gr.Textbox(lines=1, label="Head Entity")
            with gr.Column():
                tail_text = gr.Textbox(lines=1, label="Tail Name")
                tail_ent = gr.Textbox(lines=1, label="Tail Entity")
            with gr.Column():
                question_text = gr.Image(type='pil', label="Question Image")
                question_ent = gr.Textbox(lines=1, label="Question Entity")
    submit_btn = gr.Button("Submit")
    output_text = gr.Textbox(label="Output")
    
    submit_btn.click(fn=single_inference_iit, 
        inputs=[head_text, head_ent, tail_text, tail_ent, question_text, question_ent], 
        outputs=[output_text])
        
    examples=[['qinghai_lake', 'Q201294', 'inland_lake', 'Q31805992', 'examples/qinghai_lake.png', 'Q18812548']]
    ex = gr.Examples(
        examples=examples, 
        fn=single_inference_iit, 
        inputs=[head_text, head_ent, tail_text, tail_ent, question_text, question_ent], 
        outputs=[output_text],
        cache_examples=True, 
        run_on_click=True
    )
    
def blended_tab_iti():
    with gr.Column():
        gr.Markdown(""" $(I_h, T_t) : (I_q, ?)$
                    """)
        with gr.Row():
            with gr.Column():
                head_image = gr.Image(type='pil', label="Head Image")
                head_ent = gr.Textbox(lines=1, label="Head Entity")
            with gr.Column():
                tail_txt = gr.Textbox(lines=1, label="Tail Name")
                tail_ent = gr.Textbox(lines=1, label="Tail Entity")
            with gr.Column():
                question_image = gr.Image(type='pil', label="Question Image")
                question_ent = gr.Textbox(lines=1, label="Question Entity")
    submit_btn = gr.Button("Submit")
    output_text = gr.Textbox(label="Output")
    
    submit_btn.click(fn=single_inference_iit, 
        inputs=[head_image, head_ent, tail_txt, tail_ent, question_image, question_ent], 
        outputs=[output_text])
        
    examples=[['examples/qinghai_lake.png', 'Q201294', 'inland_lake', 'Q31805992', 'examples/inland_lake.png', 'Q18812548']]
    ex = gr.Examples(
        examples=examples, 
        fn=single_inference_iit, 
        inputs=[head_image, head_ent, tail_txt, tail_ent, question_image, question_ent], 
        outputs=[output_text],
        cache_examples=True, 
        run_on_click=True
    )


TITLE = """MKG Analogy"""

with gr.Blocks() as block:
    with gr.Column(elem_id="col-container"):
        gr.HTML(TITLE)

        with gr.Tab("Single Analogical Reasoning"):
            single_tab_iit()
            single_tab_tti()

        with gr.Tab("Blended Analogical Reasoning"):
            blended_tab_iti()
    
        # gr.HTML(ARTICLE)
        

block.queue(max_size=64).launch(enable_queue=True)