import os
import torch
import nibabel as nib
from flask import Flask, request, render_template, redirect, url_for, flash, jsonify
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

# Configure Matplotlib for non-GUI backend *before* importing pyplot
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# --- Preprocessing Imports ---
try:
    # Adjust import path based on Docker structure
    # Assumes HD_BET is now at /app/BrainIAC/HD_BET
    from HD_BET.run import run_hd_bet
    # Import MONAI saliency visualizer
    from monai.visualize.gradient_based import GuidedBackpropSmoothGrad
except ImportError as e:
    print(f"Could not import HD_BET or MONAI visualize: {e}. Advanced features might fail.")
    run_hd_bet = None
    GuidedBackpropSmoothGrad = None

# Import necessary components from your existing modules
from model import Backbone, SingleScanModel, Classifier
# Removed: from dataset2 import NormalSynchronizedTransform3D
# Import specific MONAI transforms needed
from monai.transforms import Resized, ScaleIntensityd # Removed ToTensord, will handle manually

app = Flask(__name__)
app.secret_key = 'supersecretkey' # Needed for flashing messages

# --- Constants for Preprocessing ---
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") # Using adult template as default
HD_BET_CONFIG_PATH = os.path.join(APP_DIR, "HD_BET", "config.py")
HD_BET_MODEL_DIR = os.path.join(APP_DIR, "hdbet_model") # Path to copied models

# --- Configuration Loading ---
def load_config():
    # Assuming config.yml is in the same directory as app.py
    config_path = os.path.join(APP_DIR, 'config.yml')
    try:
        with open(config_path, 'r') as file:
            config = yaml.safe_load(file)
        # Add default image_size if not present in config
        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"Error: Configuration file not found at {config_path}")
        # Provide default config or handle error appropriately
        config = {
            'gpu': {'device': 'cpu'},
            'infer': {'checkpoints': 'checkpoints/brainage_model_latest.pt'},
            'data': {'image_size': [128, 128, 128]} # Default image size
        }
    return config

config = load_config()
# Ensure image_size is available, e.g., from config or a default
DEFAULT_IMAGE_SIZE = (128, 128, 128)
image_size_cfg = config.get('data', {}).get('image_size', DEFAULT_IMAGE_SIZE)
# Validate image_size format
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) # Ensure it's a tuple for transforms

# --- Model Loading ---
def load_model(device, checkpoint_path):
    backbone = Backbone()
    classifier = Classifier(d_model=2048) # Make sure d_model matches your trained model
    model = SingleScanModel(backbone, classifier)

    try:
        # Construct absolute path if checkpoint_path is relative
        relative_path = config.get('infer', {}).get('checkpoints', 'checkpoints/brainage_model_latest.pt')
        # Use path relative to app.py location
        checkpoint_path_abs = os.path.join(APP_DIR, relative_path)

        checkpoint = torch.load(checkpoint_path_abs, map_location=device)
        # Adjust key if necessary based on how model was saved
        if 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
        else:
             model.load_state_dict(checkpoint)
        model.to(device)
        model.eval()
        print(f"Model loaded successfully from {checkpoint_path_abs} onto {device}.")
        return model
    except FileNotFoundError:
        print(f"Error: Checkpoint file not found at {checkpoint_path_abs}")
        return None
    except Exception as e:
        print(f"Error loading model checkpoint: {e}")
        traceback.print_exc()
        return None

device = torch.device(config.get('gpu', {}).get('device', 'cpu')) # Default to CPU
model = load_model(device, config) # Pass full config for path finding

# --- Preprocessing Functions from preprocess_utils.py ---
def bias_field_correction(img_array):
    """Performs N4 bias field correction using SimpleITK."""
    image = sitk.GetImageFromArray(img_array)
    # Ensure image is float32 for N4
    if image.GetPixelID() != sitk.sitkFloat32:
        image = sitk.Cast(image, sitk.sitkFloat32)
    maskImage = sitk.OtsuThreshold(image, 0, 1, 200)
    corrector = sitk.N4BiasFieldCorrectionImageFilter()
    numberFittingLevels = 4
    # Define iterations per level more robustly
    max_iters = [min(50 * (2**i), 200) for i in range(numberFittingLevels)]
    corrector.SetMaximumNumberOfIterations(max_iters)
    # Set convergence threshold (optional, can speed up)
    # corrector.SetConvergenceThreshold(1e-6)
    print("  Running N4 Bias Field Correction...")
    corrected_image = corrector.Execute(image, maskImage)
    print("  N4 Correction finished.")
    return sitk.GetArrayFromImage(corrected_image)

