import gradio as gr
from PIL import Image
import clipGPT
import vitGPT
import skimage.io as io
import PIL.Image
import difflib


def compare_and_highlight(text1, text2):
    matcher = difflib.SequenceMatcher(None, text1, text2)
    output = ''
    for op, a1, a2, b1, b2 in matcher.get_opcodes():
        if op == 'equal':
            output += f"**{text1[a1:a2]}**"  # Highlight matches in bold
        elif op == 'insert':
            output += f"<ins>{text2[b1:b2]}</ins>" 
        elif op == 'delete':
            output += f"<del>{text1[a1:a2]}</del>"
        elif op == 'replace':  
            # Handle replacements (more complex)
            output += f"<del>{text1[a1:a2]}</del> <ins>{text2[b1:b2]}</ins>" 
    return output


# Caption generation functions
def generate_caption_clipgpt(image):
    caption = clipGPT.generate_caption_clipgpt(image)
    return caption

def generate_caption_vitgpt(image):
    caption = vitGPT.generate_caption(image)
    return caption



with gr.Blocks() as demo:
    
    generated_captions = {  
    "CLIP-GPT2": "",
    "ViT-GPT2": "",
    }

    gr.HTML("<h1 style='text-align: center;'>MedViT: A Vision Transformer-Driven Method for Generating Medical Reports 🏥🤖</h1>")
    gr.HTML("<p style='text-align: center;'>You can generate captions by uploading an X-Ray and selecting a model of your choice below</p>")

    with gr.Row():
        sample_images = [
        "CXR191_IM-0591-1001.png",
        "CXR192_IM-0598-1001.png",
        "CXR193_IM-0601-1001.png",
        "CXR194_IM-0609-1001.png",
        "CXR195_IM-0618-1001.png"
    ]
        image = gr.Image(label="Upload Chest X-ray")
        gr.Gallery(
            value = sample_images,
            label="Sample Images",
        )
        # sample_images_gallery = gr.Gallery(
        #     value = sample_images,
        #     label="Sample Images",
        # )
    with gr.Row():
        model_choice = gr.Radio(["CLIP-GPT2", "ViT-GPT2", "ViT-CoAttention"], label="Select Model")
        generate_button = gr.Button("Generate Caption") 
    caption = gr.Textbox(label="Generated Caption") 

    def predict(img, model_name):
        if model_name == "CLIP-GPT2":
            return generate_caption_clipgpt(img)
        elif model_name == "ViT-GPT2":
            return generate_caption_vitgpt(img)
        else:
            return "Caption generation for this model is not yet implemented."     
        generated_captions[model_name] = caption 

    with gr.Row():
        caption1 = gr.Textbox(label="CLIP-GPT2")
        caption2 = gr.Textbox(label="ViT-GPT2")
        compare_button = gr.Button("Compare Captions")
    with gr.Row():
        comparison_result = gr.Textbox(label="Comparison Result")

    # Compare captions on button click
    compare_button.click(lambda: compare_and_highlight(
        generated_captions["CLIP-GPT2"], generated_captions["ViT-GPT2"]
        ), [], comparison_result) 


    generate_button.click(predict, [image, model_choice], caption)  # Trigger prediction on button click 
    # sample_images_gallery.change(predict, [sample_images_gallery, model_choice], caption)  # Handle sample images


demo.launch()