app.py

#!/usr/bin/env python3
"""
BrainGemma3D Dashboard - Gradio Interface
========================================
Interactive dashboard for generating medical reports from 3D MRI volumes (NIfTI)
using the trained BrainGemma3D model.

Usage:
    python app.py --model-dir ../BrainGemma3D

    # Custom port
    python app.py --model-dir ../BrainGemma3D --port 7860
"""

import os
import sys
import spaces
import argparse
import time
import datetime
from pathlib import Path
from typing import Optional, Tuple
import tempfile
import shutil

import gradio as gr
import torch
import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.patches import Rectangle
from matplotlib.backends.backend_pdf import PdfPages
import io
from PIL import Image

# Importa i moduli del modello da BrainGemma3D
import importlib.util
from huggingface_hub import snapshot_download

def import_module_from_path(module_name, file_path):
    """Dynamically import a module from a file path."""
    if not os.path.exists(file_path):
        raise FileNotFoundError(
            f"Module file not found: {file_path}\n"
            f"Please ensure BrainGemma3D model is downloaded from HuggingFace."
        )
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module

def ensure_braingemma3d_downloaded():
    """
    Ensure BrainGemma3D model is available locally.
    Downloads from HuggingFace if not present.
    """
    # Check if BrainGemma3D exists in parent directory
    local_path = Path(__file__).parent.parent / "BrainGemma3D"
    
    # Check required files
    arch_file = local_path / "braingemma3d_architecture.py"
    interp_file = local_path / "braingemma3d_interpretability.py"
    config_file = local_path / "model_config.json"
    
    if arch_file.exists() and interp_file.exists() and config_file.exists():
        print(f"✅ BrainGemma3D found at: {local_path}")
        return local_path
    
    # If not found, try to download from HuggingFace
    print("\n" + "="*70)
    print("📥 DOWNLOADING BRAINGEMMA3D MODEL FROM HUGGINGFACE")
    print("="*70)
    print("Repository: praiselab-picuslab/BrainGemma3D")
    print("This may take a few minutes (model size ~3.5 GB)...")
    print("")
    
    try:
        # Create parent directory
        download_dir = Path(__file__).parent.parent / "BrainGemma3D"
        download_dir.mkdir(parents=True, exist_ok=True)
        
        # Download from HuggingFace Hub
        snapshot_path = snapshot_download(
            repo_id="praiselab-picuslab/BrainGemma3D",
            local_dir=str(download_dir),
            local_dir_use_symlinks=False,
        )
        
        print(f"✅ Model downloaded successfully to: {download_dir}")
        print("="*70 + "\n")
        
        return local_path
        
    except Exception as e:
        print(f"\n❌ ERROR downloading model from HuggingFace: {e}")
        print("\n📋 MANUAL DOWNLOAD INSTRUCTIONS:")
        print("   1. Install git-lfs: apt-get install git-lfs")
        print("   2. Clone the model:")
        print("      cd /path/to/CBMS")
        print("      git clone https://huggingface.co/praiselab-picuslab/BrainGemma3D")
        print("   3. Restart the dashboard")
        print("="*70 + "\n")
        raise

# Ensure model is downloaded before importing
braingemma_path = ensure_braingemma3d_downloaded()

# Import BrainGemma3D architecture
arch_module = import_module_from_path(
    "braingemma3d_architecture",
    str(braingemma_path / "braingemma3d_architecture.py")
)
interp_module = import_module_from_path(
    "braingemma3d_interpretability",
    str(braingemma_path / "braingemma3d_interpretability.py")
)

BrainGemma3D = arch_module.BrainGemma3D
load_nifti_volume = arch_module.load_nifti_volume
CANONICAL_PROMPT = arch_module.CANONICAL_PROMPT

# Import interpretability functions
load_full_model = interp_module.load_full_model
run_interpretability = interp_module.run_interpretability
set_seed = interp_module.set_seed

# ============================================================================
# GLOBAL MODEL (caricato una sola volta all'avvio)
# ============================================================================
MODEL: Optional[BrainGemma3D] = None
MODEL_CONFIG = {}
LOAD_NIFTI_FN = None  # Riferimento alla funzione load_nifti_volume dal modello


def load_model(model_dir: str):
    """Carica il modello BrainGemma3D da directory (HuggingFace o locale)"""
    global LOAD_NIFTI_FN
    
    print("\n" + "="*70)
    print("📥 LOADING BRAINGEMMA3D MODEL")
    print("="*70)
    print(f"Model directory: {model_dir}")
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Usa la funzione load_full_model da braingemma3d_interpretability
    model, load_nifti_fn, canonical_prompt = load_full_model(model_dir, device)
    
    # Salva riferimento alla funzione load_nifti_volume
    LOAD_NIFTI_FN = load_nifti_fn
    
    print(f"✅ Model loaded successfully on {model.lm_device}")
    print("="*70 + "\n")
    
    return model


def normalize_volume_for_display(volume: torch.Tensor) -> np.ndarray:
    """
    Normalise a 3-D volume for display.

    Accepts any shape: (1,1,D,H,W), (1,D,H,W), (D,H,W).
    Returns numpy array [D, H, W] in [0, 1].
    """
    vol = volume.detach().clone().cpu()
    # Squeeze ALL leading singleton dims (handles 5-D output of load_nifti_volume)
    while vol.ndim > 3 and vol.shape[0] == 1:
        vol = vol.squeeze(0)
    # If still >3D (multi-channel), take first channel
    while vol.ndim > 3:
        vol = vol[0]
    vol_np = vol.numpy()
    if vol_np.ndim != 3:
        raise ValueError(f"Expected 3-D after squeeze, got {vol_np.shape}")
    vol_np = (vol_np - vol_np.min()) / (vol_np.max() - vol_np.min() + 1e-8)
    return vol_np


def load_nifti_native_for_display(nifti_path: str) -> np.ndarray:
    """
    Load a NIfTI volume at its native resolution for display.

    Performs canonical reorientation + depth-axis detection (same logic as
    load_nifti_volume_general) but WITHOUT resizing.  Returns a numpy array
    [D, H, W] normalised to [0, 1].
    """
    img = nib.load(nifti_path)
    img_can = nib.as_closest_canonical(img)
    vol = img_can.get_fdata(dtype=np.float32)
    zooms = np.array(img_can.header.get_zooms()[:3], dtype=np.float32)
    shape = np.array(vol.shape, dtype=np.int64)

    # ── Depth-axis detection (mirrors load_nifti_volume_general) ──────────
    median_zoom = float(np.median(zooms))
    zoom_ratios = zooms / (median_zoom + 1e-6)
    thick_axes  = np.where(zoom_ratios >= 1.5)[0]
    min_dim_ax  = int(np.argmin(shape))

    ornt = nib.orientations.io_orientation(img_can.affine)
    header_depth_ax = None
    for arr_ax, (can_code, _flip) in enumerate(ornt):
        if int(can_code) == 2:
            header_depth_ax = arr_ax
            break

    if len(thick_axes) == 1:
        depth_ax = int(thick_axes[0])
    elif len(thick_axes) > 1:
        candidates = [a for a in thick_axes if shape[a] == shape[thick_axes].min()]
        depth_ax = int(candidates[0]) if candidates else int(thick_axes[0])
    else:
        depth_ax = min_dim_ax
        sorted_dims = np.sort(shape)
        if sorted_dims[0] == sorted_dims[1] and header_depth_ax is not None:
            depth_ax = header_depth_ax

    # ── Assign H and W axes (mirrors load_nifti_volume_general) ──────────
    spatial_axes = [0, 1, 2]
    spatial_axes.remove(depth_ax)
    code_of  = {int(ornt[a, 0]): a for a in spatial_axes}
    ax_height = code_of.get(1, spatial_axes[0])   # A axis → height
    ax_width  = code_of.get(0, spatial_axes[1])   # R axis → width
    if ax_height == ax_width:
        ax_height, ax_width = spatial_axes[0], spatial_axes[1]

    # ── Transpose to (D, H, W) ────────────────────────────────────────────
    vol = np.transpose(vol, (depth_ax, ax_height, ax_width))

    # ── Robust normalisation (1st–99th percentile) ────────────────────────
    p1, p99 = np.percentile(vol, 1), np.percentile(vol, 99)
    vol = np.clip(vol, p1, p99)
    vol = (vol - p1) / (p99 - p1 + 1e-8)

    print(f"[display] native shape after transpose: {vol.shape}")
    return vol.astype(np.float32)


