src/BrainIAC/app_gradio.py

import os
import torch
import nibabel as nib
import gradio as gr
import tempfile
import yaml
import traceback # For detailed error printing
import zipfile
import dicom2nifti
import shutil
import subprocess # To run unzip command
import SimpleITK as sitk
import itk
import numpy as np
from scipy.signal import medfilt
import skimage.filters
import cv2 # For Gaussian Blur
import io # For saving plots to memory
import base64 # For encoding plots
import uuid # For unique IDs
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg') # Use non-interactive backend

# --- Custom CSS to hide buttons in logo columns ---
# REVERTING CSS CHANGES AS THEY HID THE LOGOS
# custom_css = """
# .logo-column button {
#     /* display: none !important; */ /* This seemed to hide the whole component */
#     visibility: hidden !important; /* Try making it invisible instead */
# }
# """

# --- Potential Import Issues (Check Paths in Docker/Local) ---
try:
    # Assumes HD_BET is now at /app/BrainIAC/HD_BET or adjacent in src
    from HD_BET.run import run_hd_bet
    from monai.visualize.gradient_based import GuidedBackpropSmoothGrad
except ImportError as e:
    print(f"Warning: Could not import HD_BET or MONAI visualize: {e}. Saliency/Preprocessing might fail.")
    run_hd_bet = None
    GuidedBackpropSmoothGrad = None

# Import necessary components from your existing modules
from model import Backbone, SingleScanModel, Classifier
from monai.transforms import Resized, ScaleIntensityd

# --- Constants ---
APP_DIR = os.path.dirname(__file__)
TEMPLATE_DIR = os.path.join(APP_DIR, "golden_image", "mni_templates")
PARAMS_RIGID_PATH = os.path.join(APP_DIR, "golden_image", "mni_templates", "Parameters_Rigid.txt")
DEFAULT_TEMPLATE_PATH = os.path.join(TEMPLATE_DIR, "nihpd_asym_13.0-18.5_t1w.nii")
HD_BET_CONFIG_PATH = os.path.join(APP_DIR, "HD_BET", "config.py") # May need adjustment based on actual HD_BET location
HD_BET_MODEL_DIR = os.path.join(APP_DIR, "hdbet_model") # Path to copied models

# --- Configuration Loading ---
def load_config():
    config_path = os.path.join(APP_DIR, 'config.yml')
    try:
        with open(config_path, 'r') as file:
            config = yaml.safe_load(file)
        if 'data' not in config: config['data'] = {}
        if 'image_size' not in config['data']: config['data']['image_size'] = [128, 128, 128]
    except FileNotFoundError:
        print(f"Warning: Configuration file not found at {config_path}. Using defaults.")
        config = {
            'gpu': {'device': 'cpu'},
            'infer': {'checkpoints': 'checkpoints/mci_model.pt'}, # Updated for new MCI model filename
            'data': {'image_size': [128, 128, 128]}
        }
    return config

config = load_config()
DEFAULT_IMAGE_SIZE = (128, 128, 128)
image_size_cfg = config.get('data', {}).get('image_size', DEFAULT_IMAGE_SIZE)
if not isinstance(image_size_cfg, (list, tuple)) or len(image_size_cfg) != 3:
    print(f"Warning: Invalid image_size in config ({image_size_cfg}). Using default {DEFAULT_IMAGE_SIZE}.")
    image_size = DEFAULT_IMAGE_SIZE
else:
    image_size = tuple(image_size_cfg)

# --- Model Loading ---
def load_model(cfg):
    device = torch.device(cfg.get('gpu', {}).get('device', 'cpu'))
    backbone = Backbone()
    classifier = Classifier(d_model=2048, num_classes=1)  # Binary classification for MCI
    model = SingleScanModel(backbone, classifier)  # Using BP model for MCI classification
    relative_path = cfg.get('infer', {}).get('checkpoints', 'checkpoints/mci_model.pt')
    checkpoint_path_abs = os.path.join(APP_DIR, relative_path)
    try:
        print(f"Loading MCI classification model from: {checkpoint_path_abs}")
        checkpoint = torch.load(checkpoint_path_abs, map_location=device, weights_only=False)
        state_dict = checkpoint.get('model_state_dict', checkpoint)
        model.load_state_dict(state_dict, strict=False)
        model.to(device)
        model.eval()
        print(f"MCI classification model loaded successfully onto {device}")
        return model, device
    except FileNotFoundError:
        print(f"Error: Checkpoint file not found at {checkpoint_path_abs}")
        return None, device
    except Exception as e:
        print(f"Error loading model checkpoint: {e}")
        traceback.print_exc()
        return None, device