def denoise(volume, kernel_size=3):
    """Applies median filter for denoising."""
    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):
    """Rescales intensity after removing background via Otsu."""
    print("  Rescaling intensity...")
    # Ensure input is float for Otsu and calculations
    volume_float = volume.astype(np.float32)
    try:
        t = skimage.filters.threshold_otsu(volume_float, nbins=256)
        print(f"    Otsu threshold found: {t}")
        volume_masked = np.copy(volume_float)
        volume_masked[volume_masked < t] = 0 # Apply mask based on original values
        obj_volume = volume_masked[np.where(volume_masked > 0)]
    except ValueError: # Handle cases with near-uniform intensity
        print("    Otsu failed (likely uniform image), skipping background mask.")
        obj_volume = volume_float.flatten()

    if obj_volume.size == 0:
        print("    Warning: No foreground voxels found after Otsu. Scaling full volume.")
        obj_volume = volume_float.flatten() # Fallback to full volume
        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])

    print(f"    Intensity range used for scaling: [{min_value:.2f}, {max_value:.2f}]")
    # Avoid division by zero if max == min
    denominator = max_value - min_value
    if denominator < 1e-6: denominator = 1e-6

    # Create a copy to modify for output
    output_volume = np.copy(volume_float)
    # Apply scaling only to the object volume identified (or full volume as fallback)
    if bins_num == 0:
        # Scale to 0-1 (float)
        output_volume = (volume_float - min_value) / denominator
        output_volume = np.clip(output_volume, 0.0, 1.0) # Clip results to [0, 1]
    else:
        # Scale and bin
        output_volume = np.round((volume_float - min_value) / denominator * (bins_num - 1))
        output_volume = np.clip(output_volume, 0, bins_num - 1) # Ensure within bin range

    # Ensure output is float32 for consistency
    return output_volume.astype(np.float32)

def equalize_hist(volume, bins_num=256):
    """Performs histogram equalization on non-zero voxels."""
    print("  Performing histogram equalization...")
    # Create a mask of non-zero voxels
    mask = volume > 1e-6 # Use a small epsilon for float comparison
    obj_volume = volume[mask]

    if obj_volume.size == 0:
        print("    Warning: No non-zero voxels found for histogram equalization. Skipping.")
        return volume # Return original volume if no foreground

    # Compute histogram and CDF on the non-zero voxels
    hist, bins = np.histogram(obj_volume, bins_num, range=(obj_volume.min(), obj_volume.max()))
    cdf = hist.cumsum()

    # Normalize CDF
    cdf_normalized = (bins_num - 1) * cdf / float(cdf[-1])

    # Interpolate new values for the object volume
    equalized_obj_volume = np.interp(obj_volume, bins[:-1], cdf_normalized)

    # Create a copy of the original volume to put the results back
    equalized_volume = np.copy(volume)
    equalized_volume[mask] = equalized_obj_volume

    # Ensure output is float32
    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):
    """Full enhancement pipeline from preprocess_utils."""
    print("Starting enhancement pipeline...")
    volume = img_array.astype(np.float32) # Ensure float input
    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}") # Re-raise to stop processing

# --- Registration Function (modified enhance call) ---
def register_image(input_nifti_path, output_nifti_path):
    """Registers input NIfTI to the default template using Elastix."""
    print(f"Registering {input_nifti_path} to {DEFAULT_TEMPLATE_PATH}")
    if not os.path.exists(PARAMS_RIGID_PATH):
        raise FileNotFoundError(f"Elastix parameter file not found at {PARAMS_RIGID_PATH}")
    if not os.path.exists(DEFAULT_TEMPLATE_PATH):
        raise FileNotFoundError(f"Default template file not found at {DEFAULT_TEMPLATE_PATH}")

    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 # Keep console clean
    )
    itk.imwrite(result_image, output_nifti_path)
    print(f"Registration output saved to {output_nifti_path}")

