import pandas as pd
import gradio as gr
import numpy as np

df_table_dsc = pd.read_csv('./results/df_table_dsc.csv')
df_table_prd = pd.read_csv('./results/df_table_prd.csv')
n_lines = 12

defect_dict = {'NoDefect':'No defect', 'RO': 'Defect: Roots', 'OB': 'Defect: Surface damage', 
               'RB': 'Defect: Cracks, breaks, and collapses', 'PF': 'Defect: Production error', 
               'DE':'Defect: Deformation'}


css="""
#image-out {
  height: 400px;
  width: 400px;
}
"""

def conv_int(txt):
    try:
        return (int(txt))
    except:
        return 'NaN'

def return_prev(index):
    # Update current index based on navigation input
    index -=1 
    # Ensure index is within bounds
    index = max(0, min(index, len(df_table_dsc) - 1))

    image_path, gpt4, gpt4b, cogvlm, llava, defect_class, predictions = return_image(index)
    return index, index, image_path, gpt4, gpt4b, cogvlm, llava, defect_class, predictions

def return_next(index):
    # Update current index based on navigation input
    index +=1 
    # Ensure index is within bounds
    index = max(0, min(index, len(df_table_dsc) - 1))
    image_path, gpt4, gpt4b, cogvlm, llava, defect_class, predictions = return_image(index)
    return index, index, image_path, gpt4, gpt4b, cogvlm, llava, defect_class, predictions

def return_image(index):
    # Fetch row from DataFrame
    row = df_table_dsc.iloc[index]
    image_path = f"./images/{row['img_id']}"

    defect_class =  defect_dict[row['defect_type']]    
    
    row_prd = df_table_prd.iloc[index]
    
    predictions = f"XIE:{conv_int(row_prd['Xie']>0.5)}, GPT4:{conv_int(row_prd['GPT4'])}, GPT4s:{conv_int(row_prd['GPT4_basic'])}, CogVLM:{conv_int(row_prd['CogVLM'])}, LLaVA:{conv_int(row_prd['LLaVA'])}"

    # Return iamge path and descriptions
    return image_path, row['GPT4'], row['GPT4_basic'], row['CogVLM'], row['LLaVA'], defect_class, predictions


with gr.Blocks(css=css) as demo: 
    # gr.Markdown("""
                ## Demo for 'The Potential of Generative AI for the Urban Water Sector'
                ### Riccardo Taormina, Delft University of Technology (TU Delft), Department of Water Management                
                ### email: r.taormina@tudelft.nl""")                
    gr.Markdown("""
                ## Testing Large Multimodal Models for Sewer Defect Inspection 
                Press on \<Next\> or \<Previous\> to start!
                """)                
    index = gr.Number(value=-1, visible=False)      
    with gr.Row():
        with gr.Column(scale=1):
            img_out = gr.Image(type="filepath", label="Image", elem_id="image-out")
            with gr.Row():
                txt_item = gr.Textbox(label="Sample no.", min_width= 20)
                txt_defect_class = gr.Textbox(label="Defect class", interactive=False, min_width= 150)
            txt_xie_pred = gr.Textbox(label="Predictions (XIE = benchmark, 0 = No Defect, 1 = Defect)", min_width= 300, interactive=False)
            with gr.Row():
                prev_btn = gr.Button("Previous")            
                next_btn = gr.Button("Next")            
        with gr.Column(scale=1):
            gr.Markdown('Multimodal descriptions')
            with gr.Row():
                txt_out_GPT4 = gr.Textbox(label="GPT4", lines= n_lines, max_lines=n_lines)
                txt_out_GPT4s = gr.Textbox(label="GPT4 simple", lines= n_lines, max_lines=n_lines)    
            with gr.Row():
                txt_out_CogVLM = gr.Textbox(label="CogVLM", lines= n_lines, max_lines=n_lines)
                txt_out_LLaVa = gr.Textbox(label="LLaVa", lines= n_lines, max_lines=n_lines)
    prev_btn.click(fn=return_prev, inputs=index, outputs=[index, txt_item, img_out, 
                                                            txt_out_GPT4, txt_out_GPT4s, 
                                                            txt_out_CogVLM, txt_out_LLaVa, txt_defect_class, txt_xie_pred])
    next_btn.click(fn=return_next, inputs=index, outputs=[index, txt_item, img_out, 
                                                            txt_out_GPT4, txt_out_GPT4s, 
                                                            txt_out_CogVLM, txt_out_LLaVa, txt_defect_class, txt_xie_pred])

if __name__ == "__main__":
    demo.launch()