model, device = load_model(config)

# --- Preprocessing Functions (Copied/Adapted from app.py) ---
def bias_field_correction(img_array):
    print("  Running N4 Bias Field Correction...")
    image = sitk.GetImageFromArray(img_array)
    if image.GetPixelID() != sitk.sitkFloat32:
        image = sitk.Cast(image, sitk.sitkFloat32)
    maskImage = sitk.OtsuThreshold(image, 0, 1, 200)
    corrector = sitk.N4BiasFieldCorrectionImageFilter()
    numberFittingLevels = 4
    max_iters = [min(50 * (2**i), 200) for i in range(numberFittingLevels)]
    corrector.SetMaximumNumberOfIterations(max_iters)
    corrected_image = corrector.Execute(image, maskImage)
    print("  N4 Correction finished.")
    return sitk.GetArrayFromImage(corrected_image)

def denoise(volume, kernel_size=3):
    print(f"  Applying median filter denoising (kernel={kernel_size})...")
    return medfilt(volume, kernel_size)

def rescale_intensity(volume, percentils=[0.5, 99.5], bins_num=256):
    print("  Rescaling intensity...")
    volume_float = volume.astype(np.float32)
    try:
        t = skimage.filters.threshold_otsu(volume_float, nbins=256)
        volume_masked = np.copy(volume_float)
        volume_masked[volume_masked < t] = 0
        obj_volume = volume_masked[np.where(volume_masked > 0)]
    except ValueError:
        print("    Otsu failed, skipping background mask.")
        obj_volume = volume_float.flatten()
    if obj_volume.size == 0:
        print("    Warning: No foreground voxels found. Scaling full volume.")
        obj_volume = volume_float.flatten()
        min_value = np.min(obj_volume)
        max_value = np.max(obj_volume)
    else:
        min_value = np.percentile(obj_volume, percentils[0])
        max_value = np.percentile(obj_volume, percentils[1])
    denominator = max_value - min_value
    if denominator < 1e-6: denominator = 1e-6
    output_volume = np.copy(volume_float)
    if bins_num == 0:
        output_volume = (volume_float - min_value) / denominator
        output_volume = np.clip(output_volume, 0.0, 1.0)
    else:
        output_volume = np.round((volume_float - min_value) / denominator * (bins_num - 1))
        output_volume = np.clip(output_volume, 0, bins_num - 1)
    return output_volume.astype(np.float32)

def equalize_hist(volume, bins_num=256):
    print("  Performing histogram equalization...")
    mask = volume > 1e-6
    obj_volume = volume[mask]
    if obj_volume.size == 0:
        print("    Warning: No non-zero voxels. Skipping equalization.")
        return volume
    hist, bins = np.histogram(obj_volume, bins_num, range=(obj_volume.min(), obj_volume.max()))
    cdf = hist.cumsum()
    cdf_normalized = (bins_num - 1) * cdf / float(cdf[-1])
    equalized_obj_volume = np.interp(obj_volume, bins[:-1], cdf_normalized)
    equalized_volume = np.copy(volume)
    equalized_volume[mask] = equalized_obj_volume
    return equalized_volume.astype(np.float32)

def enhance(img_array, run_bias_correction=True, kernel_size=3, percentils=[0.5, 99.5], bins_num=256, run_equalize_hist=True):
    print("Starting enhancement pipeline...")
    volume = img_array.astype(np.float32)
    try:
        if run_bias_correction: volume = bias_field_correction(volume)
        volume = denoise(volume, kernel_size)
        volume = rescale_intensity(volume, percentils, bins_num)
        if run_equalize_hist: volume = equalize_hist(volume, bins_num)
        print("Enhancement pipeline finished.")
        return volume
    except Exception as e:
        print(f"Error during enhancement: {e}")
        traceback.print_exc()
        raise RuntimeError(f"Failed enhancing image: {e}")