# --- Enhanced Image Function (calls actual enhance) ---
def run_enhance_on_file(input_nifti_path, output_nifti_path):
    """Reads NIfTI, runs enhance pipeline, saves NIfTI."""
    print(f"Running full enhancement on {input_nifti_path}")
    img_sitk = sitk.ReadImage(input_nifti_path)
    img_array = sitk.GetArrayFromImage(img_sitk)

    # Run the actual enhancement pipeline
    enhanced_array = enhance(img_array, run_bias_correction=True) # Assuming N4 is desired

    enhanced_img_sitk = sitk.GetImageFromArray(enhanced_array)
    enhanced_img_sitk.CopyInformation(img_sitk) # Preserve metadata
    sitk.WriteImage(enhanced_img_sitk, output_nifti_path)
    print(f"Enhanced image saved to {output_nifti_path}")

# --- Skull Stripping Function (Set Environment Variable) ---
def run_skull_stripping(input_nifti_path, output_dir):
    """Runs HD-BET skull stripping."""
    print(f"Running HD-BET skull stripping on {input_nifti_path}")
    if run_hd_bet is None:
        raise RuntimeError("HD-BET module could not be imported. Cannot perform skull stripping.")
    
    # Removed environment variable setting as path is fixed in HD_BET/paths.py
    # # Set environment variable *before* calling run_hd_bet
    # # Ensure the target directory exists
    # if not os.path.isdir(HD_BET_MODEL_DIR):
    #     raise FileNotFoundError(f"HD-BET model directory not found at specified path: {HD_BET_MODEL_DIR}")
    # print(f"Setting HD_BET_MODELS environment variable to: {HD_BET_MODEL_DIR}")
    # os.environ['HD_BET_MODELS'] = HD_BET_MODEL_DIR

    # Check config path
    if not os.path.exists(HD_BET_CONFIG_PATH):
        alt_config_path = os.path.join(APP_DIR, "HD_BET", "HD_BET", "config.py")
        if os.path.exists(alt_config_path):
             print(f"Warning: Using alternative HD-BET config path: {alt_config_path}")
             config_to_use = alt_config_path
        else:
            raise FileNotFoundError(f"HD-BET config file not found at {HD_BET_CONFIG_PATH} or {alt_config_path}")
    else:
        config_to_use = HD_BET_CONFIG_PATH

    # Define output paths
    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")

    # Make sure output directory exists
    os.makedirs(output_dir, exist_ok=True)

    # Run HD-BET
    run_hd_bet(input_nifti_path, output_file_path,
               mode="fast",
               device='cpu',
               config_file=config_to_use,
               postprocess=False,
               do_tta=False,
               keep_mask=True,
               overwrite=True)
    
    # Unset environment variable after use (optional, good practice)
    # del os.environ['HD_BET_MODELS'] 

    if not os.path.exists(output_file_path):
        raise RuntimeError(f"HD-BET did not produce the expected output file: {output_file_path}")

    print(f"Skull stripping output saved to {output_file_path}")
    return output_file_path, output_mask_path

# --- Image Preprocessing ---
# Define necessary MONAI transforms directly
# Keys must match the dictionary keys we create later ('image')
resize_transform = Resized(keys=["image"], spatial_size=image_size)
scale_transform = ScaleIntensityd(keys=["image"], minv=0.0, maxv=1.0)

def preprocess_nifti(nifti_path):
    """Loads and preprocesses a NIfTI file, returning a 5D tensor."""
    print(f"Preprocessing NIfTI: {nifti_path}")
    scan_data = nib.load(nifti_path).get_fdata()
    print(f"  Loaded scan data shape: {scan_data.shape}")
    scan_tensor = torch.tensor(scan_data, dtype=torch.float32).unsqueeze(0) # Add C dim
    print(f"  Shape after tensor+channel: {scan_tensor.shape}")
    sample = {"image": scan_tensor}
    sample_resized = resize_transform(sample)
    print(f"  Shape after resize: {sample_resized['image'].shape}")
    sample_scaled = scale_transform(sample_resized)
    print(f"  Shape after scaling: {sample_scaled['image'].shape}")
    input_tensor = sample_scaled["image"].unsqueeze(0).to(device) # Add B dim
    print(f"  Final shape for model: {input_tensor.shape}")
    if input_tensor.dim() != 5:
        raise ValueError(f"Preprocessing resulted in incorrect shape: {input_tensor.shape}. Expected 5D.")
    return input_tensor

