app.py

#!/usr/bin/env python3
"""
POC RMMM - Automatic Medical Report Generation with Ground Truth Comparison & Evaluation Metrics

This application provides a Gradio interface for generating medical reports from X-ray images
using the RMMM PyTorch model, with automatic evaluation metrics (BLEU-4, ROUGE-L) 
to compare against ground truth reports.
"""

# Standard library imports
import asyncio
import hashlib
import json
import os
import pickle
import re
import sys
import traceback
import warnings
from typing import Dict, List, Union

# Third-party imports
import gradio as gr
import nltk
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from rouge import Rouge
from sacrebleu import BLEU
from transformers import GPT2Tokenizer

# Configuration and warnings
warnings.filterwarnings("ignore", message=".*trust_remote_code.*")
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)

# Environment variables
os.environ["HF_DATASETS_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "0"  # Allow limited online access for core models

# Download NLTK data if needed
try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    nltk.download('punkt', quiet=True)

print("🖥️ Using device: CPU")

# Load MIMIC dataset from JSON file
def load_mimic_data() -> Dict[str, str]:
    """Load MIMIC dataset from JSON file.
    
    Returns:
        Dict[str, str]: Dictionary mapping image IDs to ground truth reports
    """
    json_path = "./data/sample_mimic_test.json"
    
    if not os.path.exists(json_path):
        print(f"Warning: {json_path} not found. Using empty dataset.")
        return {}
    
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # Convert to dictionary format for easy lookup by image ID
        ground_truth_reports = {}
        for item in data.get('sample_data', []):
            image_id = item.get('id')
            report = item.get('report', 'No report available.')
            if image_id:
                ground_truth_reports[image_id] = report
        
        print(f"Loaded {len(ground_truth_reports)} ground truth reports from MIMIC dataset")
        return ground_truth_reports
    
    except Exception as e:
        print(f"Error loading MIMIC data: {e}")
        return {}

# Load ground truth reports from JSON file
GROUND_TRUTH_REPORTS = load_mimic_data()

# Load model globally at startup
def load_rmmm_model():
    """Load RMMM model once at startup"""
    try:
        model_path = "./rmmm/rmmm_mimic_cut.pt"
        
        if not os.path.exists(model_path):
            print(f"❌ Model not found: {model_path}")
            return None
            
        print(f"🤖 Loading RMMM model from: {model_path}")
        print("🖥️ Target device: CPU")
        
        # Load model with explicit CPU mapping
        print("Loading model with CPU mapping...")
        scripted_model = torch.jit.load(model_path, map_location='cpu')
        scripted_model.eval()
        
        # Ensure all parameters are on CPU
        print("Moving all parameters to CPU...")
        scripted_model = scripted_model.cpu()
        
        # Verify device placement
        print("Verifying model device placement...")
        for param in scripted_model.parameters():
            if param.device != torch.device('cpu'):
                print(f"⚠️ Found parameter on {param.device}, moving to CPU")
                param.data = param.data.cpu()
        
        print(f"✅ RMMM model loaded successfully on CPU")
                
        return scripted_model
        
    except Exception as e:
        print(f"❌ Error loading RMMM model: {e}")
        traceback.print_exc()
        return None

# Load tokenizer globally at startup
def load_mimic_tokenizer():
    """Load MIMIC tokenizer once at startup"""
    try:
        cache_file = "./rmmm/tokenizer_cache/tokenizer.pkl"
        with open(cache_file, 'rb') as f:
            tokenizer_data = pickle.load(f)
        
        idx2token = tokenizer_data['idx2token']
        print(f"✅ Custom MIMIC tokenizer loaded with vocab size: {len(idx2token)}")
        return idx2token
        
    except Exception as e:
        print(f"⚠️ Failed to load custom tokenizer: {e}")
        return None

# Global model and tokenizer instances
print("🚀 Initializing RMMM application...")
RMMM_MODEL = load_rmmm_model()
MIMIC_TOKENIZER = load_mimic_tokenizer()

def get_available_image_paths() -> List[str]:
    """Get list of available image paths based on MIMIC JSON data.
    
    Returns:
        List[str]: List of available image file paths
    """
    json_path = "./data/sample_mimic_test.json"
    
    if not os.path.exists(json_path):
        # Fallback to scanning images directory
        images_dir = "./images"
        if os.path.exists(images_dir):
            return [os.path.join(images_dir, f) for f in os.listdir(images_dir) 
                   if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
        return []
    
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        image_paths = []
        for item in data.get('sample_data', []):
            image_id = item.get('id')
            if image_id:
                # Try different possible image extensions
                for ext in ['.jpg', '.jpeg', '.png']:
                    image_path = f"./images/{image_id}{ext}"
                    if os.path.exists(image_path):
                        image_paths.append(image_path)
                        break
        
        print(f"Found {len(image_paths)} available images from MIMIC dataset")
        return image_paths
    
    except Exception as e:
        print(f"Error loading image paths: {e}")
        return []

def preprocess_text_for_metrics(text: str) -> str:
    """Preprocess text for metric calculation.
    
    Args:
        text (str): Raw text to preprocess
        
    Returns:
        str: Cleaned and preprocessed text
    """
    if not text or text.strip() == "":
        return ""
    
    # Remove markdown formatting and special characters
    text = re.sub(r'\*\*(.*?)\*\*', r'\1', text)  # Remove **bold**
    text = re.sub(r'[📊📋🤖🩻]', '', text)  # Remove emojis
    
    # Extract only the main report content, ignore metadata
    lines = text.split('\n')
    report_lines = []
    in_report = False
    
    for line in lines:
        line = line.strip()
        if 'RADIOLOGIST REPORT:' in line or 'IMPRESSION:' in line or 'FINDINGS:' in line:
            in_report = True
            continue
        elif line.startswith('**') and ':' in line:
            # Skip metadata lines like **Study ID:** etc
            continue
        elif in_report and line:
            report_lines.append(line)
    
    # If no structured report found, use the whole cleaned text
    if not report_lines:
        report_lines = [line.strip() for line in lines if line.strip() and not line.startswith('**')]
    
    result = ' '.join(report_lines).strip()
    
    # Additional cleaning
    result = re.sub(r'\s+', ' ', result)  # Multiple spaces to single
    result = re.sub(r'[^\w\s\.\,\;\:\-\(\)]', '', result)  # Keep only basic punctuation
    
    return result

def calculate_evaluation_metrics(prediction: str, ground_truth: str) -> Dict[str, Union[float, str, None]]:
    """Calculate BLEU-4 and ROUGE-L metrics.
    
    Args:
        prediction (str): Generated prediction text
        ground_truth (str): Reference ground truth text
        
    Returns:
        Dict[str, Union[float, str, None]]: Dictionary containing metric scores and error info
    """
    
    if not prediction or not ground_truth:
        return {
            'bleu4_score': 0.0,
            'rougeL_f': 0.0,
            'error': 'Empty prediction or ground truth'
        }
    
    try:
        # Preprocess texts
        pred_clean = preprocess_text_for_metrics(prediction)
        gt_clean = preprocess_text_for_metrics(ground_truth)
        
        # Apply lowercase for better comparison
        pred_clean = pred_clean.lower()
        gt_clean = gt_clean.lower()
        
        if not pred_clean or not gt_clean:
            return {
                'bleu4_score': 0.0,
                'rougeL_f': 0.0,
                'error': 'Empty text after preprocessing'
            }
        
        # Calculate BLEU-4 score
        try:
            bleu = BLEU()
            # BLEU-4 expects list of references and hypothesis
            bleu4_score = bleu.sentence_score(pred_clean, [gt_clean]).score / 100.0  # Convert to 0-1 range
        except Exception as e:
            print(f"BLEU-4 calculation error: {e}")
            bleu4_score = 0.0
        
        # Calculate ROUGE-L score only
        try:
            rouge = Rouge()
            rouge_scores = rouge.get_scores(pred_clean, gt_clean)
            
            rougeL_f = rouge_scores[0]['rouge-l']['f']
        except Exception as e:
            print(f"ROUGE-L calculation error: {e}")
            rougeL_f = 0.0
        
        return {
            'bleu4_score': round(bleu4_score, 4),
            'rougeL_f': round(rougeL_f, 4),
            'error': None
        }
        
    except Exception as e:
        return {
            'bleu4_score': 0.0,
            'rougeL_f': 0.0,
            'error': f'Metric calculation error: {str(e)}'
        }

def format_metrics_display(metrics: Dict[str, Union[float, str, None]]) -> str:
    """Format metrics for display with modern HTML styling.
    
    Args:
        metrics (Dict[str, Union[float, str, None]]): Dictionary containing metric scores
        
    Returns:
        str: Formatted metrics display string with HTML
    """
    if metrics.get('error'):
        return f"""
        <div class="card">
            <div class="card-header">
                <div class="card-icon" style="background: var(--error-50); color: var(--error-600);">⚠️</div>
                <div>
                    <h2 class="card-title">Metrics Error</h2>
                    <p class="card-subtitle">Unable to calculate metrics</p>
                </div>
            </div>
            <div class="card-content">
                <p style="color: var(--error-600); font-weight: 500;">{metrics['error']}</p>
            </div>
        </div>
        """

    # Determine performance levels and colors
    bleu_score = metrics['bleu4_score']
    rouge_score = metrics['rougeL_f']

    def get_performance_badge(score: float) -> str:
        if score > 0.3:
            return '<span class="performance-badge badge-good">🟢 Excellent</span>'
        elif score > 0.1:
            return '<span class="performance-badge badge-fair">🟡 Fair</span>'
        else:
            return '<span class="performance-badge badge-poor">🔴 Low</span>'

    return f"""
    <div class="metrics-container">
        <div class="metrics-header">
            <h2><span>📊</span> Evaluation Metrics</h2>
        </div>

        <div class="metrics-grid">
            <div class="metric-card metric-bleu">
                <div class="metric-label">BLEU-4 Score</div>
                <div class="metric-value">{bleu_score:.4f}</div>
                {get_performance_badge(bleu_score)}
            </div>

            <div class="metric-card metric-rouge">
                <div class="metric-label">ROUGE-L F1</div>
                <div class="metric-value">{rouge_score:.4f}</div>
                {get_performance_badge(rouge_score)}
            </div>
        </div>

        <div class="metrics-info">
            💡 Higher scores indicate greater similarity to the reference report
        </div>
    </div>
    """

def inference_torch_model_fast(image_input):
    """Run inference with pre-loaded PyTorch model (fast version)"""
    try:
        # Check if model is loaded
        if RMMM_MODEL is None:
            return "❌ Error: RMMM model was not loaded correctly on startup."
        
        print(f"🔄 Running inference with pre-loaded RMMM model")
        print(f"🖼️ Image input type: {type(image_input)}")

        # Handle different input types
        if isinstance(image_input, str):
            # If it's a string, it's a file path - load it with Pillow
            print(f"📁 Loading image from path: {image_input}")
            image_input = Image.open(image_input)
        elif isinstance(image_input, np.ndarray):
            # If it's a numpy array, convert to PIL Image
            print(f"🔢 Converting numpy array to PIL Image")
            if image_input.dtype != np.uint8:
                image_input = (image_input * 255).astype(np.uint8)
            image_input = Image.fromarray(image_input)
        elif hasattr(image_input, 'mode'):
            # Already a PIL Image
            print(f"🖼️ Already a PIL Image")
        else:
            print(f"⚠️ Unknown image input type: {type(image_input)}")

        if image_input.mode != "RGB":
            image_input = image_input.convert("RGB")

        print(f"✅ Image loaded successfully: {image_input.size}")

        image_input = image_input.resize((224, 224))
        image_array = np.array(image_input).astype(np.float32)
        
        # Debug: Print image statistics to verify different images
        print(f"Image stats - Mean: {image_array.mean():.4f}, Std: {image_array.std():.4f}")
        print(f"Image range - Min: {image_array.min():.4f}, Max: {image_array.max():.4f}")
        
        # Calculate hash to verify different images
        image_hash = hashlib.md5(image_array.tobytes()).hexdigest()[:8]
        print(f"Image hash (first 8 chars): {image_hash}")
        
        # Apply ImageNet normalization (matching training preprocessing)
        image_array = image_array / 255.0
        
        # ImageNet mean and std values
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        
        # Apply normalization per channel
        for i in range(3):
            image_array[:, :, i] = (image_array[:, :, i] - mean[i]) / std[i]
        
        image_array = np.transpose(image_array, (2, 0, 1))  # HWC -> CHW
        image_tensor = torch.tensor(image_array, dtype=torch.float32, device='cpu').unsqueeze(0)

        print(f"Input tensor shape: {image_tensor.shape}")
        print(f"Input tensor device: {image_tensor.device}")
        print(f"Input tensor stats - Mean: {image_tensor.mean():.4f}, Std: {image_tensor.std():.4f}")
        print(f"Input tensor range - Min: {image_tensor.min():.4f}, Max: {image_tensor.max():.4f}")
        
        # Ensure model is on CPU before inference
        if hasattr(RMMM_MODEL, 'cpu'):
            RMMM_MODEL.cpu()

        # Run inference (model already loaded)
        with torch.no_grad():
            outputs = RMMM_MODEL(image_tensor)
            # Ensure outputs are on CPU
            outputs = outputs.cpu()
            print(f"Model output shape: {outputs.shape}")
            print(f"Model output device: {outputs.device}")
            print(f"Model output dtype: {outputs.dtype}")
            
            # Handle different output types for stats
            if outputs.dtype in [torch.float32, torch.float64]:
                print(f"Output stats - Mean: {outputs.mean():.4f}, Std: {outputs.std():.4f}")
                print(f"Output variance: {outputs.var():.6f}")
            else:
                # For integer outputs (token IDs), show basic stats
                print(f"Output stats - Min: {outputs.min()}, Max: {outputs.max()}")
                print(f"Output unique values: {len(torch.unique(outputs))}")
            
            # Debug: Print raw outputs to see if they vary
            if len(outputs.shape) >= 2:
                print(f"First few output values: {outputs.flatten()[:10]}")
            else:
                print(f"Output values: {outputs[:10] if len(outputs) > 10 else outputs}")

    # Process the output - Check if outputs are probabilities, logits, or token IDs
        if len(outputs.shape) == 3:
            # If 3D (batch, seq_len, vocab_size), take argmax
            print("Processing 3D output (batch, seq_len, vocab_size)")
            token_ids = torch.argmax(outputs, dim=-1)
        elif len(outputs.shape) == 2:
            # If 2D, check the dtype and values to determine if it's token IDs or logits
            if outputs.dtype in [torch.long, torch.int32, torch.int64]:
                print("Processing 2D output as token IDs (integer dtype)")
                token_ids = outputs
            elif outputs.max() > 1000:  # Likely token IDs already
                print("Processing 2D output as token IDs (high values)")
                token_ids = outputs.long()
            else:
                print("Processing 2D output as logits, taking argmax")
                # Treat as logits and take argmax
                token_ids = torch.argmax(outputs, dim=-1)
        elif len(outputs.shape) == 1:
            # If 1D, likely already token IDs
            print("Processing 1D output as token IDs")
            token_ids = outputs
        else:
            print(f"Unexpected output shape: {outputs.shape}")
            token_ids = outputs
            
        # Remove batch dimension if necessary
        if len(token_ids.shape) == 2:
            token_ids = token_ids[0]
            
        token_ids = token_ids.cpu().numpy().astype(np.int32)
        
        print(f"Token IDs shape: {token_ids.shape}")
        print(f"Token IDs sample: {token_ids[:10]}")
        print(f"Token IDs unique count: {len(np.unique(token_ids))}")  # Check diversity
        
        print(f"Token IDs shape: {token_ids.shape}")
        print(f"Token IDs sample: {token_ids[:10]}")

        # Decode using pre-loaded tokenizer
        if MIMIC_TOKENIZER is not None:
            # Use custom MIMIC tokenizer
            tokens = []
            for token_id in token_ids:
                if token_id == 0:  # End token
                    break
                if token_id in MIMIC_TOKENIZER:
                    tokens.append(MIMIC_TOKENIZER[token_id])
                    
            decoded_text = ' '.join(tokens).strip()
            print(f"✅ Used custom MIMIC tokenizer")
            
        else:
            # Fallback to GPT-2
            print(f"⚠️ Using GPT-2 fallback tokenizer")
            tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
            tokenizer.pad_token = tokenizer.eos_token
            
            # Limitar token_ids ao tamanho do vocabulário do GPT-2
            token_ids = np.clip(token_ids, 0, tokenizer.vocab_size - 1)
            decoded_text = tokenizer.decode(token_ids, skip_special_tokens=True).strip()

        print(f"Decoded text length: {len(decoded_text)}")
        print(f"Decoded text preview: {decoded_text[:100]}...")

    # If decoded text is too short, use fallback
        if len(decoded_text) < 10:
            decoded_text = (
                f"Medical Report - RMMM Model:\n\n"
                f"Chest X-ray analysis completed using PyTorch model. "
                f"The radiological examination has been processed successfully.\n\n"
                f"Model: rmmm_mimic_cut.pt\n"
                f"Status: Processing completed"
            )

        return decoded_text

    except Exception as e:
        error_msg = f"❌ Error processing with the RMMM model: {str(e)}"
        print(error_msg)
        traceback.print_exc()
        return error_msg

def get_ground_truth_from_filename(selected_image_filename):
    """Get ground truth report from memorized filename"""
    
    if not selected_image_filename:
        return "Ground truth not available."
    
    # Extract ID from filename
    filename = os.path.basename(selected_image_filename)
    image_id = filename.replace('.jpg', '').replace('.jpeg', '').replace('.png', '')
    
    print(f"Debug - selected_image_filename: {selected_image_filename}")
    print(f"Debug - extracted image_id: {image_id}")
    
    # Check if we have ground truth for this image
    if image_id and image_id in GROUND_TRUTH_REPORTS:
        report = GROUND_TRUTH_REPORTS[image_id]
        # Return only the clean report text without metadata
        return report.strip()
    
    return (
        f"Ground truth not available for this image (ID: {image_id}). "
        f"Upload one of the example images to see ground truth comparison."
    )

def inference_image_pipe_with_state(image_input, selected_image_filename):
    """Main inference function that uses memorized filename for ground truth"""
    
    # Get ground truth report from memorized filename
    ground_truth = get_ground_truth_from_filename(selected_image_filename)
    
    # Generate prediction using pre-loaded RMMM model
    prediction = inference_torch_model_fast(image_input)
    
    # Calculate evaluation metrics if both prediction and ground truth are available
    metrics_display = ""
    if (prediction and ground_truth and "Ground truth not available" not in ground_truth):
        metrics = calculate_evaluation_metrics(prediction, ground_truth)
        metrics_display = format_metrics_display(metrics)
    
    return prediction, ground_truth, metrics_display


def annotate_image(pil_image: Image.Image, report_text: str) -> Image.Image:
    """Create a lightweight annotated version of the input X-ray.

    This is a heuristic visualization for the demo: it searches the report
    for common radiology keywords and places labeled markers at
    deterministic locations (left lung, right lung, mediastinum) so users
    can visually inspect suggested areas. This is not a clinical
    localization method — it's a UX aid.
    """
    try:
        img = pil_image.convert("RGBA")
        overlay = Image.new('RGBA', img.size, (255, 255, 255, 0))
        draw = ImageDraw.Draw(overlay)

        keywords = {
            'consolidation': 'Consolidation',
            'effusion': 'Pleural effusion',
            'pneumothorax': 'Pneumothorax',
            'cardiomegaly': 'Cardiomegaly',
            'opacity': 'Opacity/Infiltrate',
            'infiltrate': 'Infiltrate',
            'atelectasis': 'Atelectasis'
        }

        found = []
        lower = (report_text or "").lower()
        for k, label in keywords.items():
            if k in lower:
                found.append(label)

        w, h = img.size

        # Deterministic marker positions: left lung, right lung, center
        positions = [
            (int(w * 0.25), int(h * 0.45)),
            (int(w * 0.75), int(h * 0.45)),
            (int(w * 0.5), int(h * 0.6)),
        ]

        # Draw translucent circles and text for each found keyword (up to 3)
        for i, label in enumerate(found[:3]):
            x, y = positions[i]
            r = int(min(w, h) * 0.12)
            # semi-transparent red
            draw.ellipse((x - r, y - r, x + r, y + r), fill=(255, 0, 0, 80), outline=(255, 0, 0, 180))
            # label background
            text_bg_h = 26
            text_w = 8 * len(label)
            text_x0 = x - text_w // 2 - 8
            text_y0 = y + r + 6
            draw.rectangle((text_x0, text_y0, text_x0 + text_w + 16, text_y0 + text_bg_h), fill=(0, 0, 0, 140))
            draw.text((text_x0 + 8, text_y0 + 4), label, fill=(255, 255, 255, 230))

        annotated = Image.alpha_composite(img, overlay).convert('RGB')
        return annotated
    except Exception as e:
        print(f"Annotation error: {e}")
        return pil_image


def explain_findings(ai_report: str, ground_truth: str):
    """Produce a human-friendly explanation, a step-by-step breakdown, and a detailed findings summary.

    This function focuses on clarity and educative explanations rather than clinical decision making.
    """
    try:
        ai = (ai_report or "").strip()
        gt = (ground_truth or "").strip()

        # Short summary
        if not ai:
            explanation = "No AI-generated report available to explain. Please generate a report first."
        else:
            explanation = (
                "Explanation (plain English):\n"
                f"The AI-generated report contains {len(ai.split())} words. Key findings are described below.\n\n"
            )

        # Find keywords and differences
        keywords = ['consolidation', 'effusion', 'pneumothorax', 'cardiomegaly', 'opacity', 'infiltrate', 'atelectasis']
        found = [k for k in keywords if k in ai.lower()]

        if found:
            explanation += "Detected terms in AI report: " + ", ".join(found) + ".\n"
        else:
            explanation += "No major standard keywords detected in the AI report.\n"

        # Step-by-step
        steps = [
            "1) Image preprocessing: The X-ray is resized and normalized for the model.",
            "2) Feature extraction: Visual features are encoded by the model backbone.",
            "3) Report generation: The language head outputs descriptive findings and impression.",
        ]

        if gt and "Ground truth not available" not in gt:
            steps.append("4) Evaluation: BLEU and ROUGE compare the AI text to the ground truth.")

        step_by_step_html = "<div class='card'><div class='card-header'><div class='card-icon'>🧭</div><div><h3 class='card-title'>Step-by-step</h3></div></div><div class='card-content'><ol>"
        for s in steps:
            step_by_step_html += f"<li>{s}</li>"
        step_by_step_html += "</ol></div></div>"

        # Detailed findings: show AI vs GT short diff
        detailed = "AI Report (trimmed):\n" + (ai[:150] + "..." if len(ai) > 150 else ai) + "\n\n"
        if gt and "Ground truth not available" not in gt:
            detailed += "Ground Truth (trimmed):\n" + (gt[:150] + "..." if len(gt) > 150 else gt) + "\n\n"

        detailed += "Notes:\n- This explanation is a textual aid. If localization is visible, review the annotated image.\n"

        return explanation, detailed, step_by_step_html
    except Exception as e:
        return f"Explanation generation error: {e}", "", ""


def save_report_file(selected_image_path: str, ai_report: str, ground_truth: str, metrics_html: str) -> str:
    """Save a simple text report (AI + GT + metrics) to ./reports and return the filepath.

    Returns a human-readable HTML status string for display in the status box.
    """
    try:
        os.makedirs("./reports", exist_ok=True)
        base = os.path.basename(selected_image_path) if selected_image_path else "unnamed"
        name = os.path.splitext(base)[0]
        filename = f"./reports/{name}_report.txt"
        with open(filename, 'w', encoding='utf-8') as f:
            f.write("AI-Generated Report:\n")
            f.write((ai_report or "") + "\n\n")
            f.write("Ground Truth Report:\n")
            f.write((ground_truth or "") + "\n\n")
            f.write("Evaluation Metrics (HTML):\n")
            f.write((metrics_html or "") + "\n")

        status_html = f"<div class='card'><div class='card-header'><div class='card-icon'>💾</div><div><h3 class='card-title'>Report Saved</h3></div></div><div class='card-content'><p>Saved to: <code>{filename}</code></p></div></div>"
        return status_html
    except Exception as e:
        return f"<div class='card'><div class='card-header'><div class='card-icon'>❌</div><div><h3 class='card-title'>Save Failed</h3></div></div><div class='card-content'><p>Error: {e}</p></div></div>"

with gr.Blocks(
    title="XRaySwinGen / RMMM  - AI Medical Report Generator", 
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="gray",
        neutral_hue="slate",
        font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"]
    ),
    css="""
    /* Import Google Fonts */
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&family=JetBrains+Mono:wght@400;500&display=swap');
    
    /* CSS Variables for Design System - Dark Theme Optimized */
    :root {
        --primary-50: #1e3a8a;
        --primary-100: #1d4ed8;
        --primary-500: #3b82f6;
        --primary-600: #60a5fa;
        --primary-700: #93c5fd;
        --primary-900: #dbeafe;
        
        --success-50: #14532d;
        --success-100: #166534;
        --success-500: #22c55e;
        --success-600: #4ade80;
        
        --warning-50: #92400e;
        --warning-100: #a16207;
        --warning-500: #f59e0b;
        --warning-600: #fbbf24;
        
        --error-50: #991b1b;
        --error-100: #dc2626;
        --error-500: #ef4444;
        --error-600: #f87171;
        
        --gray-50: #111827;
        --gray-100: #1f2937;
        --gray-200: #374151;
        --gray-300: #4b5563;
        --gray-400: #6b7280;
        --gray-500: #9ca3af;
        --gray-600: #d1d5db;
        --gray-700: #e5e7eb;
        --gray-800: #f3f4f6;
        --gray-900: #ffffff;
        
        --shadow-sm: 0 1px 2px 0 rgb(0 0 0 / 0.05);
        --shadow-md: 0 4px 6px -1px rgb(0 0 0 / 0.1), 0 2px 4px -2px rgb(0 0 0 / 0.1);
        --shadow-lg: 0 10px 15px -3px rgb(0 0 0 / 0.1), 0 4px 6px -4px rgb(0 0 0 / 0.1);
        --shadow-xl: 0 20px 25px -5px rgb(0 0 0 / 0.1), 0 8px 10px -6px rgb(0 0 0 / 0.1);
        
        --radius-sm: 0.375rem;
        --radius-md: 0.5rem;
        --radius-lg: 0.75rem;
        --radius-xl: 1rem;
        --radius-2xl: 1.5rem;
    }
    
    /* Global styles */
    .gradio-container {
        max-width: 1800px !important;
        margin: 0 auto !important;
        font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
        padding: 0 1rem !important;
        background-color: #0A0F1E !important;
        min-height: 100vh !important;
    }

    /* Estilo ao imprimir */
    @media print {
    body {
        background-color: #0A0F1E !important;
        -webkit-print-color-adjust: exact; /* Para WebKit (Chrome, Safari) */
        print-color-adjust: exact;         /* Padrão moderno */
        color: white; /* Garante contraste */
    }
    }
    
    /* Modern Header with Glass Effect */
    .main-header {
        background: linear-gradient(135deg, var(--primary-600) 0%, var(--primary-700) 50%, var(--primary-900) 100%) !important;
        backdrop-filter: blur(10px) !important;
        color: white !important;
        padding: 2rem 1.5rem !important;
        margin: 0 0 2rem 0 !important;
        box-shadow: var(--shadow-xl) !important;
        position: relative !important;
        overflow: hidden !important;
        border-radius: 1rem !important;
    }
    
    .main-header::before {
        content: '' !important;
        position: absolute !important;
        top: 0 !important;
        left: 0 !important;
        right: 0 !important;
        bottom: 0 !important;
        background: linear-gradient(45deg, rgba(255,255,255,0.1) 0%, rgba(255,255,255,0.05) 100%) !important;
        pointer-events: none !important;
        border-radius: 1rem !important;
    }
    
    .main-header h1 {
        font-size: 2.5rem !important;
        font-weight: 800 !important;
        margin: 0 0 0.5rem 0 !important;
        text-shadow: 0 2px 4px rgba(0,0,0,0.3) !important;
        letter-spacing: -0.025em !important;
        position: relative !important;
        z-index: 1 !important;
        border-radius: 1rem !important;
    }
    
    .main-header p {
        font-size: 1.1rem !important;
        margin: 0 0 1rem 0 !important;
        opacity: 0.95 !important;
        font-weight: 400 !important;
        position: relative !important;
        z-index: 1 !important;
        border-radius: 1rem !important;
    }
    
    .status-badges {
        display: flex !important;
        gap: 0.75rem !important;
        flex-wrap: wrap !important;
        margin-top: 1rem !important;
        position: relative !important;
        z-index: 1 !important;
    }
    
    .status-badge {
        display: inline-flex !important;
        align-items: center !important;
        gap: 0.5rem !important;
        background: rgba(255,255,255,0.2) !important;
        backdrop-filter: blur(10px) !important;
        padding: 0.5rem 1rem !important;
        border-radius: var(--radius-lg) !important;
        font-size: 0.875rem !important;
        font-weight: 500 !important;
        border: 1px solid rgba(255,255,255,0.2) !important;
    }
    
    /* Information grid layout */
    .info-grid {
        display: grid !important;
        grid-template-columns: repeat(auto-fit, minmax(280px, 1fr)) !important;
        gap: 1.5rem !important;
        justify-items: stretch !important;
        align-items: stretch !important;
    }
    
    /* Modern Card System - Dark Theme */
    .card {
        background: var(--gray-100) !important;
        border-radius: var(--radius-xl) !important;
        padding: 1.5rem !important;
        box-shadow: var(--shadow-md) !important;
        border: 1px solid var(--gray-200) !important;
        margin-bottom: 1.5rem !important;
        transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
        position: relative !important;
        overflow: hidden !important;
    }
    
    .card:hover {
        box-shadow: var(--shadow-lg) !important;
        transform: translateY(-2px) !important;
        border-color: var(--gray-300) !important;
        background: var(--gray-200) !important;
    }
    
    .card-header {
        display: flex !important;
        align-items: center !important;
        gap: 0.75rem !important;
        margin-bottom: 1rem !important;
        padding-bottom: 1rem !important;
        border-bottom: 1px solid var(--gray-200) !important;
        text-align: left !important;
    }
    
    .card-icon {
        width: 2.5rem !important;
        height: 2.5rem !important;
        border-radius: var(--radius-md) !important;
        display: flex !important;
        align-items: center !important;
        justify-content: center !important;
        font-size: 1.25rem !important;
        background: var(--primary-100) !important;
        color: var(--primary-600) !important;
        flex-shrink: 0 !important;
    }
    
    .card-title {
        color: var(--gray-900) !important;
        font-weight: 700 !important;
        font-size: 1.25rem !important;
        margin: 0 !important;
        line-height: 1.4 !important;
        text-align: left !important;
    }
    
    .card-subtitle {
        color: var(--gray-700) !important;
        font-size: 0.9375rem !important;
        margin: 0.25rem 0 0 0 !important;
        font-weight: 500 !important;
        text-align: left !important;
    }
    
    .card-content {
        color: var(--gray-800) !important;
        line-height: 1.6 !important;
        font-size: 1rem !important;
        text-align: left !important;
    }
    
    .card-content p {
        color: var(--gray-800) !important;
        margin-bottom: 0.75rem !important;
        text-align: left !important;
    }
    
    .card-content a {
        color: var(--primary-600) !important;
        font-weight: 600 !important;
        text-decoration: none !important;
    }
    
    .card-content a:hover {
        color: var(--primary-700) !important;
        text-decoration: underline !important;
    }
    
    /* Process Steps Card - Dark Theme */
    .process-steps {
        background: var(--gray-100) !important;
        border: 2px solid var(--primary-200) !important;
    }
    
    .step {
        display: flex !important;
        align-items: flex-start !important;
        gap: 1rem !important;
        margin-bottom: 1rem !important;
        padding: 1rem !important;
        border-radius: var(--radius-lg) !important;
        background: var(--gray-200) !important;
        border: 1px solid var(--gray-300) !important;
        transition: all 0.2s ease !important;
    }
    
    .step:hover {
        background: var(--gray-300) !important;
        border-color: var(--primary-300) !important;
    }
    
    .step-number {
        width: 2.5rem !important;
        height: 2.5rem !important;
        border-radius: 50% !important;
        background: var(--primary-500) !important;
        color: white !important;
        display: flex !important;
        align-items: center !important;
        justify-content: center !important;
        font-weight: 700 !important;
        font-size: 1rem !important;
        flex-shrink: 0 !important;
    }
    
    .step-content {
        text-align: left !important;
    }
    
    .step-content h4, .step-content h5 {
        color: var(--gray-900) !important;
        font-weight: 700 !important;
        font-size: 1.125rem !important;
        margin: 0 0 0.5rem 0 !important;
        text-align: left !important;
    }
    
    .step-content p {
        color: var(--gray-800) !important;
        font-size: 1rem !important;
        margin: 0 !important;
        line-height: 1.5 !important;
        text-align: left !important;
    }
    
    /* Enhanced Gallery - Dark Theme */
    .gallery-section {
        display: grid !important;
        grid-template-columns: 2fr 1fr !important;
        gap: 2rem !important;
        margin: 2rem 0 !important;
    }
    
    .gallery-card {
        background: var(--gray-100) !important;
        border-radius: var(--radius-xl) !important;
        padding: 0 !important;
        box-shadow: var(--shadow-lg) !important;
        border: 1px solid var(--gray-200) !important;
        overflow: hidden !important;
        transition: all 0.3s ease !important;
        height: fit-content !important;
    }
    
    .gallery-header {
        background: var(--gray-200) !important;
        padding: 1.5rem !important;
        border-bottom: 1px solid var(--gray-300) !important;
        text-align: center !important;
    }
    
    .gallery-title {
        color: var(--gray-900) !important;
        font-weight: 700 !important;
        font-size: 1.5rem !important;
        margin: 0 0 0.5rem 0 !important;
        display: flex !important;
        align-items: center !important;
        justify-content: center !important;
        gap: 0.75rem !important;
        text-align: center !important;
    }
    
    .gallery-subtitle {
        color: var(--gray-800) !important;
        font-size: 1rem !important;
        margin: 0 !important;
        font-weight: 500 !important;
        text-align: center !important;
    }
    
    .gradio-gallery {
        border-radius: 0 !important;
        border: none !important;
        background: var(--gray-200) !important;
        padding: 1rem !important;
        min-height: 400px !important;
    }
    
    .gradio-gallery img {
        border-radius: var(--radius-lg) !important;
        transition: all 0.3s ease !important;
        border: 2px solid transparent !important;
    }
    
    .gradio-gallery img:hover {
        transform: scale(1.02) !important;
        border-color: var(--primary-500) !important;
        box-shadow: var(--shadow-md) !important;
    }
    
    /* Status and Controls Panel - Dark Theme */
    .controls-panel {
        display: flex !important;
        flex-direction: column !important;
        gap: 1.5rem !important;
        height: fit-content !important;
        max-height: 80vh !important;
        overflow-y: auto !important;
        position: sticky !important;
        top: 2rem !important;
        padding-right: 0.5rem !important;
    }
    
    /* Status Display */
    .status-display {
        background: var(--gray-100) !important;
        border: 1px solid var(--gray-200) !important;
        border-radius: var(--radius-xl) !important;
        padding: 1.5rem !important;
        text-align: center !important;
        box-shadow: var(--shadow-sm) !important;
        margin: 1.5rem 0 !important;
    }
    
    .status-processing {
        background: var(--primary-100) !important;
        border-color: var(--primary-300) !important;
        animation: pulse-soft 2s ease-in-out infinite !important;
    }
    
    .status-success {
        background: var(--success-100) !important;
        border-color: var(--success-300) !important;
    }
    
    .status-error {
        background: var(--error-100) !important;
        border-color: var(--error-300) !important;
    }
    
    @keyframes pulse-soft {
        0%, 100% { 
            opacity: 1; 
            transform: scale(1); 
        }
        50% { 
            opacity: 0.9; 
            transform: scale(1.005); 
        }
    }
    
    /* Modern Buttons */
    .btn {
        border-radius: var(--radius-lg) !important;
        font-weight: 600 !important;
        transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1) !important;
        border: none !important;
        display: inline-flex !important;
        align-items: center !important;
        justify-content: center !important;
        gap: 0.5rem !important;
        font-size: 0.9375rem !important;
        line-height: 1 !important;
        text-decoration: none !important;
        cursor: pointer !important;
        position: relative !important;
        overflow: hidden !important;
    }
    
    .btn-primary {
        background: linear-gradient(135deg, var(--primary-500) 0%, var(--primary-600) 100%) !important;
        color: white !important;
        box-shadow: var(--shadow-md) !important;
        padding: 0.875rem 1.5rem !important;
    }
    
    .btn-primary:hover {
        background: linear-gradient(135deg, var(--primary-600) 0%, var(--primary-700) 100%) !important;
        transform: translateY(-1px) !important;
        box-shadow: var(--shadow-lg) !important;
    }
    
    .btn-primary:active {
        transform: translateY(0) !important;
        box-shadow: var(--shadow-md) !important;
    }
    
    .btn-secondary {
        background: white !important;
        color: var(--gray-700) !important;
        border: 1px solid var(--gray-300) !important;
        box-shadow: var(--shadow-sm) !important;
        padding: 0.75rem 1.25rem !important;
    }
    
    .btn-secondary:hover {
        background: var(--gray-50) !important;
        border-color: var(--gray-400) !important;
        color: var(--gray-900) !important;
        box-shadow: var(--shadow-md) !important;
    }
    
    .gradio-button {
        border-radius: var(--radius-lg) !important;
        font-weight: 600 !important;
        transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1) !important;
        border: none !important;
        box-shadow: var(--shadow-md) !important;
    }
    
    .gradio-button.primary {
        background: linear-gradient(135deg, var(--primary-500) 0%, var(--primary-600) 100%) !important;
        color: white !important;
        font-size: 1rem !important;
        padding: 0.875rem 1.5rem !important;
    }
    
    .gradio-button.primary:hover {
        background: linear-gradient(135deg, var(--primary-600) 0%, var(--primary-700) 100%) !important;
        transform: translateY(-1px) !important;
        box-shadow: var(--shadow-lg) !important;
    }
    
    .gradio-button.secondary {
        background: white !important;
        color: var(--gray-700) !important;
        border: 1px solid var(--gray-300) !important;
    }
    
    .gradio-button.secondary:hover {
        background: var(--gray-50) !important;
        border-color: var(--gray-400) !important;
        color: var(--gray-900) !important;
    }
    
    /* Enhanced Textboxes - Dark Theme */
    .gradio-textbox {
        border-radius: var(--radius-xl) !important;
        border: 1px solid var(--gray-200) !important;
        box-shadow: var(--shadow-sm) !important;
        background: var(--gray-100) !important;
        transition: all 0.2s ease !important;
        overflow: hidden !important;
    }
    
    .gradio-textbox:focus-within {
        box-shadow: 0 0 0 3px var(--primary-100), var(--shadow-md) !important;
        border-color: var(--primary-400) !important;
        background: var(--gray-200) !important;
    }
    
    .gradio-textbox textarea {
        font-family: 'Inter', sans-serif !important;
        line-height: 1.6 !important;
        font-size: 0.9375rem !important;
        border: none !important;
        padding: 1.25rem !important;
        color: var(--gray-900) !important;
        background: transparent !important;
        resize: vertical !important;
        min-height: 120px !important;
    }
    
    /* Controls panel textboxes - more compact */
    .controls-panel .gradio-textbox textarea {
        min-height: 100px !important;
        padding: 1rem !important;
        font-size: 0.875rem !important;
        line-height: 1.5 !important;
        color: var(--gray-900) !important;
    }
    
    .gradio-textbox label {
        color: var(--gray-900) !important;
        font-weight: 700 !important;
        font-size: 1.125rem !important;
        margin-bottom: 0.5rem !important;
        display: flex !important;
        align-items: center !important;
        gap: 0.5rem !important;
        text-align: center !important;
    }
    
    .report-ai {
        background: var(--primary-100) !important;
        border-color: var(--primary-300) !important;
    }
    
    .report-ai textarea {
        color: var(--gray-900) !important;
    }
    
    .report-ground-truth {
        background: var(--success-100) !important;
        border-color: var(--success-300) !important;
    }
    
    .report-ground-truth textarea {
        color: var(--gray-900) !important;
    }
    
    /* Enhanced Metrics Display - Dark Theme */
    .metrics-container {
        background: var(--gray-100) !important;
        border-radius: var(--radius-2xl) !important;
        padding: 0 !important;
        box-shadow: var(--shadow-lg) !important;
        border: 1px solid var(--gray-200) !important;
        overflow: hidden !important;
        margin: 1.5rem 0 !important;
    }
    
    .metrics-header {
        background: var(--primary-500) !important;
        color: white !important;
        padding: 1.5rem !important;
        position: relative !important;
        text-align: center !important;
    }
    
    .metrics-header h2 {
        color: white !important;
        font-weight: 700 !important;
        font-size: 1.5rem !important;
        margin: 0 !important;
        position: relative !important;
        z-index: 1 !important;
        display: flex !important;
        align-items: center !important;
        justify-content: center !important;
        gap: 0.75rem !important;
        text-align: center !important;
    }
    
    .metrics-grid {
        display: grid !important;
        grid-template-columns: 1fr 1fr !important;
        gap: 0 !important;
        margin: 0 !important;
    }
    
    .metric-card {
        padding: 1.5rem !important;
        position: relative !important;
        transition: all 0.2s ease !important;
        text-align: center !important;
        background: var(--gray-100) !important;
    }
    
    .metric-card:first-child {
        border-right: 1px solid var(--gray-300) !important;
    }
    
    .metric-card:hover {
        background: var(--gray-200) !important;
    }
    
    .metric-label {
        color: var(--gray-800) !important;
        font-size: 1rem !important;
        font-weight: 700 !important;
        text-transform: uppercase !important;
        letter-spacing: 0.05em !important;
        margin: 0 0 0.5rem 0 !important;
        text-align: center !important;
    }
    
    .metric-value {
        font-size: 2.5rem !important;
        font-weight: 800 !important;
        line-height: 1 !important;
        margin: 0 0 0.75rem 0 !important;
        font-family: 'JetBrains Mono', monospace !important;
        text-align: center !important;
    }
    
    .metric-bleu .metric-value {
        color: var(--primary-600) !important;
    }
    
    .metric-rouge .metric-value {
        color: var(--success-600) !important;
    }
    
    .performance-badge {
        display: inline-flex !important;
        align-items: center !important;
        justify-content: center !important;
        gap: 0.5rem !important;
        padding: 0.75rem 1.25rem !important;
        border-radius: var(--radius-lg) !important;
        font-size: 1rem !important;
        font-weight: 700 !important;
        border: 1px solid transparent !important;
        text-align: center !important;
    }
    
    .badge-good {
        background: var(--success-200) !important;
        color: var(--success-900) !important;
        border-color: var(--success-400) !important;
    }
    
    .badge-fair {
        background: var(--warning-200) !important;
        color: var(--warning-900) !important;
        border-color: var(--warning-400) !important;
    }
    
    .badge-poor {
        background: var(--error-200) !important;
        color: var(--error-900) !important;
        border-color: var(--error-400) !important;
    }
    
    .metrics-info {
        background: var(--gray-200) !important;
        padding: 1.25rem 1.5rem !important;
        border-top: 1px solid var(--gray-300) !important;
        color: var(--gray-900) !important;
        font-size: 1rem !important;
        font-weight: 600 !important;
        text-align: center !important;
    }
    
    /* Responsive Design */
    @media (max-width: 1200px) {
        .gallery-section {
            grid-template-columns: 1fr !important;
            gap: 1.5rem !important;
        }
        
        .controls-panel {
            position: relative !important;
            top: auto !important;
        }
        
        .info-grid {
            grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)) !important;
            gap: 1rem !important;
        }
    }
    
    @media (max-width: 1024px) {
        .gradio-container {
            padding: 0 0.75rem !important;
        }
        
        .main-header {
            padding: 1.5rem 1rem !important;
            border-radius: 1rem !important;
        }
        
        .main-header h1 {
            font-size: 2rem !important;
            border-radius: 1rem !important;
        }
        
        .metrics-grid {
            grid-template-columns: 1fr !important;
        }
        
        .metric-card:first-child {
            border-right: none !important;
            border-bottom: 1px solid var(--gray-200) !important;
        }
        
        .card {
            padding: 1rem !important;
        }
    }
    
    @media (max-width: 768px) {
        .main-header h1 {
            font-size: 1.75rem !important;
            border-radius: 1rem !important;
        }
        
        .main-header p {
            font-size: 1rem !important;
            border-radius: 1rem !important;
        }
        
        .status-badges {
            flex-direction: column !important;
            align-items: stretch !important;
        }
        
        .gradio-gallery {
            padding: 0.75rem !important;
        }
        
        .card {
            margin-bottom: 1rem !important;
        }
        
        .metric-value {
            font-size: 1.75rem !important;
        }
        
        .gallery-section {
            gap: 1rem !important;
        }
        
        .info-grid {
            grid-template-columns: 1fr !important;
            gap: 1rem !important;
        }
        
        .section-header {
            text-align: center !important;
            margin-bottom: 1.5rem !important;
        }
    }
    
    @media (max-width: 640px) {
        .gradio-container {
            padding: 0 0.5rem !important;
        }
        
        .main-header {
            padding: 1rem !important;
            border-radius: 1rem !important;
        }
        
        .main-header h1 {
            font-size: 1.5rem !important;
            border-radius: 1rem !important;
        }
        
        .card {
            padding: 0.75rem !important;
        }
        
        .gradio-textbox textarea {
            padding: 1rem !important;
            min-height: 150px !important;
        }
    }
    
    /* Animations and Loading States */
    @keyframes slideInUp {
        from {
            opacity: 0;
            transform: translateY(20px);
        }
        to {
            opacity: 1;
            transform: translateY(0);
        }
    }
    
    @keyframes fadeIn {
        from { opacity: 0; }
        to { opacity: 1; }
    }
    
    @keyframes shimmer {
        0% {
            background-position: -200px 0;
        }
        100% {
            background-position: calc(200px + 100%) 0;
        }
    }
    
    .loading-shimmer {
        background: linear-gradient(90deg, var(--gray-100) 25%, var(--gray-50) 50%, var(--gray-100) 75%) !important;
        background-size: 200px 100% !important;
        animation: shimmer 1.5s infinite !important;
    }
    
    .card-animate {
        animation: slideInUp 0.4s ease-out !important;
    }
    
    .fade-in {
        animation: fadeIn 0.3s ease-out !important;
    }
    
    /* Scrollbar Styling */
    ::-webkit-scrollbar {
        width: 8px !important;
        height: 8px !important;
    }
    
    ::-webkit-scrollbar-track {
        background: var(--gray-100) !important;
        border-radius: var(--radius-lg) !important;
    }
    
    ::-webkit-scrollbar-thumb {
        background: var(--gray-400) !important;
        border-radius: var(--radius-lg) !important;
        transition: background 0.2s ease !important;
    }
    
    ::-webkit-scrollbar-thumb:hover {
        background: var(--gray-500) !important;
    }
    
    /* Custom scrollbar for Firefox */
    * {
        scrollbar-width: thin !important;
        scrollbar-color: var(--gray-400) var(--gray-100) !important;
    }
    
    /* Section Headers - Dark Theme */
    .section-header {
        text-align: center !important;
        margin: 3rem 0 2rem 0 !important;
    }
    
    .section-title {
        color: var(--gray-900) !important;
        font-size: 2.25rem !important;
        font-weight: 800 !important;
        margin: 0 0 0.5rem 0 !important;
        letter-spacing: -0.025em !important;
        text-align: center !important;
    }
    
    .section-subtitle {
        color: var(--gray-800) !important;
        font-size: 1.25rem !important;
        margin: 0 !important;
        font-weight: 500 !important;
        text-align: center !important;
    }
    
    /* Utility Classes */
    .text-center { text-align: center !important; }
    .text-left { text-align: left !important; }
    .text-right { text-align: right !important; }
    
    .mb-0 { margin-bottom: 0 !important; }
    .mb-1 { margin-bottom: 0.25rem !important; }
    .mb-2 { margin-bottom: 0.5rem !important; }
    .mb-3 { margin-bottom: 0.75rem !important; }
    .mb-4 { margin-bottom: 1rem !important; }
    .mb-6 { margin-bottom: 1.5rem !important; }
    .mb-8 { margin-bottom: 2rem !important; }
    
    .p-0 { padding: 0 !important; }
    .p-1 { padding: 0.25rem !important; }
    .p-2 { padding: 0.5rem !important; }
    .p-3 { padding: 0.75rem !important; }
    .p-4 { padding: 1rem !important; }
    .p-6 { padding: 1.5rem !important; }
    
    .flex { display: flex !important; }
    .inline-flex { display: inline-flex !important; }
    .items-center { align-items: center !important; }
    .justify-center { justify-content: center !important; }
    .gap-2 { gap: 0.5rem !important; }
    .gap-3 { gap: 0.75rem !important; }
    .gap-4 { gap: 1rem !important; }
    """
) as demo:
    # Modern header with gradient background
    model_status = "✅ RMMM model loaded and ready" if RMMM_MODEL is not None else "❌ Error loading RMMM model"
    tokenizer_status = "✅ MIMIC tokenizer loaded" if MIMIC_TOKENIZER is not None else "⚠️ Using GPT-2 tokenizer as fallback"
    
    gr.HTML(f"""
    <div class="main-header">
        <h1>🩻 XRaySwinGen / RMMM</h1>
        <p>Automatic Medical Report Generation System with Real-Time Evaluation Metrics</p>
        <div class="status-badges">
            <div class="status-badge">
                <span>🤖</span>
                <span>{model_status}</span>
            </div>
            <div class="status-badge">
                <span>📝</span>
                <span>{tokenizer_status}</span>
            </div>
        </div>
    </div>
    """)
    
    # Instructions and Metrics section with improved cards
    with gr.Row():
        with gr.Column(scale=0.75, variant='compact', elem_classes=["controls-panel"]):
            gr.HTML("""
            <div class="card process-steps">
                <div class="card-header">
                    <div class="card-icon">📖</div>
                    <div>
                        <h3 class="card-title">How to Use the System</h3>
                        <p class="card-subtitle">Simplified 3-step process</p>
                    </div>
                </div>
                <div class="card-content">
                    <div class="step">
                        <div class="step-number">1</div>
                        <div class="step-content">
                            <h5>Select Image</h5>
                            <p>Click on any X-ray in the gallery</p>
                        </div>
                    </div>
                    <div class="step">
                        <div class="step-number">2</div>
                        <div class="step-content">
                            <h5>Automatic Processing</h5>
                            <p>The model will automatically create the report</p>
                        </div>
                    </div>
                    <div class="step">
                        <div class="step-number">3</div>
                        <div class="step-content">
                            <h5>Analyze Results</h5>
                            <p>Review the AI report and quality metrics</p>
                        </div>
                    </div>
                </div>
            </div>
            """)
        
        with gr.Column(scale=0.75, variant='compact', elem_classes=["controls-panel"]):
            # Main interface with optimized layout - Gallery + Controls side by side
            example_images = get_available_image_paths()
            
            gr.HTML("""
            <div class="gallery-card">
                <div class="gallery-header">
                    <h4 class="gallery-title">
                        <span>🩻</span>
                        <span>MIMIC-CXR Gallery</span>
                    </h4>
                    <p class="gallery-subtitle">
                        Click any image to automatically generate the medical report
                    </p>
                </div>
            </div>
            """)
            gallery = gr.Gallery(
                value=example_images,
                columns=3,
                height='400px',
                object_fit="contain",
                allow_preview=True,
                show_label=False,
                show_download_button=False,
                interactive=True,
                container=True,
                elem_classes=["gallery-card"]
            )
        
        with gr.Column(scale=0.75, variant='compact', elem_classes=["controls-panel"]):
            metrics_display = gr.HTML(
                value="""
                <div class="card">
                    <div class="card-header">
                        <div class="card-icon">📊</div>
                        <div>
                            <h3 class="card-title">Evaluation Metrics</h3>
                            <p class="card-subtitle">Comparison with reference reports</p>
                        </div>
                    </div>
                    <div class="card-content">
                        <p style="color: var(--gray-600); font-style: italic; text-align: center; margin: 2rem 0;">
                            Select an image to see real-time evaluation metrics
                        </p>
                    </div>
                </div>
                """,
                label="Evaluation Metrics"
            )
            # Status display
            status_display = gr.HTML(
                value="""
                <div class="card">
                    <div class="card-header">
                        <div class="card-icon">💡</div>
                        <div>
                            <h3 class="card-title">System Status</h3>
                            <p class="card-subtitle">Monitor processing</p>
                        </div>
                    </div>
                    <div class="card-content">
                        <p style="color: var(--gray-600); text-align: center; margin: 0; font-weight: 500;">
                            Select an image from the gallery to start
                        </p>
                    </div>
                </div>
                """,
                label="Status"
            )
    
    # Reports and Analysis Section - separate row
    with gr.Row():
        with gr.Column(scale=1):
            # AI Report
            ai_report = gr.Textbox(
                label="🤖 AI-Generated Report",
                lines=6,
                max_lines=20,
                placeholder="Select an image from the gallery above to see the AI analysis...",
                container=True,
                show_copy_button=True,
                elem_classes=["report-ai"]
            )

            # Ground Truth Report
            ground_truth = gr.Textbox(
                label="📋 Ground Truth Report",
                lines=8,
                max_lines=20,
                placeholder="The reference report will appear here when you select an image...",
                container=True,
                show_copy_button=True,
                elem_classes=["report-ground-truth"]
            )
        
        with gr.Column(scale=1):
            # Annotated X-ray (visual overlay of suggested findings)
            annotated_image = gr.Image(
                label="🩻 Annotated X-ray",
                interactive=False,
                visible=True
            )

            # Explain findings button and outputs
            explain_btn = gr.Button("🔍 Explain Findings", size="md", elem_classes=["btn-primary"]) 

            explanation_box = gr.Textbox(
                label="🧾 Explanation (Plain English)",
                lines=6,
                interactive=False,
                container=True,
                show_copy_button=True
            )

            detailed_findings = gr.Textbox(
                label="🧩 Detailed Findings",
                lines=8,
                interactive=False,
                container=True,
                show_copy_button=True
            )

            step_by_step_html = gr.HTML(value="", label="Step-by-step")

            # Download report helper - writes a local file and returns path via status
            download_report_btn = gr.Button("⬇️ Save Report")

            # Clear button (placed here so it's defined before wiring)
            clear_btn = gr.Button(
                "🗑️ Clear Results",
                size="lg",
                elem_classes=["secondary"]
            )
    
    # Hidden components for internal processing
    image_input = gr.Image(visible=False)  # Hidden image component for processing
    selected_image_state = gr.State(value="")  # Hidden state to store the selected image filename
    
    
    # Gallery click handler to load selected image and automatically generate report
    def load_and_generate_report(evt: gr.SelectData):
        selected_image_path = example_images[evt.index]
        print(f"Gallery selection - Loading image: {selected_image_path}")
        
        # Update status to show processing
        processing_status = """
        <div class="card">
            <div class="card-header">
                <div class="card-icon" style="background: var(--primary-50); color: var(--primary-600);">🔄</div>
                <div>
                    <h3 class="card-title">Processing...</h3>
                    <p class="card-subtitle">Generating medical report</p>
                </div>
            </div>
            <div class="card-content">
                <p style="color: var(--primary-600); text-align: center; margin: 0; font-weight: 500;">
                    Analyzing image and generating report...
                </p>
            </div>
        </div>
        """
        
        try:
            # Load the image explicitly to ensure it's properly loaded
            loaded_image = Image.open(selected_image_path).convert('RGB')
            print(f"✅ Successfully loaded image: {loaded_image.size}")
            
            # Generate reports immediately
            ai_report_text, ground_truth_text, metrics_html = inference_image_pipe_with_state(
                loaded_image, selected_image_path
            )

            # Create an annotated image (UX overlay) and an explanation summary
            annotated_img = annotate_image(loaded_image.copy(), ai_report_text)
            expl_text, detailed_text, step_html = explain_findings(ai_report_text, ground_truth_text)
            
            # Update status to show completion
            completion_status = f"""
            <div class="card">
                <div class="card-header">
                    <div class="card-icon" style="background: var(--success-50); color: var(--success-600);">✅</div>
                    <div>
                        <h3 class="card-title">Completed!</h3>
                        <p class="card-subtitle">Report generated successfully</p>
                    </div>
                </div>
                <div class="card-content">
                    <p style="color: var(--success-600); text-align: center; margin: 0; font-weight: 500; font-size: 0.875rem;">
                        {os.path.basename(selected_image_path)}
                    </p>
                </div>
            </div>
            """
            
            return (
                loaded_image,           # image_input (hidden)
                selected_image_path,    # selected_image_state
                ai_report_text,         # ai_report
                ground_truth_text,      # ground_truth
                annotated_img,          # annotated_image
                expl_text,              # explanation_box
                detailed_text,          # detailed_findings
                step_html,              # step_by_step_html
                metrics_html,           # metrics_display
                completion_status       # status_display
            )
            
        except Exception as e:
            print(f"❌ Error loading/processing image: {e}")
            error_status = f"""
            <div class="card">
                <div class="card-header">
                    <div class="card-icon" style="background: var(--error-50); color: var(--error-600);">❌</div>
                    <div>
                        <h3 class="card-title">Error</h3>
                        <p class="card-subtitle">Processing failed</p>
                    </div>
                </div>
                <div class="card-content">
                    <p style="color: var(--error-600); text-align: center; margin: 0; font-weight: 500; font-size: 0.875rem;">
                        {str(e)}
                    </p>
                </div>
            </div>
            """
            # On error, keep the annotated and explanation outputs empty
            return None, selected_image_path, "", "", None, "", "", "", "", error_status
    
    gallery.select(
        fn=load_and_generate_report,
        outputs=[
            image_input,
            selected_image_state,
            ai_report,
            ground_truth,
            annotated_image,
            explanation_box,
            detailed_findings,
            step_by_step_html,
            metrics_display,
            status_display,
        ]
    )

    # Wire Save Report button to write combined report to ./reports and update status
    download_report_btn.click(
        fn=lambda s, a, g, m: save_report_file(s, a, g, m),
        inputs=[selected_image_state, ai_report, ground_truth, metrics_display],
        outputs=[status_display]
    )

    # Wire Explain Findings button
    explain_btn.click(
        fn=lambda a, g: explain_findings(a, g),
        inputs=[ai_report, ground_truth],
        outputs=[explanation_box, detailed_findings, step_by_step_html]
    )
    
    # Clear button functionality — reset all outputs created by gallery.select
    def _clear_all():
        metrics_placeholder = """
        <div class="card">
            <div class="card-header">
                <div class="card-icon">📊</div>
                <div>
                    <h2 class="card-title">Evaluation Metrics</h2>
                    <p class="card-subtitle">Comparison with reference reports</p>
                </div>
            </div>
            <div class="card-content">
                <p style="color: var(--gray-600); font-style: italic; text-align: center; margin: 2rem 0;">
                    Select an image to see real-time evaluation metrics
                </p>
            </div>
        </div>
        """

        status_placeholder = """
        <div class="card">
            <div class="card-header">
                <div class="card-icon">💡</div>
                <div>
                    <h3 class="card-title">System Status</h3>
                    <p class="card-subtitle">Track the processing</p>
                </div>
            </div>
            <div class="card-content">
                <p style="color: var(--gray-600); text-align: center; margin: 0; font-weight: 500;">
                    Select an image from the gallery to start
                </p>
            </div>
        </div>
        """

        return (
            None,   # image_input
            "",    # selected_image_state
            "",    # ai_report
            "",    # ground_truth
            None,   # annotated_image
            "",    # explanation_box
            "",    # detailed_findings
            "",    # step_by_step_html
            metrics_placeholder,
            status_placeholder,
        )

    clear_btn.click(
        fn=_clear_all,
        inputs=[],
        outputs=[
            image_input,
            selected_image_state,
            ai_report,
            ground_truth,
            annotated_image,
            explanation_box,
            detailed_findings,
            step_by_step_html,
            metrics_display,
            status_display,
        ]
    )
    # (End of Blocks) Launch when executed directly
if __name__ == "__main__":
    # Fix for Windows asyncio connection issues
    if sys.platform.startswith('win'):
        try:
            asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
        except Exception:
            pass

    demo.launch(
        show_error=True,
        quiet=False,
        share=True
    )