def register_image(input_nifti_path, output_nifti_path):
    print(f"Registering {input_nifti_path} to {DEFAULT_TEMPLATE_PATH}")
    if not all(os.path.exists(p) for p in [PARAMS_RIGID_PATH, DEFAULT_TEMPLATE_PATH]):
        raise FileNotFoundError("Elastix parameter or template file not found.")
    fixed_image = itk.imread(DEFAULT_TEMPLATE_PATH, itk.F)
    moving_image = itk.imread(input_nifti_path, itk.F)
    parameter_object = itk.ParameterObject.New()
    parameter_object.AddParameterFile(PARAMS_RIGID_PATH)
    result_image, _ = itk.elastix_registration_method(fixed_image, moving_image, parameter_object=parameter_object, log_to_console=False)
    itk.imwrite(result_image, output_nifti_path)
    print(f"Registration output saved to {output_nifti_path}")

def run_enhance_on_file(input_nifti_path, output_nifti_path):
    print(f"Running full enhancement on {input_nifti_path}")
    img_sitk = sitk.ReadImage(input_nifti_path)
    img_array = sitk.GetArrayFromImage(img_sitk)
    enhanced_array = enhance(img_array, run_bias_correction=True)
    enhanced_img_sitk = sitk.GetImageFromArray(enhanced_array)
    enhanced_img_sitk.CopyInformation(img_sitk)
    sitk.WriteImage(enhanced_img_sitk, output_nifti_path)
    print(f"Enhanced image saved to {output_nifti_path}")

def run_skull_stripping(input_nifti_path, output_dir):
    print(f"Running HD-BET skull stripping on {input_nifti_path}")
    if run_hd_bet is None: raise RuntimeError("HD-BET module not imported.")
    if not os.path.exists(HD_BET_CONFIG_PATH): raise FileNotFoundError(f"HD-BET config not found at {HD_BET_CONFIG_PATH}")
    if not os.path.isdir(HD_BET_MODEL_DIR): raise FileNotFoundError(f"HD-BET models not found at {HD_BET_MODEL_DIR}")

    base_name = os.path.basename(input_nifti_path).replace(".nii.gz", "").replace(".nii", "")
    output_file_path = os.path.join(output_dir, f"{base_name}_bet.nii.gz")
    output_mask_path = os.path.join(output_dir, f"{base_name}_bet_mask.nii.gz")
    os.makedirs(output_dir, exist_ok=True)

    try:
        run_hd_bet(input_nifti_path, output_file_path, mode="fast", device='cpu', config_file=HD_BET_CONFIG_PATH, postprocess=False, do_tta=False, keep_mask=True, overwrite=True)
    finally:
        pass

    if not os.path.exists(output_file_path): raise RuntimeError("HD-BET did not produce output file.")
    print(f"Skull stripping output saved to {output_file_path}")
    return output_file_path, output_mask_path

# --- MONAI Transforms ---
resize_transform = Resized(keys=["image"], spatial_size=image_size)
scale_transform = ScaleIntensityd(keys=["image"], minv=0.0, maxv=1.0)

def preprocess_nifti_for_model(nifti_path):
    print(f"Preprocessing NIfTI for model: {nifti_path}")
    scan_data = nib.load(nifti_path).get_fdata()
    scan_tensor = torch.tensor(scan_data, dtype=torch.float32).unsqueeze(0) # Add C dim
    sample = {"image": scan_tensor}
    sample_resized = resize_transform(sample)
    sample_scaled = scale_transform(sample_resized)
    input_tensor = sample_scaled["image"].unsqueeze(0).to(device) # Add B dim
    if input_tensor.dim() != 5: raise ValueError(f"Preprocessing resulted in incorrect shape: {input_tensor.shape}")
    print(f"  Final shape for model: {input_tensor.shape}")
    return input_tensor

# --- Saliency Generation ---
def generate_saliency(model_to_use, input_tensor_5d):
    if GuidedBackpropSmoothGrad is None: raise ImportError("MONAI visualize components not imported.")
    if model_to_use is None: raise ValueError("Model not loaded.")
    print("Generating saliency map...")
    input_tensor_5d.requires_grad_(True)
    visualizer = GuidedBackpropSmoothGrad(model=model_to_use.backbone.to(device), stdev_spread=0.15, n_samples=10, magnitude=True)
    try:
        with torch.enable_grad():
            saliency_map_5d = visualizer(input_tensor_5d.to(device))
        input_3d = input_tensor_5d.squeeze().cpu().detach().numpy()
        saliency_3d = saliency_map_5d.squeeze().cpu().detach().numpy()
        print("Saliency map generated.")
        return input_3d, saliency_3d
    except Exception as e:
        print(f"Error during saliency map generation: {e}")
        traceback.print_exc()
        return None, None
    finally:
        input_tensor_5d.requires_grad_(False)