# --- Final NIfTI Preprocessing for Model ---
def preprocess_nifti_for_model(nifti_path):
    """Loads final NIfTI and prepares 5D tensor for the model."""
    # ... (Same as previous preprocess_nifti function) ...
    print(f"Preprocessing NIfTI for model: {nifti_path}")
    scan_data = nib.load(nifti_path).get_fdata()
    print(f"  Loaded scan data shape: {scan_data.shape}")
    scan_tensor = torch.tensor(scan_data, dtype=torch.float32).unsqueeze(0) # Add C dim
    print(f"  Shape after tensor+channel: {scan_tensor.shape}")
    sample = {"image": scan_tensor}
    sample_resized = resize_transform(sample)
    print(f"  Shape after resize: {sample_resized['image'].shape}")
    sample_scaled = scale_transform(sample_resized)
    print(f"  Shape after scaling: {sample_scaled['image'].shape}")
    input_tensor = sample_scaled["image"].unsqueeze(0).to(device) # Add B dim
    print(f"  Final shape for model: {input_tensor.shape}")
    if input_tensor.dim() != 5:
        raise ValueError(f"Preprocessing resulted in incorrect shape: {input_tensor.shape}. Expected 5D.")
    return input_tensor

# --- Saliency Map Generation ---
def generate_saliency(model, input_tensor_5d):
    """Generates saliency map using GuidedBackpropSmoothGrad."""
    if GuidedBackpropSmoothGrad is None:
        raise ImportError("MONAI visualize components not imported. Cannot generate saliency map.")
    if model is None:
         raise ValueError("Model not loaded. Cannot generate saliency map.")
    
    print("Generating saliency map...")
    input_tensor_5d.requires_grad_(True)
    # Use the backbone for saliency as in the original script
    # Ensure model and backbone are on the correct device (CPU in this case)
    visualizer = GuidedBackpropSmoothGrad(model=model.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))
        print("Saliency map generated.")
        
        # Detach, move to CPU, remove Batch and Channel dims for processing/plotting -> (D, H, W)
        input_3d = input_tensor_5d.squeeze().cpu().detach().numpy()
        saliency_3d = saliency_map_5d.squeeze().cpu().detach().numpy()
        
        return input_3d, saliency_3d
    
    except Exception as e:
        print(f"Error during saliency map generation: {e}")
        traceback.print_exc()
        # Return None or empty arrays if generation fails
        return None, None
    finally:
        # Ensure requires_grad is turned off if it was modified
        input_tensor_5d.requires_grad_(False)

