import gradio as gr
import json
from pathlib import Path

def load_json(file_path):
    with open(file_path, 'r') as file:
        return json.load(file)

def display_data(index, data):
    item = data[index]
    note = item.get('Note', 'No note available')
    summary = item.get('Summary', {})
    mcqs = item.get('MCQS', {})
    image_path = item.get('image', '')
    
    summary_html = "<h3>Summary</h3>"
    if 'chest_xray_summary' in summary:
        for topic in summary['chest_xray_summary']:
            summary_html += f"<p><strong>{topic['topic']}:</strong> {topic['description']}</p>"
    else:
        summary_html += "<p>No summary available</p>"
    
    mcqs_html = "<h3>MCQs</h3>"
    if isinstance(mcqs, dict) and 'questions' in mcqs:
        for i, q in enumerate(mcqs['questions'], 1):
            mcqs_html += f"<p><strong>Q{i}: {q['question']}</strong></p>"
            mcqs_html += "<ul>"
            for option in q['options']:
                mcqs_html += f"<li>{option}</li>"
            mcqs_html += "</ul>"
            mcqs_html += f"<p>Correct Answer: {q['correct_answer']}</p>"
            if 'location_in_text' in q:
                mcqs_html += f"<p><em>Location in text: {q['location_in_text']}</em></p>"
            mcqs_html += "<hr>"
    elif isinstance(mcqs, list):
        for i, q in enumerate(mcqs, 1):
            mcqs_html += f"<p><strong>Q{i}: {q['question']}</strong></p>"
            mcqs_html += "<ul>"
            for option in q['options']:
                mcqs_html += f"<li>{option}</li>"
            mcqs_html += "</ul>"
            mcqs_html += f"<p>Correct Answer: {q['correct_answer']}</p>"
            if 'location_in_text' in q:
                mcqs_html += f"<p><em>Location in text: {q['location_in_text']}</em></p>"
            mcqs_html += "<hr>"
    else:
        mcqs_html += "<p>No MCQs available for this item.</p>"
    
    return note, summary_html, mcqs_html, image_path, f"Item {index + 1} of {len(data)}"

def next_item(index, data):
    return min(index + 1, len(data) - 1)

def prev_item(index, data):
    return max(index - 1, 0)

def create_app(json_path):
    data = load_json(json_path)
    
    with gr.Blocks() as app:
        gr.Markdown("# Chest X-Ray Analysis Viewer")
        
        with gr.Row():
            with gr.Column(scale=1):
                image = gr.Image(label="X-Ray Image")
            with gr.Column(scale=2):
                note = gr.Textbox(label="Note", lines=4)
                summary = gr.HTML(label="Summary")
                mcqs = gr.HTML(label="MCQs")
        
        index = gr.State(0)
        item_count = gr.Textbox(label="Item Count")
        
        with gr.Row():
            prev_btn = gr.Button("Previous")
            next_btn = gr.Button("Next")
        
        def update_display(index):
            return display_data(index, data)
        
        next_btn.click(
            next_item,
            inputs=[index, gr.State(data)],
            outputs=[index]
        ).then(
            update_display,
            inputs=[index],
            outputs=[note, summary, mcqs, image, item_count]
        )
        
        prev_btn.click(
            prev_item,
            inputs=[index, gr.State(data)],
            outputs=[index]
        ).then(
            update_display,
            inputs=[index],
            outputs=[note, summary, mcqs, image, item_count]
        )
        
        app.load(
            update_display,
            inputs=[index],
            outputs=[note, summary, mcqs, image, item_count]
        )
    
    return app


json_path = Path("mimic_data.json")
app = create_app(json_path)
app.launch()