# --- Plotting Function (Returns NumPy arrays for Gradio) ---
def create_slice_plots(mri_data_3d, saliency_data_3d, slice_index):
    print(f"  Generating plots for slice index: {slice_index}")
    if mri_data_3d is None or saliency_data_3d is None: return None, None, None
    
    # Change from the third dimension (axis 2, sagittal) to the first dimension (axis 0, axial)
    if not (0 <= slice_index < mri_data_3d.shape[0]):
        print(f"    Error: Slice index {slice_index} out of bounds (0-{mri_data_3d.shape[0]-1}).")
        return None, None, None

    # Function to save plot to NumPy array
    def save_plot_to_numpy(fig):
        with io.BytesIO() as buf:
            fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=75) # Adjust DPI as needed
            plt.close(fig)
            buf.seek(0)
            img_arr = plt.imread(buf, format='png')
        # Return RGBA array, can be simplified if only grayscale needed for input
        return (img_arr * 255).astype(np.uint8)

    try:
        # Extract axial slice instead of sagittal
        mri_slice = mri_data_3d[slice_index, :, :]
        saliency_slice_orig = saliency_data_3d[slice_index, :, :]

        # Normalize MRI Slice (using volume stats)
        p1_vol, p99_vol = np.percentile(mri_data_3d, (1, 99))
        mri_norm_denom = max(p99_vol - p1_vol, 1e-6)
        mri_slice_norm = np.clip((mri_slice - p1_vol) / mri_norm_denom, 0, 1)

        # Process Saliency Slice
        saliency_slice = np.copy(saliency_slice_orig)
        saliency_slice[saliency_slice < 0] = 0
        saliency_slice_blurred = cv2.GaussianBlur(saliency_slice, (15, 15), 0)
        s_max_vol = max(np.max(saliency_data_3d[saliency_data_3d >= 0]), 1e-6)
        saliency_slice_norm = saliency_slice_blurred / s_max_vol
        saliency_slice_thresholded = np.where(saliency_slice_norm > 0.0, saliency_slice_norm, 0) # Threshold slightly > 0

        # Plot 1: Input Slice
        fig1, ax1 = plt.subplots(figsize=(6, 6))
        ax1.imshow(mri_slice_norm, cmap='gray', interpolation='none', origin='lower')
        ax1.axis('off')
        input_plot_np = save_plot_to_numpy(fig1)

        # Plot 2: Saliency Heatmap
        fig2, ax2 = plt.subplots(figsize=(6, 6))
        ax2.imshow(saliency_slice_thresholded, cmap='magma', interpolation='none', origin='lower', vmin=0) # Set vmin
        ax2.axis('off')
        heatmap_plot_np = save_plot_to_numpy(fig2)

        # Plot 3: Overlay
        fig3, ax3 = plt.subplots(figsize=(6, 6))
        ax3.imshow(mri_slice_norm, cmap='gray', interpolation='none', origin='lower')
        if np.max(saliency_slice_thresholded) > 0:
             ax3.contour(saliency_slice_thresholded, cmap='magma', origin='lower', linewidths=1.0, levels=np.linspace(saliency_slice_thresholded.min(), saliency_slice_thresholded.max(), 5)) # Adjust levels
        ax3.axis('off')
        overlay_plot_np = save_plot_to_numpy(fig3)

        print(f"    Generated numpy plots successfully for slice {slice_index}.")
        return input_plot_np, heatmap_plot_np, overlay_plot_np

    except Exception as e:
        print(f"Error generating numpy plots for slice {slice_index}: {e}")
        traceback.print_exc()
        return None, None, None