# --- Plotting Function for Single Slice ---
def create_plot_images_for_slice(mri_data_3d, saliency_data_3d, slice_index):
    """Creates base64 encoded PNGs for a specific axial slice index."""
    print(f"  Generating plots for slice index: {slice_index}")
    if mri_data_3d is None or saliency_data_3d is None:
        print("    Input or Saliency data is None, cannot generate plot.")
        return None
    if slice_index < 0 or slice_index >= mri_data_3d.shape[2]:
        print(f"    Error: Slice index {slice_index} out of bounds (0-{mri_data_3d.shape[2]-1}).")
        return None

    # Function to save plot to base64 string (copied from previous version)
    def save_plot_to_base64(fig):
        buf = io.BytesIO()
        fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=75)
        plt.close(fig) # Close the figure immediately
        buf.seek(0)
        img_str = base64.b64encode(buf.read()).decode('utf-8')
        buf.close()
        return img_str

    try:
        mri_slice = mri_data_3d[:, :, slice_index]
        saliency_slice_orig = saliency_data_3d[:, :, slice_index]

        # --- Normalize MRI Slice (using volume stats if available, otherwise slice stats) ---
        # For consistency, ideally pass volume stats, but recalculating per slice is fallback
        p1_vol, p99_vol = np.percentile(mri_data_3d, (1, 99))
        mri_norm_denom = p99_vol - p1_vol
        if mri_norm_denom < 1e-6: mri_norm_denom = 1e-6
        mri_slice_norm = np.clip(mri_slice, p1_vol, p99_vol)
        mri_slice_norm = (mri_slice_norm - p1_vol) / mri_norm_denom

        # --- Process Saliency Slice ---
        saliency_slice = np.copy(saliency_slice_orig)
        saliency_slice[saliency_slice < 0] = 0 # Ensure non-negative
        saliency_slice_blurred = cv2.GaussianBlur(saliency_slice, (15, 15), 0)
        # Use volume max for normalization if possible, fallback to slice max
        s_max_vol = np.max(saliency_data_3d[saliency_data_3d >= 0]) # Max of non-negative values in volume
        if s_max_vol < 1e-6: s_max_vol = 1e-6
        # --- Add logging for the calculated global max ---
        print(f"    Calculated Global Max Saliency (s_max_vol) for normalization: {s_max_vol:.4f}")
        # --------------------------------------------------
        saliency_slice_norm = saliency_slice_blurred / s_max_vol
        threshold_value = 0.0
        saliency_slice_thresholded = np.where(saliency_slice_norm > threshold_value, saliency_slice_norm, 0)

        # --- Generate Plots ---
        slice_plots = {}

        # Plot 1: Input Slice
        fig1, ax1 = plt.subplots(figsize=(3, 3))
        ax1.imshow(mri_slice_norm, cmap='gray', interpolation='none', origin='lower')
        ax1.axis('off')
        slice_plots['input_slice'] = save_plot_to_base64(fig1)

        # Plot 2: Saliency Heatmap
        fig2, ax2 = plt.subplots(figsize=(3, 3))
        ax2.imshow(saliency_slice_thresholded, cmap='magma', interpolation='none', origin='lower')
        ax2.axis('off')
        slice_plots['heatmap_slice'] = save_plot_to_base64(fig2)

        # Plot 3: Overlay
        fig3, ax3 = plt.subplots(figsize=(3, 3))
        ax3.imshow(mri_slice_norm, cmap='gray', interpolation='none', origin='lower')
        if np.max(saliency_slice_thresholded) > 0:
             # Remove fixed levels to let contour auto-determine levels based on slice data
             ax3.contour(saliency_slice_thresholded, cmap='magma', origin='lower', linewidths=1.0)
        ax3.axis('off')
        slice_plots['overlay_slice'] = save_plot_to_base64(fig3)

        print(f"    Generated plots successfully for slice {slice_index}.")
        return slice_plots

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