def create_slice_image(vol_np: np.ndarray, plane: str, slice_idx: int) -> Image.Image:
    """
    Creates a single slice image from a 3D volume.

    Args:
        vol_np: Normalised volume [D, H, W]
        plane: 'axial', 'coronal', or 'sagittal'
        slice_idx: Index of the slice to display

    Returns:
        PIL Image of the slice (fixed aspect ratio, no stretching)
    """
    D, H, W = vol_np.shape
    OUTPUT_PX = 512   # target output size so all views are equally sized

    if plane == 'axial':
        slice_idx = max(0, min(D - 1, slice_idx))
        img_data = vol_np[slice_idx, :, :]   # H x W
        n_rows, n_cols = H, W
        title = f'Axial  \u00b7  slice {slice_idx + 1} / {D}'
    elif plane == 'coronal':
        slice_idx = max(0, min(H - 1, slice_idx))
        img_data = vol_np[:, slice_idx, :]   # D x W  (depth vertical)
        n_rows, n_cols = D, W
        title = f'Coronal  \u00b7  slice {slice_idx + 1} / {H}'
    elif plane == 'sagittal':
        slice_idx = max(0, min(W - 1, slice_idx))
        img_data = vol_np[:, :, slice_idx]   # D x H  (depth vertical)
        n_rows, n_cols = D, H
        title = f'Sagittal  \u00b7  slice {slice_idx + 1} / {W}'
    else:
        raise ValueError(f"Unknown plane: {plane}")

    # Pad depth dimension symmetrically so the slice is square (no stretching)
    if n_rows != n_cols:
        if n_rows < n_cols:
            pad = (n_cols - n_rows) // 2
            extra = (n_cols - n_rows) - 2 * pad
            img_data = np.pad(img_data, ((pad, pad + extra), (0, 0)), mode='constant')
        else:
            pad = (n_rows - n_cols) // 2
            extra = (n_rows - n_cols) - 2 * pad
            img_data = np.pad(img_data, ((0, 0), (pad, pad + extra)), mode='constant')
        n_rows, n_cols = img_data.shape

    # --- Render at true voxel aspect ratio, then pad to square OUTPUT_PX ---
    # figure size in inches at 100 dpi: width = n_cols/100, height = n_rows/100
    dpi = 100
    fig_w = max(n_cols / dpi, 2.0)
    fig_h = max(n_rows / dpi, 2.0) + 0.35   # +0.35 inch title headroom

    fig, ax = plt.subplots(1, 1, figsize=(fig_w, fig_h), dpi=dpi,
                           facecolor='#0d0d0d')
    ax.set_facecolor('#0d0d0d')

    # aspect='equal' preserves square voxels; figsize already accounts for shape
    ax.imshow(img_data, cmap='gray', vmin=0, vmax=1,
              origin='lower', aspect='equal',
              interpolation='lanczos')

    ax.set_title(title, fontsize=11, fontweight='bold',
                 color='white', pad=5,
                 fontfamily='DejaVu Sans')
    ax.axis('off')

    plt.tight_layout(pad=0.25)

    buf = io.BytesIO()
    plt.savefig(buf, format='png', dpi=dpi, bbox_inches='tight',
                facecolor=fig.get_facecolor())
    buf.seek(0)
    raw = Image.open(buf).copy()
    plt.close(fig)

    # Pad shorter axis to make a pure-black square so Gradio doesn't re-stretch
    rw, rh = raw.size
    side = max(rw, rh)
    padded = Image.new('RGB', (side, side), color=(0, 0, 0))
    padded.paste(raw, ((side - rw) // 2, (side - rh) // 2))
    return padded


# ============================================================================
# PDF RADIOLOGIST REPORT
# ============================================================================

def generate_pdf_fn(
    vol_np,
    report_text: str,
    patient_name: str,
    patient_id: str,
    exam_date: str,
    referring_physician: str,
    institution: str,
):
    """
    Generates a professional single-page A4 neuroradiology report.

    Layout (top → bottom):
      [header]       two-tone band: institution left, report type right
      [patient table] 3 × 2 grid with cell borders
      [MRI panel]    dark full-width band with three labelled views
      [findings]     framed box with justified text
      [impression]   light-amber notice box
      [signature]    two-column sign-off area
      [footer]       disclaimer band with timestamp
    """
    if vol_np is None or not report_text.strip():
        return None

    import textwrap as _tw

    checkpoint_name = Path(MODEL_CONFIG.get('checkpoint_dir', 'MedGemma3D')).name
    exam_date_str   = exam_date.strip() or datetime.date.today().strftime('%Y-%m-%d')
    generated_on    = datetime.datetime.now().strftime('%Y-%m-%d  %H:%M')
    report_id       = f"RPT-{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}"

    pdf_path = tempfile.mktemp(suffix='_brain_mri_report.pdf')

    # ── Color palette ────────────────────────────────────────────────────
    NAVY    = '#0d2137'   # deep navy
    TEAL    = '#1565a0'   # accent blue
    LTEAL   = '#1976d2'   # lighter accent
    STEEL   = '#dce8f4'   # very light blue
    LGRAY   = '#f5f7fa'   # panel bg
    WHITE   = '#ffffff'
    TEXTDK  = '#111122'   # body text
    TEXTMD  = '#445566'   # secondary text
    TEXTLT  = '#7788aa'   # field labels
    RULE    = '#a8c4dc'   # divider lines
    ORANGE  = '#b83200'   # disclaimer accent
    AMBER   = '#fff8e1'   # impression panel bg
    AMBERB  = '#f9a825'   # impression panel border
    DARKBG  = '#0a0a12'   # MRI panel bg

    # ── Helpers ──────────────────────────────────────────────────────────
    def _pad_sq(img):
        r, c = img.shape
        if r == c:
            return img
        if r < c:
            p, e = (c - r) // 2, (c - r) % 2
            return np.pad(img, ((p, p + e), (0, 0)), mode='constant')
        p, e = (r - c) // 2, (r - c) % 2
        return np.pad(img, ((0, 0), (p, p + e)), mode='constant')

    def _section_accent(fig, y, h, title):
        """Left-accent section header: vertical teal bar + bold label."""
        # Vertical accent bar
        bar = fig.add_axes([0.03, y, 0.005, h])
        bar.set_facecolor(TEAL)
        bar.axis('off')
        # Horizontal rule (full width, very thin)
        rule = fig.add_axes([0.038, y, 0.932, 0.0012])
        rule.set_facecolor(RULE)
        rule.axis('off')
        # Title text
        fig.text(0.042, y + h * 0.15, title,
                 fontsize=7.2, fontweight='bold', color=TEAL,
                 fontfamily='DejaVu Sans', transform=fig.transFigure,
                 va='bottom')

    def _bounded_box(fig, left, bottom, width, height,
                     fc=LGRAY, ec=RULE, lw=0.7):
        """Draw a filled rectangle directly on fig coordinates."""
        box_ax = fig.add_axes([left, bottom, width, height])
        box_ax.set_facecolor(fc)
        box_ax.add_patch(Rectangle((0, 0), 1, 1,
                                    linewidth=lw, edgecolor=ec,
                                    facecolor=fc, transform=box_ax.transAxes))
        box_ax.axis('off')
        return box_ax

    def _render_justified(ax, text, fontsize=7.6, color='#111122',
                           linespacing=1.55, x0=0.016):
        """Render fully-justified text using per-glyph width measurement."""
        import matplotlib.font_manager as _fm

        prop = _fm.FontProperties(family='DejaVu Sans', size=fontsize)
        fig_w_pts  = fig.get_figwidth() * 72
        fig_h_pts  = fig.get_figheight() * 72
        ax_pos     = ax.get_position()
        ax_w_pts   = fig_w_pts * ax_pos.width
        ax_h_pts   = fig_h_pts * ax_pos.height
        usable_pts = ax_w_pts * (1.0 - x0 - 0.008)

        # PIL truetype `size` is in PIXELS at PIL's default 96 DPI.
        # We need widths in typographic POINTS (1/72 inch) to match usable_pts.
        # Conversion: pil_size_px = fontsize_pts * 96/72
        #             width_pts   = width_px    * 72/96
        _PIL_DPI   = 96.0
        _PTS2PX    = _PIL_DPI / 72.0          # multiply pts → PIL pixels
        _PX2PTS    = 72.0 / _PIL_DPI          # multiply PIL pixels → pts
        _pil_size  = max(1, int(fontsize * _PTS2PX))   # font size in px for PIL

        from PIL import ImageFont as _ImageFont
        _font_path = _fm.findfont(prop)
        try:
            _pil_font = _ImageFont.truetype(_font_path, size=_pil_size)
        except Exception:
            _pil_font = None

        def _ww(w):
            """Return width of string w in typographic POINTS."""
            if not w:
                return 0.0
            if _pil_font is not None:
                try:
                    bbox = _pil_font.getbbox(w)   # pixels
                    return float(bbox[2] - bbox[0]) * _PX2PTS
                except AttributeError:
                    try:
                        return float(_pil_font.getlength(w)) * _PX2PTS
                    except Exception:
                        pass
            # Fallback: ~0.55 em per character
            return len(w) * fontsize * 0.55

        space_w = _ww(' ')
        words   = text.split()

        # ── Word-wrap ────────────────────────────────────────────────────
        lines = []
        cur_words, cur_w = [], 0.0
        for word in words:
            ww = _ww(word)
            needed = ww + (space_w if cur_words else 0.0)
            if cur_words and cur_w + needed > usable_pts:
                lines.append(cur_words)
                cur_words, cur_w = [word], ww
            else:
                cur_words.append(word)
                cur_w += needed
        if cur_words:
            lines.append(cur_words)

        # ── Render lines top-down ────────────────────────────────────────
        n_lines    = len(lines)
        line_h_ax  = fontsize * linespacing / ax_h_pts   # axes-fraction per line
        y_cursor   = 0.97

        for i, lw_list in enumerate(lines):
            is_last = (i == n_lines - 1)
            gaps    = len(lw_list) - 1

            if is_last or gaps == 0:
                ax.text(x0, y_cursor, ' '.join(lw_list),
                        fontsize=fontsize, color=color,
                        fontfamily='DejaVu Sans', va='top',
                        transform=ax.transAxes)
            else:
                total_word_w = sum(_ww(w) for w in lw_list)
                extra_each   = (usable_pts - total_word_w) / gaps
                x_cur_pts    = x0 * ax_w_pts
                for j, word in enumerate(lw_list):
                    ax.text(x_cur_pts / ax_w_pts, y_cursor, word,
                            fontsize=fontsize, color=color,
                            fontfamily='DejaVu Sans', va='top',
                            transform=ax.transAxes)
                    x_cur_pts += _ww(word) + extra_each

            y_cursor -= line_h_ax

    # ═════════════════════════════════════════════════════════════════════
    with PdfPages(pdf_path) as pdf:
        fig = plt.figure(figsize=(8.27, 11.69), facecolor=WHITE)
        fig.patch.set_linewidth(0)

        # ── 1. HEADER BAND  y = 0.915 → 1.00 ────────────────────────────
        # Deep-navy left block (65 %) + dark-teal right block (35 %)
        hdr_l = fig.add_axes([0.0, 0.915, 0.65, 0.085])
        hdr_l.set_facecolor(NAVY)
        hdr_l.axis('off')

        hdr_r = fig.add_axes([0.65, 0.915, 0.35, 0.085])
        hdr_r.set_facecolor(TEAL)
        hdr_r.axis('off')

        # Thin teal accent stripe at very left edge
        acc = fig.add_axes([0.0, 0.915, 0.012, 0.085])
        acc.set_facecolor(LTEAL)
        acc.axis('off')

        hdr_l.text(0.04, 0.72, (institution or 'Radiology Department').upper(),
                   fontsize=11.5, fontweight='bold', color=WHITE,
                   fontfamily='DejaVu Sans', va='center')
        hdr_l.text(0.04, 0.28, 'NEURORADIOLOGY  ·  BRAIN MRI  ·  DIAGNOSTIC REPORT',
                   fontsize=7.2, color='#a8cce8',
                   fontfamily='DejaVu Sans', va='center', style='italic')

        hdr_r.text(0.96, 0.70, exam_date_str,
                   fontsize=10, fontweight='bold', color=WHITE,
                   fontfamily='DejaVu Sans', ha='right', va='center')
        hdr_r.text(0.96, 0.28, f'Report ID: {report_id}',
                   fontsize=6.5, color='#c8dff0',
                   fontfamily='DejaVu Sans', ha='right', va='center')

        # Thin white separator line between left and right header blocks
        _sep = fig.add_axes([0.648, 0.915, 0.001, 0.085])
        _sep.set_facecolor('#ffffff40')
        _sep.axis('off')

        # ── 2. PATIENT TABLE  y = 0.825 → 0.910 ─────────────────────────
        # Outer frame
        _bounded_box(fig, 0.03, 0.825, 0.94, 0.085, fc=LGRAY, ec=RULE, lw=1.0)

        # Row 1 cells (top row — slightly darker bg)
        cell_w  = 0.94 / 3
        cell_h1 = 0.043
        cell_h2 = 0.040
        cell_y1 = 0.825 + cell_h2          # top row bottom
        cell_y2 = 0.825                    # bottom row bottom
        _lkw = dict(fontsize=5.8, color=TEXTLT, fontfamily='DejaVu Sans')
        _vkw = dict(fontsize=8.2, color=TEXTDK, fontweight='bold',
                    fontfamily='DejaVu Sans')
        _vkws= dict(fontsize=6.5, color=TEXTDK, fontweight='bold',
                    fontfamily='DejaVu Sans')

        table_data = [
            ('PATIENT NAME',        patient_name or '—',    _vkw),
            ('PATIENT ID / MRN',    patient_id   or '—',    _vkw),
            ('DATE OF EXAM',        exam_date_str,          _vkw),
            ('REFERRING PHYSICIAN', referring_physician or '—', _vkw),
            ('MODALITY',            'Brain MRI 3D',         _vkw),
            ('INSTITUTION',         institution or '—',     _vkw),
        ]
        for col in range(3):
            lx = 0.03 + col * cell_w
            # top row cell bg
            _bounded_box(fig, lx, cell_y1, cell_w, cell_h1,
                         fc=STEEL, ec=RULE, lw=0.5)
            # bottom row cell bg
            _bounded_box(fig, lx, cell_y2, cell_w, cell_h2,
                         fc=LGRAY, ec=RULE, lw=0.5)
            # content: top row
            lbl, val, kw = table_data[col]
            fig.text(lx + 0.010, cell_y1 + cell_h1 - 0.005, lbl,
                     **_lkw, transform=fig.transFigure, va='top')
            fig.text(lx + 0.010, cell_y1 + 0.005, val,
                     **kw, transform=fig.transFigure, va='bottom')
            # content: bottom row
            lbl2, val2, kw2 = table_data[col + 3]
            fig.text(lx + 0.010, cell_y2 + cell_h2 - 0.004, lbl2,
                     **_lkw, transform=fig.transFigure, va='top')
            fig.text(lx + 0.010, cell_y2 + 0.004, val2,
                     **kw2, transform=fig.transFigure, va='bottom')

        # ── 3. MRI VIEWS PANEL  y = 0.540 → 0.820 ───────────────────────
        mri_bottom = 0.540
        mri_height = 0.280

        # Outer dark background
        mri_bg = fig.add_axes([0.03, mri_bottom, 0.94, mri_height])
        mri_bg.set_facecolor(DARKBG)
        mri_bg.add_patch(Rectangle((0, 0), 1, 1, linewidth=1.2,
                                    edgecolor=TEAL, facecolor=DARKBG))
        mri_bg.axis('off')

        # Header strip inside MRI panel
        mri_bg.add_patch(Rectangle((0, 0.895), 1, 0.105,
                                    facecolor=TEAL, edgecolor='none'))
        mri_bg.text(0.5, 0.947, 'MULTI-PLANAR RECONSTRUCTION  ·  MID-VOLUME SLICES',
                    ha='center', va='center', fontsize=7, fontweight='bold',
                    color=WHITE, fontfamily='DejaVu Sans')

        D, H, W = vol_np.shape
        mid_d, mid_h, mid_w = D // 2, H // 2, W // 2

        panel_specs = [
            ('AXIAL',    _pad_sq(vol_np[mid_d, :, :]), f'z = {mid_d + 1} / {D}'),
            ('CORONAL',  _pad_sq(vol_np[:, mid_h, :]), f'y = {mid_h + 1} / {H}'),
            ('SAGITTAL', _pad_sq(vol_np[:, :, mid_w]), f'x = {mid_w + 1} / {W}'),
        ]
        outer_pad  = 0.015   # fraction inside mri_bg axes
        inner_gap  = 0.012
        p_w = (1.0 - 2 * outer_pad - 2 * inner_gap) / 3
        img_top_margin  = 0.105   # below title strip

        for i, (lbl, img_data, coord) in enumerate(panel_specs):
            x0 = outer_pad + i * (p_w + inner_gap)
            # Image axes (inside mri_bg)
            img_ax = mri_bg.inset_axes([x0, img_top_margin + 0.04,
                                         p_w, 1.0 - img_top_margin - 0.085 - 0.04])
            img_ax.set_facecolor(DARKBG)
            img_ax.imshow(img_data, cmap='gray', vmin=0, vmax=1,
                          origin='lower', aspect='equal', interpolation='lanczos')
            img_ax.axis('off')
            # Label below image
            mri_bg.text(x0 + p_w / 2,
                        img_top_margin + 0.022,
                        f'{lbl}   {coord}',
                        ha='center', va='center',
                        fontsize=6.2, color='#99bbdd',
                        fontfamily='DejaVu Sans', fontweight='bold')

        # ── 4. RADIOLOGICAL FINDINGS  y = 0.268 → 0.535 ─────────────────
        find_bottom = 0.268
        find_top    = 0.535
        find_height = find_top - find_bottom

        _section_accent(fig, find_top - 0.016, 0.016, 'RADIOLOGICAL FINDINGS')

        # Framed text box
        find_box = _bounded_box(fig, 0.03, find_bottom,
                                 0.94, find_height - 0.022,
                                 fc=LGRAY, ec=RULE, lw=0.8)

        # Left accent stripe inside box
        find_box.add_patch(Rectangle((0, 0), 0.008, 1.0,
                                      facecolor=LTEAL, edgecolor='none'))

        # Pixel-accurate justified text using per-glyph TextPath measurement
        clean_text = report_text.strip().replace('\n', ' ')
        _render_justified(find_box, clean_text,
                          fontsize=7.6, color=TEXTDK, linespacing=1.55, x0=0.016)

        # ── 5. IMPRESSION  y = 0.190 → 0.263 ────────────────────────────
        imp_bottom = 0.190
        imp_top    = 0.263
        imp_height = imp_top - imp_bottom

        _section_accent(fig, imp_top - 0.016, 0.016, 'IMPRESSION  &  CLINICAL CORRELATION')

        imp_box = _bounded_box(fig, 0.03, imp_bottom,
                                0.94, imp_height - 0.022,
                                fc=AMBER, ec=AMBERB, lw=0.9)
        _imp_text = (
            '\u26a0  This report was produced by an AI system (BrainGemma3D) for research '
            'purposes only.  Clinical correlation with patient history and review by a '
            'qualified radiologist is mandatory before any clinical use.  The findings '
            'above must NOT be used for diagnosis or treatment.'
        )
        _render_justified(imp_box, _imp_text,
                          fontsize=6.8, color='#5a3000', linespacing=1.6, x0=0.015)

        # ── 6. REPORTING RADIOLOGIST  y = 0.103 → 0.185 ─────────────────
        sig_bottom = 0.103
        sig_top    = 0.185

        _section_accent(fig, sig_top - 0.016, 0.016, 'REPORTING RADIOLOGIST')

        sig_area = _bounded_box(fig, 0.03, sig_bottom,
                                 0.94, sig_top - sig_bottom - 0.022,
                                 fc=WHITE, ec=RULE, lw=0.7)

        # Left column: name + signature line
        sig_area.text(0.01, 0.88, 'NAME  /  QUALIFICATION',
                      fontsize=5.8, color=TEXTLT, fontfamily='DejaVu Sans',
                      va='top', transform=sig_area.transAxes)
        sig_area.add_patch(Rectangle((0.01, 0.42), 0.43, 0.004,
                                      facecolor='#aaaaaa', edgecolor='none'))
        sig_area.text(0.01, 0.35, 'SIGNATURE  &  STAMP',
                      fontsize=5.8, color=TEXTLT, fontfamily='DejaVu Sans',
                      va='top', transform=sig_area.transAxes)
        sig_area.add_patch(Rectangle((0.01, 0.05), 0.43, 0.004,
                                      facecolor='#aaaaaa', edgecolor='none'))

        # Right column: date of report + countersign
        sig_area.text(0.55, 0.88, 'DATE OF REPORT',
                      fontsize=5.8, color=TEXTLT, fontfamily='DejaVu Sans',
                      va='top', transform=sig_area.transAxes)
        sig_area.text(0.55, 0.60, exam_date_str,
                      fontsize=9, fontweight='bold', color=TEXTDK,
                      fontfamily='DejaVu Sans', va='top', transform=sig_area.transAxes)
        sig_area.text(0.55, 0.35, 'COUNTERSIGNATURE',
                      fontsize=5.8, color=TEXTLT, fontfamily='DejaVu Sans',
                      va='top', transform=sig_area.transAxes)
        sig_area.add_patch(Rectangle((0.55, 0.05), 0.44, 0.004,
                                      facecolor='#aaaaaa', edgecolor='none'))

        # ── 7. DISCLAIMER FOOTER  y = 0.00 → 0.097 ──────────────────────
        foot = fig.add_axes([0, 0, 1, 0.097])
        foot.set_facecolor('#0d2137')
        # Thin ORANGE rule at top of footer
        foot.add_patch(Rectangle((0, 0.90), 1, 0.10,
                                   facecolor=ORANGE, edgecolor='none'))
        foot.axis('off')

        foot.text(0.5, 0.72,
                  'NOT FOR CLINICAL USE  ·  RESEARCH PROTOTYPE  ·  AWAITING RADIOLOGIST VALIDATION',
                  ha='center', va='center', fontsize=6.5, fontweight='bold',
                  color='#ffccaa', fontfamily='DejaVu Sans')
        foot.text(0.5, 0.38,
                  'This document was generated automatically by BrainGemma3D.  '
                  'It has not been validated by a licensed physician and must not be used for '
                  'clinical decision-making.',
                  ha='center', va='center', fontsize=6.2, color='#8899bb',
                  fontfamily='DejaVu Sans', linespacing=1.4)
        foot.text(0.03, 0.10, f'Generated: {generated_on}',
                  fontsize=5.8, color='#556677', fontfamily='DejaVu Sans', va='bottom')
        foot.text(0.97, 0.10, 'BrainGemma3D  ·  Page 1 of 1',
                  ha='right', fontsize=5.8, color='#556677',
                  fontfamily='DejaVu Sans', va='bottom')

        pdf.savefig(fig, dpi=200, bbox_inches=None)
        plt.close(fig)

    return pdf_path


@spaces.GPU
def generate_interpretability_grid(
    vol_np: np.ndarray,
    report: str,
    nifti_path: str,
    lime_samples: int = 50,
    n_segments: int = 20,
    alpha: float = 0.5,
    seed: int = 42,
) -> Image.Image:
    """
    Genera il plot 2x3 grid per l'interpretabilità LIME.
    
    Args:
        vol_np: Volume numpy (D, H, W) normalizzato [0, 1]
        report: Report generato dal modello
        nifti_path: Path al file NIfTI originale
        lime_samples: Numero di campioni LIME
        n_segments: Numero di supervoxel
        alpha: Trasparenza overlay
        seed: Random seed
    
    Returns:
        PIL Image del plot 2x3 grid
    """
    from scipy.ndimage import binary_erosion
    
    print(f"\n🔬 Running LIME interpretability analysis...")
    print(f"   Samples: {lime_samples}, Supervoxels: {n_segments}")
    
    # Set seed
    set_seed(seed)
    
    # 1. Create supervoxels
    segments, brain_mask = interp_module.big_supervoxels_brain_only(
        vol_np, n_segments=n_segments
    )
    
    # 2. Prepare LIME  
    from lime import lime_image
    segmentation_fn = interp_module.make_segmentation_fn(segments)
    explainer = lime_image.LimeImageExplainer()
    
    # 3. Prediction function
    def predict_fn(vols_4d):
        """vols_4d: (n_samples, D, H, W)"""
        vols_5d = vols_4d[:, np.newaxis, :, :, :]  # Add channel dim
        scores = interp_module.lime_score_report_nll(
            vols_5d,
            MODEL,
            prompt=CANONICAL_PROMPT,
            report_ref=report,
            batch_size=1,
        )
        return scores
    
    # 4. Run LIME
    explanation = explainer.explain_instance(
        vol_np,
        predict_fn,
        top_labels=1,
        hide_color=0.0,
        num_samples=lime_samples,
        segmentation_fn=segmentation_fn,
    )
    
    # 5. Get weights
    label = explanation.top_labels[0]
    weights = dict(explanation.local_exp[label])
    
    print(f"✅ LIME completed! Weights range: [{min(weights.values()):.4f}, {max(weights.values()):.4f}]")
    
    # 6. Create 2x3 grid figure
    D = vol_np.shape[0]
    # Select 3 representative slices
    lo = int(0.30 * D)
    hi = int(0.70 * D)
    selected_slices = np.linspace(lo, hi, 3, dtype=int).tolist()
    
    n_slices = len(selected_slices)
    fig, axes = plt.subplots(2, n_slices, figsize=(n_slices * 4, 2 * 4), facecolor='white')
    
    for col, slice_idx in enumerate(selected_slices):
        # Extract axial slice
        img_slice = vol_np[slice_idx, :, :]
        seg_slice = segments[slice_idx, :, :]
        
        # Row 0: Original
        axes[0, col].imshow(img_slice, cmap='gray', origin='lower', interpolation='bilinear')
        axes[0, col].set_title(f'Slice {slice_idx}', fontsize=12, fontweight='bold')
        axes[0, col].axis('off')
        
        # Row 1: LIME Overlay
        axes[1, col].imshow(img_slice, cmap='gray', origin='lower', interpolation='bilinear')
        
        # Create overlay using the helper function from interpretability
        overlay = np.zeros((*seg_slice.shape, 4), dtype=np.float32)
        all_weights = [float(v) for k, v in weights.items() if int(k) != 0]
        if all_weights:
            max_abs_weight = max(abs(w) for w in all_weights)
            if max_abs_weight > 1e-8:
                for seg_id_str, weight in weights.items():
                    seg_id = int(seg_id_str)
                    if seg_id == 0:
                        continue
                    mask = (seg_slice == seg_id)
                    if not mask.any():
                        continue
                    norm_weight = weight / max_abs_weight
                    edge_mask = mask & (~binary_erosion(mask))
                    
                    if weight > 0:  # Red
                        overlay[mask, 0] = 1.0
                        overlay[mask, 3] = alpha * abs(norm_weight)
                        overlay[edge_mask, 3] = min(1.0, alpha * abs(norm_weight) * 2.0)
                    else:  # Blue
                        overlay[mask, 1] = 0.4
                        overlay[mask, 2] = 1.0
                        overlay[mask, 3] = alpha * abs(norm_weight)
                        overlay[edge_mask, 3] = min(1.0, alpha * abs(norm_weight) * 2.0)
        
        axes[1, col].imshow(overlay, origin='lower', interpolation='nearest')
        axes[1, col].axis('off')
    
    # Add row labels
    axes[0, 0].text(-0.15, 0.5, 'Original', transform=axes[0, 0].transAxes,
                    fontsize=14, fontweight='bold', va='center', rotation=90)
    axes[1, 0].text(-0.15, 0.5, 'LIME Overlay', transform=axes[1, 0].transAxes,
                    fontsize=14, fontweight='bold', va='center', rotation=90)
    
    plt.tight_layout()
    
    # Convert to PIL Image
    buf = io.BytesIO()
    plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
    buf.seek(0)
    img = Image.open(buf).copy()
    plt.close(fig)
    
    print(f"✅ Interpretability plot generated (slices: {selected_slices})")
    
    return img


# ============================================================================
# GRADIO INTERFACE
# ============================================================================

@spaces.GPU
def generate_report_fn(
    nifti_file,
    custom_prompt: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
    vol_np_preloaded=None,
    progress=gr.Progress()
):
    """
    Main function: generate a diagnostic report from a NIfTI volume.

    Returns:
        Tuple: (formatted_report, raw_report, volume_normalized,
                axial_slider, coronal_slider, sagittal_slider,
                axial_img, coronal_img, sagittal_img, interpretability_img)
    """
    global MODEL
    
    empty_img = Image.new('RGB', (400, 400), color='gray')
    
    if MODEL is None:
        return "❌ ERROR: Model not loaded!", "", None, 0, 0, 0, empty_img, empty_img, empty_img
    
    if nifti_file is None:
        return "⚠️  Please upload a NIfTI file (.nii or .nii.gz)", "", None, 0, 0, 0, empty_img, empty_img, empty_img
    
    start_time = time.time()
    
    try:
        # Step 1: Load volume
        progress(0.1, desc="📂 Step 1/7: Loading NIfTI file...")

        # Resolve the path Gradio gave us
        nifti_path = nifti_file if isinstance(nifti_file, str) else (
            nifti_file.name if hasattr(nifti_file, 'name') else str(nifti_file)
        )
        nifti_path = str(nifti_path).strip()

        # Server-side extension validation (Gradio can't check double extensions)
        _basename = os.path.basename(nifti_path)
        if not (_basename.endswith('.nii') or _basename.endswith('.nii.gz')):
            return (
                "❌ Invalid file format. Please upload a **.nii** or **.nii.gz** file.",
                "", None, 0, 0, 0, empty_img, empty_img, empty_img
            )

        # --- DEBUG: show exactly what Gradio delivered ---
        import gzip as _gzip
        _exists = os.path.exists(nifti_path)
        _size   = os.path.getsize(nifti_path) if _exists else -1
        with open(nifti_path, 'rb') as _fh:
            _magic4 = _fh.read(4).hex()
        print(f"📂 Gradio path  : {nifti_path}")
        print(f"   exists={_exists}  size={_size}  magic={_magic4}")

        # nibabel identifies the reader by file extension.
        # If Gradio stripped/changed the extension, add the correct one.
        target_size = tuple(MODEL_CONFIG.get("target_size", [64, 128, 128]))
        _tmp_copy = None

        if nifti_path.endswith('.nii') or nifti_path.endswith('.nii.gz'):
            # Extension already correct — pass straight to nibabel
            load_path = nifti_path
        else:
            # Detect format from magic bytes and create a correctly-named symlink/copy
            _is_gz = (_magic4[:4] == '1f8b')  # gzip magic = 1f 8b
            _ext   = '.nii.gz' if _is_gz else '.nii'
            _tmp_copy = nifti_path + _ext
            shutil.copy2(nifti_path, _tmp_copy)
            load_path = _tmp_copy
            print(f"   → copied to: {load_path}")

        try:
            volume = load_nifti_volume(load_path, target_size=target_size)
            # Reuse the vol_np already loaded at upload time (avoids a redundant NIfTI read)
            if vol_np_preloaded is not None:
                vol_np = vol_np_preloaded
                print("[display] reusing preloaded vol_np from upload preview")
            else:
                vol_np = load_nifti_native_for_display(load_path)
        finally:
            if _tmp_copy and os.path.exists(_tmp_copy):
                try:
                    os.remove(_tmp_copy)
                except Exception:
                    pass
        
        # Debug: stampa la forma del volume
        print(f"📊 Model volume: {volume.shape}  |  Display volume: {vol_np.shape}")
        
        progress(0.25, desc="📊 Step 2/7: Normalizing volume for visualization...")
        D, H, W = vol_np.shape
        
        # Crea slice iniziali centrali
        mid_d, mid_h, mid_w = D // 2, H // 2, W // 2
        axial_img = create_slice_image(vol_np, 'axial', mid_d)
        coronal_img = create_slice_image(vol_np, 'coronal', mid_h)
        sagittal_img = create_slice_image(vol_np, 'sagittal', mid_w)
        
        # Step 3: Prepare prompt
        progress(0.4, desc="📝 Step 3/7: Preparing prompt...")
        # LOGICA NUOVA: Se custom_prompt è vuoto → solo CANONICAL
        #                Se custom_prompt non è vuoto → CANONICAL + custom (trasparente)
        if custom_prompt.strip():
            prompt = CANONICAL_PROMPT + "\n" + custom_prompt.strip()
            prompt_info = f"Canonical + Custom: \n'{prompt}'"
        else:
            prompt = CANONICAL_PROMPT
            prompt_info = "Canonical prompt ONLY"
        
        print(f"\n{'='*60}")
        print(f"🔮 GENERATING REPORT")
        print(f"{'='*60}")
        print(f"Prompt: {prompt_info}")
        print(f"Params: max_tokens={max_tokens}, temp={temperature}, top_p={top_p}")
        print(f"        rep_penalty={repetition_penalty}")
        
        # Step 4: Generate
        progress(0.5, desc="🤖 Step 4/7: Running BrainGemma3D inference...")
        
        # Prepara il volume per il modello (deve avere la forma corretta)
        # Il modello si aspetta il volume come restituito da load_nifti_volume
        # Assicurati che abbia la batch dimension se necessario
        model_volume = volume
        if model_volume.ndim == 3:
            # Aggiungi dimensione batch: [D, H, W] -> [1, D, H, W]
            model_volume = model_volume.unsqueeze(0)
        elif model_volume.ndim == 4 and model_volume.shape[0] == 1:
            # È già nella forma corretta [1, D, H, W] o [C=1, D, H, W]
            pass
        
        print(f"🤖 Volume shape for model.generate_report: {model_volume.shape}")
        
        with torch.no_grad():
            report = MODEL.generate_report(
                model_volume,
                prompt=prompt,
                max_new_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                no_repeat_ngram_size=3,
            )
        
        progress(0.8, desc="✨ Step 5/7: Formatting report...")
        
        # Format output
        elapsed = time.time() - start_time
        
        formatted_report = f"""<div>
<details>
<summary><strong>📊 Generation Info</strong> &nbsp;<small style="opacity:0.7;font-weight:normal">(click to expand)</small></summary>

<div>

| Field | Value |
|---|---|
| **Generation Time** | {elapsed:.2f} s |
| **Model** | `{Path(MODEL_CONFIG['checkpoint_dir']).name}` |
| **Prompt Mode** | {'Canonical + Custom Instructions' if custom_prompt.strip() else 'Canonical Only'} |
| **Max Tokens** | {max_tokens} |
| **Temperature** | {temperature} |
| **Top-p** | {top_p} |
| **Repetition Penalty** | {repetition_penalty} |

</div>
</details>

---

<h3>📋 Generated Report</h3>

<div style="line-height:1.6;">

{report}

</div>

---

<p style="color:#dc2626;">⚠️ <strong>Disclaimer:</strong> This report is generated by an AI model for research purposes only.<br>
Not intended for clinical diagnosis or treatment decisions.</p>

</div>""".strip()
        
        # Step 6: Generate Interpretability
        progress(0.85, desc="🔬 Step 6/7: Running LIME interpretability...")
        print(f"\n{'='*60}")
        print("🔬 GENERATING INTERPRETABILITY")
        print(f"{'='*60}")
        
        try:
            interp_img = generate_interpretability_grid(
                vol_np=vol_np,
                report=report,
                nifti_path=load_path,
                lime_samples=15,  # Reduced for faster dashboard
                n_segments=20,
                alpha=0.5,
                seed=42,
            )
            print("✅ Interpretability generated successfully")
        except Exception as interp_error:
            print(f"⚠️  Interpretability failed: {interp_error}")
            import traceback
            traceback.print_exc()
            # Create empty image if interpretability fails
            interp_img = Image.new('RGB', (800, 600), color='gray')
            fig, ax = plt.subplots(figsize=(8, 6))
            ax.text(0.5, 0.5, f'⚠️ Interpretability failed:\n{str(interp_error)}',
                    ha='center', va='center', fontsize=12, color='red',
                    transform=ax.transAxes)
            ax.axis('off')
            buf = io.BytesIO()
            plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
            buf.seek(0)
            interp_img = Image.open(buf).copy()
            plt.close()
        
        progress(1.0, desc=f"✅ Step 7/7: Complete! ({time.time() - start_time:.1f}s)")
        
        print(f"✅ Report + Interpretability generated in {time.time() - start_time:.2f}s")
        print(f"{'='*60}\n")
        
        return formatted_report, report, vol_np, mid_d, mid_h, mid_w, axial_img, coronal_img, sagittal_img, interp_img
        
    except Exception as e:
        import traceback
        error_msg = f"❌ ERROR during generation:\n\n{str(e)}\n\n{traceback.format_exc()}"
        print(error_msg)
        return error_msg, "", None, 0, 0, 0, empty_img, empty_img, empty_img, empty_img


def preview_nifti_fn(nifti_file):
    """
    Load and display NIfTI slices immediately after upload, without running inference.
    Returns (status_msg, vol_np, mid_d, mid_h, mid_w, axial_img, coronal_img, sagittal_img).
    """
    empty_img = Image.new('RGB', (400, 400), color='#0d0d0d')
    if nifti_file is None:
        return gr.update(), None, 0, 0, 0, empty_img, empty_img, empty_img

    nifti_path = nifti_file if isinstance(nifti_file, str) else (
        nifti_file.name if hasattr(nifti_file, 'name') else str(nifti_file)
    )
    nifti_path = str(nifti_path).strip()

    _basename = os.path.basename(nifti_path)
    if not (_basename.endswith('.nii') or _basename.endswith('.nii.gz')):
        return (
            "❌ Invalid file format. Please upload a **.nii** or **.nii.gz** file.",
            None, 0, 0, 0, empty_img, empty_img, empty_img
        )

    _tmp_copy = None
    try:
        with open(nifti_path, 'rb') as _fh:
            _magic4 = _fh.read(4).hex()

        if nifti_path.endswith('.nii') or nifti_path.endswith('.nii.gz'):
            load_path = nifti_path
        else:
            _is_gz  = (_magic4[:4] == '1f8b')
            _ext    = '.nii.gz' if _is_gz else '.nii'
            _tmp_copy = nifti_path + _ext
            shutil.copy2(nifti_path, _tmp_copy)
            load_path = _tmp_copy

        vol_np = load_nifti_native_for_display(load_path)
        D, H, W = vol_np.shape
        mid_d, mid_h, mid_w = D // 2, H // 2, W // 2

        axial_img    = create_slice_image(vol_np, 'axial',    mid_d)
        coronal_img  = create_slice_image(vol_np, 'coronal',  mid_h)
        sagittal_img = create_slice_image(vol_np, 'sagittal', mid_w)

        print(f"[preview] {os.path.basename(nifti_path)}  →  shape {vol_np.shape}")
        return gr.update(), vol_np, mid_d, mid_h, mid_w, axial_img, coronal_img, sagittal_img

    except Exception as e:
        import traceback
        tb = traceback.format_exc()
        print(f"[preview] Failed to load NIfTI: {e}\n{tb}")
        err = f"❌ ERROR loading NIfTI file:\n\n{str(e)}\n\n{tb}"
        return err, None, 0, 0, 0, empty_img, empty_img, empty_img
    finally:
        if _tmp_copy and os.path.exists(_tmp_copy):
            try:
                os.remove(_tmp_copy)
            except Exception:
                pass


# ============================================================================
# GRADIO INTERFACE
# ============================================================================

def create_interface():
    """Build the Gradio interface."""

    PROFESSIONAL_CSS = """
    /* ============================================================
       BrainGemma3D Dashboard  ·  Professional Medical UI
    ============================================================ */

    /* ---- Font stack (Google Fonts loaded by gr.themes.GoogleFont) ---- */
    *, *::before, *::after {
        font-family: 'Inter', 'Segoe UI', 'Helvetica Neue', Arial, sans-serif !important;
        -webkit-font-smoothing: antialiased;
    }
    code, pre, kbd, .monospace {
        font-family: 'IBM Plex Mono', 'Consolas', 'Courier New', monospace !important;
        font-size: 0.82em !important;
    }

    /* ---- Page background ---- */
    body, .gradio-container { background: var(--background-fill-primary); }

    /* ---- Headings ---- */
    h1 {
        font-size: 1.5rem !important;
        font-weight: 700 !important;
        letter-spacing: -0.02em !important;
        color: var(--body-text-color) !important;
    }
    h2 {
        font-size: 0.95rem !important;
        font-weight: 600 !important;
        text-transform: uppercase !important;
        letter-spacing: 0.06em !important;
        border-bottom: 2px solid var(--border-color-primary) !important;
        padding-bottom: 4px !important;
        margin-bottom: 10px !important;
        color: var(--body-text-color) !important;
    }
    h3 {
        font-size: 0.88rem !important;
        font-weight: 600 !important;
        color: var(--body-text-color) !important;
    }

    /* ---- Buttons ---- */
    button.primary {
        background: linear-gradient(135deg, #14325a 0%, #1e5799 100%) !important;
        color: #fff !important;
        font-weight: 600 !important;
        font-size: 0.82rem !important;
        letter-spacing: 0.08em !important;
        text-transform: uppercase !important;
        border-radius: 6px !important;
        border: none !important;
        box-shadow: 0 2px 8px rgba(20,50,90,0.30) !important;
        transition: opacity 0.15s, box-shadow 0.15s !important;
    }
    button.primary:hover {
        opacity: 0.88 !important;
        box-shadow: 0 4px 16px rgba(20,50,90,0.40) !important;
    }
    button.secondary {
        background: var(--button-secondary-background-fill) !important;
        color: var(--button-secondary-text-color) !important;
        border: 1.5px solid var(--border-color-primary) !important;
        border-radius: 6px !important;
        font-weight: 500 !important;
        font-size: 0.82rem !important;
        letter-spacing: 0.04em !important;
        transition: background 0.15s !important;
    }
    button.secondary:hover { 
        background: var(--button-secondary-background-fill-hover) !important;
    }

    /* ---- Sliders ---- */
    input[type='range'] { accent-color: #1a3a5c !important; }

    /* ---- Tab bar ---- */
    .tab-nav button {
        font-weight: 500 !important;
        font-size: 0.80rem !important;
        letter-spacing: 0.05em !important;
        text-transform: uppercase !important;
        color: var(--body-text-color-subdued) !important;
    }
    .tab-nav button.selected {
        border-bottom: 3px solid #1e5799 !important;
        color: var(--body-text-color) !important;
        font-weight: 700 !important;
    }

    /* ---- Report area ---- */
    .report-box p, .report-box li {
        font-size: 0.875rem !important;
        line-height: 1.75 !important;
        color: var(--body-text-color) !important;
    }
    .report-box h3 {
        color: var(--body-text-color) !important;
        border-bottom: 1px solid var(--border-color-primary) !important;
        padding-bottom: 2px !important;
    }

    /* ---- Misc ---- */
    .progress-bar { background-color: #1e5799 !important; }
    footer { visibility: hidden !important; }

    /* ---- MRI viewer: force black background on image containers ---- */
    .viewer-tab .image-container,
    .viewer-tab img,
    .viewer-image > div,
    .viewer-image .image-container {
        background-color: #000000 !important;
    }
    """

    with gr.Blocks(
        title="BrainGemma3D — Brain MRI Report",
        css=PROFESSIONAL_CSS,
    ) as demo:

        # ---- Page header -----------------------------------------------
        gr.Markdown(
            '<h1>🧠 BrainGemma3D — Brain Report Automation via Inflated Vision Transformers in 3D</h1>'
        )
        gr.Markdown(
            f'''
            Upload a <strong>NIfTI volume</strong> (.nii / .nii.gz) to generate an AI-assisted \
brain MRI diagnostic report using <strong>BrainGemma3D</strong>.
            <blockquote>⚠️ <strong>Research use only.</strong> Not intended for clinical diagnosis or treatment decisions.</blockquote>
            '''
        )

        # ---- Hidden states ---------------------------------------------
        volume_state      = gr.State(None)
        raw_report_state  = gr.State("")

        with gr.Row(equal_height=False):
            # ============================================================
            # LEFT — Upload, patient info, parameters
            # ============================================================
            with gr.Column(scale=1, min_width=380):

                gr.Markdown('<h2>📂 Upload</h2>')
                nifti_input = gr.File(
                    label="NIfTI File  (.nii / .nii.gz)",
                    file_types=[".nii", ".gz"],  # .gz needed to pass .nii.gz through Gradio's client-side check
                    type="filepath",
                )
                gr.Markdown(
                    '<blockquote>📁 <strong>Accepted formats:</strong> <code>.nii</code> and <code>.nii.gz</code> (NIfTI). '
                    'Upload a 3D brain MRI volume. '
                    'Once uploaded, mid-volume slices will be previewed automatically.</blockquote>'
                )

                gr.Markdown('<h2>🏥 Patient Information</h2>')
                with gr.Group():
                    with gr.Row():
                        patient_name = gr.Textbox(
                            label="Patient Name",
                            placeholder="Last, First",
                            max_lines=1,
                        )
                        patient_id = gr.Textbox(
                            label="Patient ID / MRN",
                            placeholder="e.g. 0012345",
                            max_lines=1,
                        )
                    with gr.Row():
                        exam_date = gr.Textbox(
                            label="Exam Date",
                            placeholder=datetime.date.today().strftime('%Y-%m-%d'),
                            max_lines=1,
                        )
                        referring_physician = gr.Textbox(
                            label="Referring Physician",
                            placeholder="Dr. ...",
                            max_lines=1,
                        )
                    institution = gr.Textbox(
                        label="Institution / Department",
                        placeholder="e.g. Neuroradiology Dept. — General Hospital",
                        max_lines=1,
                    )

                with gr.Accordion("⚙️  Generation Parameters", open=False):
                    custom_prompt = gr.Textbox(
                        label="Additional instructions or questions about the image (optional)",
                        placeholder="E.g. Focus on white-matter lesions or mass effect; or ask a specific question about the scan (e.g. 'Is there surrounding edema?')",
                        lines=4,
                        value="",
                        info="Leave blank to use the canonical prompt; if provided, this text will be appended as additional instructions and may include specific questions about the scan.",
                    )
                    max_tokens = gr.Slider(
                        label="Max Tokens",
                        minimum=50, maximum=512, value=160, step=10,
                        info="Maximum length of the generated text.",
                    )
                    temperature = gr.Slider(
                        label="Temperature",
                        minimum=0.0, maximum=2.0, value=0.1, step=0.05,
                        info="Lower → more deterministic.",
                    )
                    top_p = gr.Slider(
                        label="Top-p (Nucleus Sampling)",
                        minimum=0.0, maximum=1.0, value=0.9, step=0.05,
                    )
                    repetition_penalty = gr.Slider(
                        label="Repetition Penalty",
                        minimum=1.0, maximum=2.0, value=1.2, step=0.1,
                    )

                with gr.Row():
                    generate_btn = gr.Button(
                        "Generate", variant="primary", size="lg"
                    )
                    clear_btn = gr.Button("🗑️  Clear All", variant="secondary")

                # -- 3D Volume Viewer (hidden until a volume is loaded) --
                viewer_section = gr.Group(visible=False)
                with viewer_section:
                    gr.Markdown('<h2>🔬 3D Volume Viewer</h2>')
                    with gr.Tabs(elem_classes=["viewer-tab"]):
                        with gr.Tab("Axial"):
                            axial_slider = gr.Slider(
                                minimum=0, maximum=100, value=50, step=1,
                                label="Slice index", interactive=True,
                            )
                            axial_image = gr.Image(
                                label="Axial View", type="pil", height=360,
                                elem_classes=["viewer-image"],
                            )
                        with gr.Tab("Coronal"):
                            coronal_slider = gr.Slider(
                                minimum=0, maximum=100, value=50, step=1,
                                label="Slice index", interactive=True,
                            )
                            coronal_image = gr.Image(
                                label="Coronal View", type="pil", height=360,
                                elem_classes=["viewer-image"],
                            )
                        with gr.Tab("Sagittal"):
                            sagittal_slider = gr.Slider(
                                minimum=0, maximum=100, value=50, step=1,
                                label="Slice index", interactive=True,
                            )
                            sagittal_image = gr.Image(
                                label="Sagittal View", type="pil", height=360,
                                elem_classes=["viewer-image"],
                            )

            # ============================================================
            # RIGHT — Report output and actions
            # ============================================================
            with gr.Column(scale=1, min_width=380):

                gr.Markdown('<h2>📋 Diagnostic Report</h2>')

                report_output = gr.Markdown(
                    value='<em style="opacity: 0.7;">The generated report will appear here.</em>',
                    elem_classes=["report-box"],
                )
                
                gr.Markdown('<h2>🔬 LIME Interpretability Analysis</h2>')
                gr.Markdown(
                    '<blockquote>This shows which brain regions most strongly support (🔴 red) or contradict (🔵 blue) the generated diagnosis.</blockquote>'
                )
                
                interpretability_output = gr.Image(
                    label="LIME Interpretability (2×3 Grid: Original + Overlay)",
                    type="pil",
                    height=500,
                    elem_classes=["interpretability-plot"],
                )

                with gr.Row():
                    pdf_btn  = gr.Button(
                        "📄  Download PDF Report", size="sm", variant="primary"
                    )

                pdf_output = gr.File(
                    label="PDF Report",
                    visible=False,
                    interactive=False,
                )

        # ================================================================
        # EVENT HANDLERS
        # ================================================================

        # -- Generate report --
        generate_btn.click(
            fn=generate_report_fn,
            inputs=[
                nifti_input,
                custom_prompt,
                max_tokens,
                temperature,
                top_p,
                repetition_penalty,
                volume_state,
            ],
            outputs=[
                report_output,
                raw_report_state,
                volume_state,
                axial_slider,
                coronal_slider,
                sagittal_slider,
                axial_image,
                coronal_image,
                sagittal_image,
                interpretability_output,
            ],
            show_progress="full",
        ).then(
            fn=lambda vol: (
                gr.update(maximum=vol.shape[0] - 1, value=vol.shape[0] // 2)
                    if vol is not None else gr.update(),
                gr.update(maximum=vol.shape[1] - 1, value=vol.shape[1] // 2)
                    if vol is not None else gr.update(),
                gr.update(maximum=vol.shape[2] - 1, value=vol.shape[2] // 2)
                    if vol is not None else gr.update(),
                gr.update(visible=True)
                    if vol is not None else gr.update(visible=False),
            ),
            inputs=[volume_state],
            outputs=[axial_slider, coronal_slider, sagittal_slider, viewer_section],
        )

        # -- Slice sliders --
        def update_axial(vol_np, slice_idx):
            if vol_np is None:
                return Image.new('RGB', (800, 400), color='#0d0d0d')
            return create_slice_image(vol_np, 'axial', int(slice_idx))

        def update_coronal(vol_np, slice_idx):
            if vol_np is None:
                return Image.new('RGB', (800, 400), color='#0d0d0d')
            return create_slice_image(vol_np, 'coronal', int(slice_idx))

        def update_sagittal(vol_np, slice_idx):
            if vol_np is None:
                return Image.new('RGB', (800, 400), color='#0d0d0d')
            return create_slice_image(vol_np, 'sagittal', int(slice_idx))

        axial_slider.release(
            fn=update_axial,
            inputs=[volume_state, axial_slider],
            outputs=axial_image,
        )
        coronal_slider.release(
            fn=update_coronal,
            inputs=[volume_state, coronal_slider],
            outputs=coronal_image,
        )
        sagittal_slider.release(
            fn=update_sagittal,
            inputs=[volume_state, sagittal_slider],
            outputs=sagittal_image,
        )

        # -- Preview on upload --
        nifti_input.upload(
            fn=preview_nifti_fn,
            inputs=[nifti_input],
            outputs=[
                report_output,
                volume_state,
                axial_slider,
                coronal_slider,
                sagittal_slider,
                axial_image,
                coronal_image,
                sagittal_image,
            ],
            show_progress="full",
        ).then(
            fn=lambda vol: (
                gr.update(maximum=vol.shape[0] - 1, value=vol.shape[0] // 2)
                    if vol is not None else gr.update(),
                gr.update(maximum=vol.shape[1] - 1, value=vol.shape[1] // 2)
                    if vol is not None else gr.update(),
                gr.update(maximum=vol.shape[2] - 1, value=vol.shape[2] // 2)
                    if vol is not None else gr.update(),
                gr.update(visible=True)
                    if vol is not None else gr.update(visible=False),
            ),
            inputs=[volume_state],
            outputs=[axial_slider, coronal_slider, sagittal_slider, viewer_section],
        )

        # -- Clear --
        def clear_all_fn():
            empty = Image.new('RGB', (800, 400), color='#0d0d0d')
            return (
                None,
                "*The generated report will appear here.*",
                "",
                None,
                gr.update(value=0, maximum=100),
                gr.update(value=0, maximum=100),
                gr.update(value=0, maximum=100),
                empty, empty, empty,
                gr.update(visible=False),
                gr.update(visible=False),
            )

        clear_btn.click(
            fn=clear_all_fn,
            inputs=[],
            outputs=[
                nifti_input,
                report_output,
                raw_report_state,
                volume_state,
                axial_slider,
                coronal_slider,
                sagittal_slider,
                axial_image,
                coronal_image,
                sagittal_image,
                viewer_section,
                pdf_output,
            ],
        )

        # -- Generate PDF --
        def on_pdf_click(vol_np, report_text, p_name, p_id, e_date, ref_phys, inst):
            path = generate_pdf_fn(
                vol_np, report_text, p_name, p_id, e_date, ref_phys, inst
            )
            if path is None:
                return gr.update(visible=False)
            return gr.update(value=path, visible=True)

        pdf_btn.click(
            fn=on_pdf_click,
            inputs=[
                volume_state,
                raw_report_state,
                patient_name,
                patient_id,
                exam_date,
                referring_physician,
                institution,
            ],
            outputs=pdf_output,
        )

        # ================================================================
        # FOOTER
        # ================================================================
        gr.HTML("""
        <div style="
            margin-top: 2.5rem;
            padding: 1.6rem 2rem;
            border-radius: 10px;
            border-top: 3px solid #1976d2;
            border: 1px solid var(--border-color-primary);
            background: var(--background-fill-secondary);
            font-family: 'Inter', 'Segoe UI', Arial, sans-serif;
        ">
            <!-- Top rule -->
            <div style="display:flex; align-items:center; gap:0.75rem; margin-bottom:1.1rem;">
            <div style="flex:1; height:1px; background:#1976d2; opacity:0.5;"></div>
            <span style="font-size:0.70rem; font-weight:700; letter-spacing:0.14em;
                     text-transform:uppercase; color: var(--body-text-color);">
                About this project
            </span>
            <div style="flex:1; height:1px; background:#1976d2; opacity:0.5;"></div>
            </div>

            <!-- Team -->
            <p style="margin:0 0 0.35rem 0; font-size:0.78rem; font-weight:600;
                  letter-spacing:0.06em; text-transform:uppercase; color: var(--body-text-color);">
            Development Team
            </p>
            <p style="margin:0 0 0.9rem 0; font-size:0.88rem; line-height:1.6; color: var(--body-text-color);">
            Mariano Barone &nbsp;·&nbsp; Francesco Di Serio &nbsp;·&nbsp;
            Giuseppe Riccio &nbsp;·&nbsp; Antonio Romano &nbsp;·&nbsp; Vincenzo Moscato
            </p>

            <!-- Institution -->
            <p style="margin:0 0 1.2rem 0; font-size:0.80rem; line-height:1.5; color: var(--body-text-color-subdued);">
            Department of Electrical Engineering and Information Technology<br>
            <strong style="color: var(--body-text-color);">University of Naples Federico II</strong>,
            Italy
            </p>

            <!-- Divider -->
            <div style="height:1px; background: var(--border-color-primary); margin-bottom:1.1rem;"></div>

            <!-- Badge row -->
            <div style="display:flex; flex-wrap:wrap; align-items:center;
                gap:1.2rem; justify-content:space-between;">

            <div style="display:flex; align-items:center; gap:0.5rem;">
                <span style="font-size:1.1rem;">❤️</span>
                <span style="font-size:0.80rem; line-height:1.4; color: var(--body-text-color);">
                Built for the
                <strong style="color:#ffd570;">MedGemma Impact Challenge 🏆</strong><br>
                <span style="font-size:0.74rem; color: var(--body-text-color-subdued);">
                    Advancing Medical AI with Google's Health AI Developer Foundations
                </span>
                </span>
            </div>

            <a href="https://github.com/PRAISELab-PicusLab/BrainGemma3D/"
               target="_blank" rel="noopener noreferrer"
               style="
                   display:inline-flex; align-items:center; gap:0.45rem;
                   padding:0.45rem 1.0rem;
                   background: var(--button-secondary-background-fill);
                   border:1.5px solid #3a6a9a;
                   border-radius:6px; text-decoration:none;
                   font-size:0.78rem; font-weight:600;
                   color: var(--body-text-color);
                   letter-spacing:0.03em; transition:opacity 0.15s;
                   white-space:nowrap;
               "
               onmouseover="this.style.opacity='0.8'"
               onmouseout="this.style.opacity='1'">
                <!-- GitHub mark SVG -->
                <svg height="16" viewBox="0 0 16 16" width="16" fill="currentColor"
                 xmlns="http://www.w3.org/2000/svg">
                <path d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38
                     0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13
                     -.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66
                     .07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15
                     -.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27
                     .68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12
                     .51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48
                     0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0016 8c0-4.42-3.58-8-8-8z"/>
                </svg>
                View Source on GitHub
            </a>
            </div>
        </div>
        """)

    return demo


# ============================================================================
# MAIN
# ============================================================================

# 1) Model config: default for Spaces
DEFAULT_MODEL_DIR = str(Path(__file__).parent.parent / "BrainGemma3D")
MODEL_CONFIG.update({
    "checkpoint_dir": os.environ.get("MODEL_DIR", DEFAULT_MODEL_DIR),
    "target_size": [64, 128, 128],
})

# 2) Load model
MODEL = load_model(DEFAULT_MODEL_DIR)

# 3) Create the demo at module-level
demo = create_interface()

# 4) Queue + launch with PORT
demo.queue(default_concurrency_limit=5)

demo.launch(
    server_name="0.0.0.0",
    server_port=int(os.environ.get("PORT", 7860)),
    show_error=True,
    inbrowser=False,
    # ssr_mode=False,
    theme=gr.themes.Base(
        font=[gr.themes.GoogleFont("Inter"), "Helvetica Neue", "Arial", "sans-serif"],
        font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "Fira Code", "Consolas", "monospace"],
        primary_hue=gr.themes.colors.blue,
        secondary_hue=gr.themes.colors.slate,
        neutral_hue=gr.themes.colors.slate,
        text_size=gr.themes.sizes.text_md,
        radius_size=gr.themes.sizes.radius_md,
        spacing_size=gr.themes.sizes.spacing_md,
    ),
)