# Add this function after the create_slice_plots function
def create_probability_gauge(probability):
    """
    Creates a gauge visualization for the MCI probability.
    
    Args:
        probability (float): A value between 0 and 1 representing the MCI probability.
        
    Returns:
        numpy.ndarray: A numpy array containing the gauge visualization image.
    """
    # Create a figure with a polar projection
    fig = plt.figure(figsize=(8, 4))
    ax = fig.add_subplot(111, polar=True)
    
    # Set the min and max angles for the gauge (in radians)
    # -pi/2 to pi/2 creates a half-circle (180 degrees)
    theta_min = -np.pi/2
    theta_max = np.pi/2
    
    # Calculate the angle for the needle based on probability (0 to 1)
    needle_angle = theta_min + probability * (theta_max - theta_min)
    
    # Create a color gradient for the gauge background
    cmap = plt.cm.RdYlGn_r  # Red-Yellow-Green colormap (reversed)
    
    # Draw the gauge background
    theta = np.linspace(theta_min, theta_max, 100)
    radii = np.ones_like(theta)
    
    # Create color array for the gauge segments
    norm = plt.Normalize(0, 1)
    colors = cmap(np.linspace(0, 1, len(theta)))
    
    # Draw colored bars for the gauge
    width = (theta_max - theta_min) / len(theta)
    bars = ax.bar(theta, radii, width=width, bottom=0.0, alpha=0.8, linewidth=0)
    
    # Set the color for each bar
    for bar, color in zip(bars, colors):
        bar.set_facecolor(color)
    
    # Add the needle
    needle_length = 0.9
    ax.annotate('', xy=(needle_angle, needle_length), xytext=(needle_angle, 0),
                arrowprops=dict(arrowstyle='wedge', color='black', lw=2))
    
    # Add boundary markers and labels - move them lower by adjusting y position (1.1 -> 1.3)
    ax.text(theta_min, 2.2, 'Healthy Control', ha='left', va='center', fontsize=12)
    ax.text(theta_max, 1.3, 'MCI', ha='right', va='center', fontsize=12)
    
    # Add the probability text below the gauge
    prob_text = f"Probability: {probability:.2f}"
    fig.text(0.5, 0.15, prob_text, ha='center', va='center', fontsize=14, fontweight='bold')
    
    # Set the limits and remove unnecessary elements
    ax.set_ylim(0, 1.4)  # Increased the upper limit to accommodate the lower labels
    ax.set_theta_zero_location('N')  # 0 at the top
    ax.set_theta_direction(-1)  # clockwise
    ax.set_thetagrids([])  # Remove angle labels
    ax.grid(False)  # Remove grid
    ax.set_rgrids([])  # Remove radial labels
    ax.spines['polar'].set_visible(False)  # Remove the outer circle
    
    # Convert figure to numpy array
    with io.BytesIO() as buf:
        fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.1, dpi=100)
        plt.close(fig)
        buf.seek(0)
        img_arr = plt.imread(buf, format='png')
    
    return (img_arr * 255).astype(np.uint8)