# --- Flask Routes ---
@app.route('/', methods=['GET'])
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if model is None:
        flash('Model not loaded. Cannot perform prediction.', 'error')
        return redirect(url_for('index'))

    # Get form data
    file_type = request.form.get('file_type')
    run_preprocess_flag = request.form.get('preprocess') == 'yes'
    generate_saliency_flag = request.form.get('generate_saliency') == 'yes' # Get saliency flag
    file = request.files.get('scan_file')

    # --- Basic Input Validation ---
    if not file_type:
        flash('Please select a file type (NIfTI or DICOM).', 'error')
        return redirect(url_for('index'))
    if not file or file.filename == '':
        flash('No scan file selected', 'error')
        return redirect(url_for('index'))

    print(f"Received upload: type='{file_type}', filename='{file.filename}', preprocess={run_preprocess_flag}, saliency={generate_saliency_flag}")

    # --- Setup Temporary Directory --- 
    # temp_dir_obj = tempfile.TemporaryDirectory() # <--- PROBLEM: Cleans up automatically
    # Use mkdtemp to create a persistent temporary directory
    # NOTE: Requires a manual cleanup strategy later!
    try:
        temp_dir = tempfile.mkdtemp() 
    except Exception as e:
        print(f"Error creating temporary directory: {e}")
        flash("Server error: Could not create temporary directory.", "error")
        return redirect(url_for('index'))
        
    # Generate a unique ID based on the temp directory name
    unique_id = os.path.basename(temp_dir) 
    print(f"Created persistent temp directory: {temp_dir} (ID: {unique_id})")
    nifti_for_preprocessing_path = None # Path to the NIfTI before optional preprocessing

    try:
        # --- Handle Upload and DICOM Conversion ---
        # --- Handle NIfTI Upload ---
        if file_type == 'nifti':
            if not file.filename.endswith('.nii.gz'):
                flash('Invalid file type for NIfTI selection. Please upload .nii.gz', 'error')
                # temp_dir_obj.cleanup() # No object to cleanup, need manual rmtree
                shutil.rmtree(temp_dir, ignore_errors=True) 
                return redirect(url_for('index'))
            uploaded_file_path = os.path.join(temp_dir, "uploaded_scan.nii.gz")
            file.save(uploaded_file_path)
            print(f"Saved uploaded NIfTI file to: {uploaded_file_path}")
            nifti_for_preprocessing_path = uploaded_file_path

        # --- Handle DICOM Upload ---
        elif file_type == 'dicom':
            if not file.filename.endswith('.zip'):
                flash('Invalid file type for DICOM selection. Please upload a .zip file.', 'error')
                # temp_dir_obj.cleanup()
                shutil.rmtree(temp_dir, ignore_errors=True) 
                return redirect(url_for('index'))
            uploaded_zip_path = os.path.join(temp_dir, "dicom_files.zip")
            file.save(uploaded_zip_path)
            print(f"Saved uploaded 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:
                # Use shutil.unpack_archive for better cross-platform compatibility potentially
                shutil.unpack_archive(uploaded_zip_path, dicom_input_dir)
                print(f"Unzip successful.")
            except Exception as e:
                print(f"Unzip failed: {e}")
                flash(f'Error unzipping DICOM file: {e}', 'error')
                # temp_dir_obj.cleanup()
                shutil.rmtree(temp_dir, ignore_errors=True) 
                return redirect(url_for('index'))
            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 file: {nifti_for_preprocessing_path}")
            except Exception as e:
                print(f"DICOM to NIfTI conversion failed: {e}")
                flash(f'Error converting DICOM to NIfTI: {e}', 'error')
                # temp_dir_obj.cleanup()
                shutil.rmtree(temp_dir, ignore_errors=True) 
                return redirect(url_for('index'))
        else:
            flash('Invalid file type selected.', 'error')
            # temp_dir_obj.cleanup()
            shutil.rmtree(temp_dir, ignore_errors=True) 
            return redirect(url_for('index'))
        
        if not nifti_for_preprocessing_path or not os.path.exists(nifti_for_preprocessing_path):
             flash('Error: Could not find the NIfTI file for processing.', 'error')
             # temp_dir_obj.cleanup()
             shutil.rmtree(temp_dir, ignore_errors=True) 
             return redirect(url_for('index'))

        # --- Optional Preprocessing Steps ---
        nifti_to_predict_path = nifti_for_preprocessing_path
        if run_preprocess_flag:
            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:
                print(f"Error during optional preprocessing pipeline: {e}")
                traceback.print_exc()
                flash(f'Error during preprocessing: {e}', 'error')
                # temp_dir_obj.cleanup()
                shutil.rmtree(temp_dir, ignore_errors=True) 
                return redirect(url_for('index'))
        else:
             print("--- Skipping Optional Preprocessing Pipeline ---")

        # --- Final Preprocessing for Model & Prediction ---
        input_tensor_5d = preprocess_nifti_for_model(nifti_to_predict_path)
        print("Performing prediction...")
        with torch.no_grad():
            output = model(input_tensor_5d)
            predicted_age = output.item()
            predicted_age_years = predicted_age / 12 # Adjust if needed
        print(f"Prediction successful: {predicted_age_years:.2f} years")

        # --- Saliency Data Handling (Generate, Save, Get Initial Plot) ---
        saliency_output_for_template = None # Initialize
        if generate_saliency_flag:
            print("--- Generating & Saving Saliency Data ---")
            try:
                input_3d_for_plot, saliency_3d = generate_saliency(model, input_tensor_5d)
                
                if input_3d_for_plot is not None and saliency_3d is not None:
                    num_slices = input_3d_for_plot.shape[2]
                    center_slice_index = num_slices // 2

                    # Save the numpy arrays for the dynamic route
                    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_for_plot)
                    np.save(saliency_array_path, saliency_3d)
                    print(f"Saved input array to {input_array_path}")
                    print(f"Saved saliency array to {saliency_array_path}")

                    # Generate ONLY the center slice plots for the initial view
                    center_slice_plots = create_plot_images_for_slice(input_3d_for_plot, saliency_3d, center_slice_index)

                    if center_slice_plots:
                        # Prepare data structure for the template
                        saliency_output_for_template = {
                            'center_slice_plots': center_slice_plots,
                            'num_slices': num_slices,
                            'center_slice_index': center_slice_index,
                            'unique_id': unique_id, # Pass the ID for filenames
                            'temp_dir_path': temp_dir # Pass the full path for lookup
                        }
                        print("--- Saliency Data Saved & Initial Plot Generated ---")
                    else:
                         print("--- Center Slice Plotting Failed ---")
                         flash('Failed to generate initial saliency plot.', 'warning')
                else:
                     print("--- Saliency Generation Failed --- ")
                     flash('Saliency map generation failed.', 'warning')

            except Exception as e:
                print(f"Error during saliency processing/saving: {e}")
                traceback.print_exc()
                flash('Could not generate or save saliency maps due to an error.', 'error')

        # Render result, passing prediction and potentially the NEW saliency structure
        return render_template('index.html', 
                               prediction=f"{predicted_age_years:.2f} years", 
                               saliency_info=saliency_output_for_template) # Pass the new dict

    except Exception as e:
        flash(f'Error processing file: {e}', 'error')
        print(f"Caught Exception during prediction process: {e}")
        traceback.print_exc()
        # Ensure cleanup happens even if exception occurs mid-process
        # temp_dir_obj.cleanup()
        if temp_dir and os.path.exists(temp_dir):
            shutil.rmtree(temp_dir, ignore_errors=True) # Manual cleanup on general error
        return redirect(url_for('index'))

    # NOTE: Temporary directory created with mkdtemp is NOT automatically cleaned.
    # Need a separate mechanism (e.g., cron job, background task) to remove old directories
    # from the system's temporary location (e.g., /tmp) based on age.
    # Leaving the directory here so /get_slice can access the files.

