import numpy as np
import gradio as gr
import json
from datasets import load_dataset


theme = gr.themes.Base(
    secondary_hue="cyan",
).set(
    button_primary_background_fill='*primary_100',
    button_primary_text_color='*neutral_950',
    border_color_accent="*primary_100",
    body_text_color="*neutral_950",
    block_border_color="*primary_100",
    block_title_text_color="*neutral_950",
    block_title_text_size="text_xl",
    panel_border_color="*primary_950",
    form_gap_width="2px",
    
)

# Daten laden
directory_path = load_dataset("HuberDa/OCTTestdaten")
disease_1 = "Alterbedinge Makulardegeneration (AMD)"
disease_2 = "Diabetische Netzhauterkrankung (DR)"
disease_3 = "Drusen"
disease_4 = "Chorodiale Neovaskularisation (CNV)"
disease_5 = "Diabetisches Makulaödem (DME)"
disease_6 = "Normal / Keine Krankheit"

# Bilder generieren
def image_generator(hf_dataset, split="train"):
    for index, example in enumerate(hf_dataset[split]):
        image = example['image']  # Access the image column
        image_array = np.array(image)  # Convert to a NumPy array
        yield index, image_array  # Yield the index and image

# Generator aktivieren
image_gen = image_generator(directory_path)


# Input der Bilder sammeln
def record_input(label: str, index: int, free_text: str):
    # Den Output mit index, label, and free text sammeln
    output = {
        "index": index,
        "label": label,
        "free_text": free_text
    }
    # Output notieren
    with open("output.txt", 'a') as f:
        json.dump(output, f)
        f.write("\n")
    return f"Bild {index} wurde als {label} klassifiziert. \nZusätzliche Notizen: {free_text}"

def start():
    try:
        # Nächstes Bild und Index holen
        index, image = next(image_gen)  # Index und Bild entpacken
        return str(index), image, ""  # Rückgabe von index, image, and leeren der Textbox
    except StopIteration:
        return "No more images", None, ""  # Falls keine Bilder übrig
    except Exception as e:
        return str(e), None, ""  # Für andere Probleme

with gr.Blocks(fill_height=True, fill_width=True, theme = theme) as demo:
    gr.Markdown("# OCTis Training Center")
       
    with gr.Row():
        # Bildanzeige
        img_block = gr.Image(visible=True, width=500, height=500, show_fullscreen_button=True, show_share_button=False, show_download_button=False, show_label=False)

    with gr.Row():
        output_message = gr.Textbox(label="Ergebnis", interactive=False)
        free_text_input = gr.Textbox(label="zusätzliche Informationen", placeholder="Hier Text eingeben...", lines=2)  
        
        # versteckte textbox für index
        index_box = gr.Textbox(label="Index", visible=False)
       
    with gr.Row():
        # Buttons 1 - 3 
        disease1_btn = gr.Button(value=disease_1,variant="primary")
        disease2_btn = gr.Button(value=disease_2,variant="primary")
        disease3_btn = gr.Button(value=disease_3,variant="primary")
        
    with gr.Row():
        # Button 4 - 6
        disease4_btn = gr.Button(value=disease_4,variant="primary")
        disease5_btn = gr.Button(value=disease_5,variant="primary")
        nodisease_btn = gr.Button(value=disease_6,variant="primary")

    with gr.Row():
        # Start und skip Button
        start_btn = gr.Button("Start")
        skip_btn = gr.Button("Überspringen")

    # Starten der App
    start_btn.click(fn=start, outputs=[index_box, img_block, free_text_input])
    skip_btn.click(fn=start, outputs=[index_box, img_block, free_text_input])

    # Labeling mit Buttons
    disease1_btn.click(fn=lambda idx, free_text: record_input(disease_1, idx, free_text), inputs=[index_box, free_text_input], outputs=output_message)
    disease2_btn.click(fn=lambda idx, free_text: record_input(disease_2, idx, free_text), inputs=[index_box, free_text_input], outputs=output_message)
    disease3_btn.click(fn=lambda idx, free_text: record_input(disease_3, idx, free_text), inputs=[index_box, free_text_input], outputs=output_message)
    disease4_btn.click(fn=lambda idx, free_text: record_input(disease_4, idx, free_text), inputs=[index_box, free_text_input], outputs=output_message)
    disease5_btn.click(fn=lambda idx, free_text: record_input(disease_5, idx, free_text), inputs=[index_box, free_text_input], outputs=output_message)
    nodisease_btn.click(fn=lambda idx, free_text: record_input(disease_6, idx, free_text), inputs=[index_box, free_text_input], outputs=output_message)

    # Nächstes Bild
    disease1_btn.click(fn=start, outputs=[index_box, img_block, free_text_input])
    disease2_btn.click(fn=start, outputs=[index_box, img_block, free_text_input])
    disease3_btn.click(fn=start, outputs=[index_box, img_block, free_text_input])
    disease4_btn.click(fn=start, outputs=[index_box, img_block, free_text_input])
    disease5_btn.click(fn=start, outputs=[index_box, img_block, free_text_input])
    nodisease_btn.click(fn=start, outputs=[index_box, img_block, free_text_input])

demo.queue()
demo.launch(share=True, debug=True)