# --- Gradio Processing Function ---
def process_scan(file_type, uploaded_file, run_preprocess, generate_saliency_flag):
    if model is None:
        raise gr.Error("Model is not loaded. Cannot perform prediction.")
    if uploaded_file is None:
        raise gr.Error("No file uploaded.")

    temp_dir = tempfile.mkdtemp()
    print(f"Created temp directory: {temp_dir}")
    nifti_for_preprocessing_path = None
    error_message = None
    prediction_text = "Processing..."
    # Initialize outputs to None or placeholder images/values
    input_slice_img, heatmap_slice_img, overlay_slice_img = None, None, None
    probability_gauge = None
    saliency_state = {"input_path": None, "saliency_path": None, "num_slices": 0}
    slider_update = gr.Slider(value=0, minimum=0, maximum=1, visible=False) # Initially hidden, use max=1 to avoid log(0) error

    try:
        # --- Handle Upload and DICOM Conversion ---
        file_path = uploaded_file.name # Get path from Gradio file object
        filename = os.path.basename(file_path)
        print(f"Processing '{filename}' (type: {file_type})")

        if file_type == 'NIfTI':
            # Check if the filename ends with either .nii or .nii.gz
            if not (filename.lower().endswith('.nii.gz') or filename.lower().endswith('.nii')):
                raise gr.Error("Invalid NIfTI file. Please upload .nii or .nii.gz")
            
            # Define the destination path (always .nii.gz for consistency)
            dest_path = os.path.join(temp_dir, "uploaded_scan.nii.gz")
            nifti_for_preprocessing_path = dest_path

            # Check if the uploaded file is uncompressed .nii
            if filename.lower().endswith('.nii') and not filename.lower().endswith('.nii.gz'):
                print(f"Detected uncompressed .nii file: {filename}. Compressing to {dest_path}")
                try:
                    # Load the uncompressed .nii file
                    img = nib.load(file_path)
                    # Save it as a compressed .nii.gz file
                    nib.save(img, dest_path)
                    print(f"Successfully compressed and saved to: {dest_path}")
                except Exception as e:
                    raise gr.Error(f"Failed to load or compress .nii file: {e}")
            else:
                # If it's already .nii.gz, just copy it
                print(f"Copying compressed NIfTI {filename} to: {dest_path}")
                shutil.copy(file_path, dest_path)
            
            # nifti_for_preprocessing_path is already set to dest_path
            # print(f"NIfTI path for preprocessing: {nifti_for_preprocessing_path}") # Redundant logging

        elif file_type == 'DICOM (zip)':
            if not filename.endswith('.zip'):
                raise gr.Error("Invalid DICOM file. Please upload a .zip archive.")
            uploaded_zip_path = os.path.join(temp_dir, "dicom_files.zip")
            shutil.copy(file_path, uploaded_zip_path)
            print(f"Copied DICOM zip to: {uploaded_zip_path}")
            dicom_input_dir = os.path.join(temp_dir, "dicom_input")
            nifti_output_dir = os.path.join(temp_dir, "nifti_output")
            os.makedirs(dicom_input_dir, exist_ok=True)
            os.makedirs(nifti_output_dir, exist_ok=True)
            try:
                shutil.unpack_archive(uploaded_zip_path, dicom_input_dir)
                print("Unzip successful.")
            except Exception as e:
                raise gr.Error(f"Error unzipping DICOM file: {e}")
            try:
                dicom2nifti.convert_directory(dicom_input_dir, nifti_output_dir, compression=True, reorient=True)
                nifti_files = [f for f in os.listdir(nifti_output_dir) if f.endswith('.nii.gz')]
                if not nifti_files: raise RuntimeError("dicom2nifti did not produce a .nii.gz file.")
                nifti_for_preprocessing_path = os.path.join(nifti_output_dir, nifti_files[0])
                print(f"DICOM conversion successful. NIfTI: {nifti_for_preprocessing_path}")
            except Exception as e:
                raise gr.Error(f"Error converting DICOM to NIfTI: {e}")
        else:
             raise gr.Error("Invalid file type selected.")

        if not nifti_for_preprocessing_path or not os.path.exists(nifti_for_preprocessing_path):
             raise gr.Error("Could not find the NIfTI file after initial processing.")

        # --- Optional Preprocessing ---
        nifti_to_predict_path = nifti_for_preprocessing_path
        if run_preprocess:
            print("--- Running Optional Preprocessing Pipeline ---")
            try:
                registered_path = os.path.join(temp_dir, "registered.nii.gz")
                register_image(nifti_for_preprocessing_path, registered_path)
                enhanced_path = os.path.join(temp_dir, "enhanced.nii.gz")
                run_enhance_on_file(registered_path, enhanced_path)
                skullstrip_output_dir = os.path.join(temp_dir, "skullstripped")
                skullstripped_path, _ = run_skull_stripping(enhanced_path, skullstrip_output_dir)
                nifti_to_predict_path = skullstripped_path
                print("--- Optional Preprocessing Pipeline Complete ---")
            except Exception as e:
                raise gr.Error(f"Error during preprocessing: {e}")
        else:
             print("--- Skipping Optional Preprocessing Pipeline ---")

        # --- Prediction (Changed for MCI Classification) ---
        input_tensor_5d = preprocess_nifti_for_model(nifti_to_predict_path)
        print("Performing MCI classification prediction...")
        with torch.no_grad():
            try:
                output = model(input_tensor_5d)
                
                # Convert output to probability
                if isinstance(output, torch.Tensor):
                    logit = output.item()
                else:
                    logit = output
                
                # Apply sigmoid to get probability
                probability = torch.sigmoid(torch.tensor(logit)).item()
                predicted_class = 1 if probability > 0.5 else 0
                class_label = "MCI" if predicted_class == 1 else "Healthy Control"
                
                # Create the probability gauge visualization
                probability_gauge = create_probability_gauge(probability)
                
                # Format prediction text for classification
                prediction_text = f"Prediction: {class_label} "
            except Exception as pred_error:
                print(f"Error during prediction: {pred_error}")
                traceback.print_exc()
                raise gr.Error(f"Failed to make prediction: {pred_error}")
        
        print(prediction_text)

        # --- Saliency Map Generation ---
        if generate_saliency_flag:
            print("--- Generating Saliency Data ---")
            try:
                input_3d, saliency_3d = generate_saliency(model, input_tensor_5d)
                if input_3d is not None and saliency_3d is not None:
                    num_slices = input_3d.shape[0]  # Using axial slices now (first dimension)
                    center_slice_index = num_slices // 2

                    # Save numpy arrays to the temp dir for the slider callback
                    unique_id = str(uuid.uuid4())
                    input_array_path = os.path.join(temp_dir, f"{unique_id}_input.npy")
                    saliency_array_path = os.path.join(temp_dir, f"{unique_id}_saliency.npy")
                    np.save(input_array_path, input_3d)
                    np.save(saliency_array_path, saliency_3d)
                    print(f"Saved input array: {input_array_path}")
                    print(f"Saved saliency array: {saliency_array_path}")

                    # Generate initial plots for the center slice
                    input_slice_img, heatmap_slice_img, overlay_slice_img = create_slice_plots(input_3d, saliency_3d, center_slice_index)

                    # Update state for the slider callback
                    saliency_state = {
                        "input_path": input_array_path,
                        "saliency_path": saliency_array_path,
                        "num_slices": num_slices
                    }
                    # Update and show the slider
                    slider_update = gr.Slider(value=center_slice_index, minimum=0, maximum=num_slices - 1, step=1, label="Select Slice", visible=True)
                    print("--- Saliency Generated and Initial Plot Created ---")
                else:
                    error_message = "Saliency map generation failed."
                    print(f"Warning: {error_message}")

            except ImportError as e:
                 error_message = f"Cannot generate saliency: {e}"
                 print(f"Warning: {error_message}")
            except Exception as e:
                error_message = f"Error during saliency processing: {e}"
                traceback.print_exc()
                print(f"Warning: {error_message}")

    except Exception as e:
        print(f"Error in process_scan: {e}")
        traceback.print_exc()
        # Use gr.Warning for non-fatal errors shown to user
        if error_message: # Prepend specific error if available
             gr.Warning(f"{error_message}. General error: {e}")
        else:
             gr.Warning(f"An error occurred: {e}")
        # Return default/error states for outputs
        return "Error during processing", None, None, None, None, gr.Slider(visible=False), {"input_path": None, "saliency_path": None, "num_slices": 0}

    finally:
        # Optional: Schedule cleanup of the temp_dir if files aren't needed long-term
        # Be cautious if files ARE needed by slider state. Gradio might handle this?
        # shutil.rmtree(temp_dir, ignore_errors=True)
        # print(f"Cleaned up temp directory: {temp_dir}") # <--- Defer cleanup
        pass


    # Return results including the probability gauge
    return prediction_text, input_slice_img, heatmap_slice_img, overlay_slice_img, probability_gauge, slider_update, saliency_state