# --- New Route for Dynamic Slice Loading ---
@app.route('/get_slice/<unique_id>/<int:slice_index>')
def get_slice(unique_id, slice_index):
    # Get the actual temporary directory path from query parameter
    temp_dir_path = request.args.get('path')
    if not temp_dir_path:
        print("Error: 'path' query parameter missing in /get_slice request")
        return jsonify({"error": "Required path information missing."}), 400
    
    # Construct paths using the provided directory path and unique ID
    input_array_path = os.path.join(temp_dir_path, f"{unique_id}_input.npy")
    saliency_array_path = os.path.join(temp_dir_path, f"{unique_id}_saliency.npy")
    print(f"Attempting to load slice {slice_index} for ID {unique_id} from actual path: {temp_dir_path}")

    try:
        # Check using the exact paths constructed above
        if not os.path.exists(input_array_path) or not os.path.exists(saliency_array_path):
             print(f"Error: .npy files not found for ID {unique_id} at {temp_dir_path}")
             return jsonify({"error": "Saliency data not found. It might have expired or failed to save."}), 404
        
        input_3d = np.load(input_array_path)
        saliency_3d = np.load(saliency_array_path)
        print(f"Loaded arrays for ID {unique_id}. Input shape: {input_3d.shape}, Saliency shape: {saliency_3d.shape}")

        # Generate plots for the requested slice using the helper function
        slice_plots = create_plot_images_for_slice(input_3d, saliency_3d, slice_index)

        if slice_plots:
            return jsonify(slice_plots) # Return plot data as JSON
        else:
            return jsonify({"error": f"Failed to generate plots for slice {slice_index}."}), 500

    except Exception as e:
        print(f"Error in /get_slice for ID {unique_id}, slice {slice_index}: {e}")
        traceback.print_exc()
        return jsonify({"error": "An internal error occurred while fetching the slice data."}), 500

if __name__ == '__main__':
    # Use '0.0.0.0' to make it accessible outside the container
    app.run(host='0.0.0.0', port=5000, debug=False) # Turn off debug for production/docker