# --- Gradio Slider Update Function ---
def update_slice_viewer(slice_index, current_state):
    input_path = current_state.get("input_path")
    saliency_path = current_state.get("saliency_path")

    if not input_path or not saliency_path or not os.path.exists(input_path) or not os.path.exists(saliency_path):
        print(f"Warning: Cannot update slice viewer. Missing or invalid numpy array paths in state: {current_state}")
        # Return None or placeholder images to indicate error
        return None, None, None

    try:
        input_3d = np.load(input_path)
        saliency_3d = np.load(saliency_path)
        num_slices = input_3d.shape[0]  # Using axial slices (first dimension)

        # Ensure slice_index is valid (Gradio slider should handle bounds, but double-check)
        slice_index = int(slice_index)
        if not (0 <= slice_index < num_slices):
             print(f"Warning: Invalid slice index {slice_index} received by update function.")
             return None, None, None # Or return previous plots?

        # Generate new plots for the selected slice
        input_slice_img, heatmap_slice_img, overlay_slice_img = create_slice_plots(input_3d, saliency_3d, slice_index)

        return input_slice_img, heatmap_slice_img, overlay_slice_img

    except Exception as e:
        print(f"Error updating slice viewer for index {slice_index}: {e}")
        traceback.print_exc()
        # Return None or indicate error
        return None, None, None


# --- Build Gradio Interface ---
with gr.Blocks(css="""
#header-row {
    min-height: 150px;
    align-items: center;
}
.logo-img img {
    height: 150px;
    object-fit: contain;
}
.probability-gauge {
    display: flex;
    justify-content: center;
    margin-top: 1rem;
}
""") as demo:
    # Header Row with Logos and Title
    with gr.Row(elem_id="header-row"):
        with gr.Column(scale=1):
            gr.Image(os.path.join(APP_DIR, "static/images/kannlab.png"),
                     show_label=False, interactive=False,
                     show_download_button=False,
                     container=False,
                     elem_classes=["logo-img"])
        with gr.Column(scale=3):
            gr.Markdown(
                "<h1 style='text-align: center; margin-bottom: 2.5rem'>"
                "BrainIAC: MCI Classification"
                "</h1>"
            )
        with gr.Column(scale=1):
            gr.Image(os.path.join(APP_DIR, "static/images/brainiac.jpeg"),
                     show_label=False, interactive=False,
                     show_download_button=False,
                     container=False,
                     elem_classes=["logo-img"])

    # --- Add model description section ---
    with gr.Accordion("ℹ️ Model Details and Usage Guide", open=False):
        gr.Markdown("""
### 🧠 BrainIAC: MCI Classification

**Model Description**  
A 3D ResNet50 model trained to predict Mild Cognitive Impairment (MCI) from T1-weighted MRI scans.

**Training Dataset**  
- **Subjects**: Trained on T1-weighted MRI scans from subjects with MCI and healthy controls
- **Imaging Modality**: T1-weighted MRI  
- **Preprocessing**: Registration to MNI, N4 bias correction, histogram equalization, skull stripping

**Input**  
- Format: NIfTI or zipped DICOM  
- Required sequence: T1w (3D)  

**Output**  
- Binary classification: MCI or Healthy Control
- Probability score for MCI

**Intended Use**  
- Research use only!

**NOTE**  
- Not validated on T2, FLAIR, DWI or other sequences
- Not validated on pathological cases beyond MCI
- Upload PHI data at own risk!
- The model is hosted on a cloud-based CPU instance.
- The data is not stored, shared or collected for any purpose!

""")

    # Use gr.State to store paths to numpy arrays for the slider callback
    saliency_state = gr.State({"input_path": None, "saliency_path": None, "num_slices": 0})

    # Main Content Row (Controls Left, Output Right)
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown("### Controls")
                file_type = gr.Radio(["NIfTI", "DICOM (zip)"], label="Select Input File Type", value="NIfTI")
                scan_file = gr.File(label="Upload Scan File")
                run_preprocess = gr.Checkbox(label="Run Preprocessing Pipeline ", value=False)
                generate_saliency_checkbox = gr.Checkbox(label="Generate Saliency Maps ", value=True)
                submit_btn = gr.Button("Classify MCI", variant="primary")

        with gr.Column(scale=3):
            with gr.Group():
                gr.Markdown("### Classification Result")
                prediction_output = gr.Label(label="Classification Result")
                # Add the probability gauge visualization
                gr.Markdown("<p style='text-align: center;'>MCI Probability</p>")
                probability_gauge = gr.Image(label="Probability Gauge", type="numpy", show_label=False, elem_classes=["probability-gauge"])

            with gr.Group():
                gr.Markdown("### Saliency Map Viewer (Axial Slice)")
                slice_slider = gr.Slider(label="Select Slice", minimum=0, maximum=0, step=1, value=0, visible=False)
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("<p style='text-align: center;'>Input Slice</p>")
                        input_slice_img = gr.Image(label="Input Slice", type="numpy", show_label=False)
                    with gr.Column():
                        gr.Markdown("<p style='text-align: center;'>Saliency Heatmap</p>")
                        heatmap_slice_img = gr.Image(label="Saliency Heatmap", type="numpy", show_label=False)
                    with gr.Column():
                        gr.Markdown("<p style='text-align: center;'>Overlay</p>")
                        overlay_slice_img = gr.Image(label="Overlay", type="numpy", show_label=False)

    # --- Wire Components ---
    submit_btn.click(
        fn=process_scan,
        inputs=[file_type, scan_file, run_preprocess, generate_saliency_checkbox],
        outputs=[prediction_output, input_slice_img, heatmap_slice_img, overlay_slice_img, probability_gauge, slice_slider, saliency_state]
    )

    slice_slider.change(
        fn=update_slice_viewer,
        inputs=[slice_slider, saliency_state],
        outputs=[input_slice_img, heatmap_slice_img, overlay_slice_img]
    )

# --- Launch the App ---
if __name__ == "__main__":
    if model is None:
        print("ERROR: Model failed to load. Gradio app cannot start.")
    else:
        print("Launching Gradio Interface...")
        demo.launch(server_name="0.0.0.0", server_port=7860, debug=False, share=False)