src/data/preprocess_dataset.py

#!/usr/bin/env python3
"""
LUNA16 Preprocessing Pipeline
Phase 1, Week 2: Patch Extraction and Preprocessing

This script implements memory-efficient preprocessing for the DCA-Net model:
1. Nodule patch extraction (64x64x64)
2. Context patch extraction (96x96x96 → 48x48x48)
3. HU windowing and normalization
4. Data balancing strategy
5. Train/Val/Test split

Hardware-aware: Processes one scan at a time for 8GB RAM constraint.
"""

import os
import numpy as np
import pandas as pd
import SimpleITK as sitk
from pathlib import Path
from tqdm import tqdm
import json
import logging
from datetime import datetime
from scipy.ndimage import zoom
from concurrent.futures import ProcessPoolExecutor, as_completed
import multiprocessing
import psutil
import gc
import signal
import sys

# ============== Configuration ==============
DATA_DIR = Path("data")
OUTPUT_DIR = Path("preprocessed_data")
NODULE_PATCH_SIZE = 64  # 64x64x64 cube
CONTEXT_PATCH_SIZE = 96  # Extract 96x96x96, then downsample
CONTEXT_DOWNSAMPLE = 48  # Downsample to 48x48x48

# HU windowing for lung tissue
HU_MIN = -1000
HU_MAX = 400

# Data balancing: 1 positive : 7 negatives (from roadmap)
HARD_NEG_RATIO = 3  # Hard negatives per positive (reduced from 5 — more generalizable)
RANDOM_NEG_RATIO = 4  # Random negatives per positive (increased from 2 — better diversity)

RANDOM_SEED = 42
np.random.seed(RANDOM_SEED)

# Parallel processing configuration
# Use 32 of available CPU cores (leave some for system)
NUM_WORKERS = min(32, max(1, multiprocessing.cpu_count() - 4))
BATCH_SIZE = 64  # Process in batches for efficient checkpointing


# ─────────────────────────────────────────────────────────────
# Display
# ─────────────────────────────────────────────────────────────
RED = "\033[91m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
BLUE = "\033[94m"
CYAN = "\033[96m"
BOLD = "\033[1m"
DIM = "\033[2m"
RESET = "\033[0m"


def banner():
    print(f"""
{BOLD}{CYAN}╔════════════════════════════════════════════════════════════════════╗
║                                                                    ║
║    ██████╗ ███╗   ██╗ ██████╗ ██████╗                              ║
║   ██╔═══██╗████╗  ██║██╔════╝██╔═══██╗                             ║
║   ██║   ██║██╔██╗ ██║██║     ██║   ██║                             ║
║   ██║   ██║██║╚██╗██║██║     ██║   ██║                             ║
║   ╚██████╔╝██║ ╚████║╚██████╗╚██████╔╝                             ║
║    ╚═════╝ ╚═╝  ╚═══╝ ╚═════╝ ╚═════╝                              ║
║   ██╗   ██╗██╗███████╗██╗ ██████╗ ███╗   ██╗    ██╗  ██╗           ║
║   ██║   ██║██║██╔════╝██║██╔═══██╗████╗  ██║    ╚██╗██╔╝           ║
║   ██║   ██║██║███████╗██║██║   ██║██╔██╗ ██║     ╚███╔╝            ║
║   ╚██╗ ██╔╝██║╚════██║██║██║   ██║██║╚██╗██║     ██╔██╗            ║
║    ╚████╔╝ ██║███████║██║╚██████╔╝██║ ╚████║    ██╔╝ ██╗           ║
║     ╚═══╝  ╚═╝╚══════╝╚═╝ ╚═════╝ ╚═╝  ╚═══╝    ╚═╝  ╚═╝           ║
║                                                                    ║
║   Dual-Context Attention Network                                   ║
║   AI-Powered Lung Cancer Detection — Data Preprocessing            ║
║                                                                    ║
╚════════════════════════════════════════════════════════════════════╝{RESET}
""")


# ─────────────────────────────────────────────────────────────
# Logging Setup
# ─────────────────────────────────────────────────────────────
logger = logging.getLogger('preprocessing')

def setup_logging():
    """Setup dual logging: console (colored) + file (clean text) in logs/ folder."""
    log_dir = Path("logs")
    log_dir.mkdir(parents=True, exist_ok=True)
    
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    log_file = log_dir / f"preprocessing_{timestamp}.log"
    
    logger.setLevel(logging.DEBUG)
    
    # File handler — clean text (no ANSI colors)
    fh = logging.FileHandler(log_file, encoding='utf-8')
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(logging.Formatter('%(asctime)s | %(levelname)-8s | %(message)s', datefmt='%Y-%m-%d %H:%M:%S'))
    logger.addHandler(fh)
    
    # Console handler — minimal (print() handles colored console output)
    ch = logging.StreamHandler(sys.stdout)
    ch.setLevel(logging.WARNING)  # Only warnings+ to avoid double-printing
    ch.setFormatter(logging.Formatter('%(message)s'))
    logger.addHandler(ch)
    
    logger.info(f"Preprocessing log started: {log_file}")
    print(f"  {DIM}Log file:{RESET} {log_file}")
    
    return log_file


def _strip_ansi(text):
    """Remove ANSI escape codes for clean log file output."""
    import re
    return re.sub(r'\033\[[0-9;]*m', '', text)


def section(title):
    print(f"\n{BOLD}{BLUE}{'─' * 60}")
    print(f"  {title}")
    print(f"{'─' * 60}{RESET}")
    logger.info(f"{'─' * 60}")
    logger.info(f"  {title}")
    logger.info(f"{'─' * 60}")


def info(label, value):
    print(f"  {DIM}{label}:{RESET} {value}")
    logger.info(f"  {label}: {_strip_ansi(str(value))}")


def success(msg):
    print(f"  {GREEN}✓ {msg}{RESET}")
    logger.info(f"  ✓ {msg}")


def warn(msg):
    print(f"  {YELLOW}! {msg}{RESET}")
    logger.warning(f"  ! {msg}")


def get_scan_paths():
    """Get all .mhd scan file paths organized by subset (supports 0-9)."""
    scan_paths = {}
    for i in range(10):  # Changed from range(5) to support full LUNA16 (0-9)
        subset_dir = DATA_DIR / f"subset{i}" / f"subset{i}"
        if subset_dir.exists():
            mhd_files = list(subset_dir.glob("*.mhd"))
            if len(mhd_files) > 0:  # Only add if scans found
                scan_paths[f"subset{i}"] = mhd_files
    return scan_paths


def load_scan(mhd_path):
    """Load a CT scan from .mhd file."""
    img = sitk.ReadImage(str(mhd_path))
    arr = sitk.GetArrayFromImage(img)  # (z, y, x)
    
    # Get metadata for coordinate conversion
    origin = np.array(img.GetOrigin())  # (x, y, z)
    spacing = np.array(img.GetSpacing())  # (x, y, z)
    direction = np.array(img.GetDirection()).reshape(3, 3)
    
    return arr, origin, spacing, direction


def world_to_voxel(world_coord, origin, spacing, direction=None):
    """Convert world coordinates (mm) to voxel indices."""
    # Simple conversion (assuming identity direction matrix)
    voxel_coord = (world_coord - origin) / spacing
    return np.round(voxel_coord).astype(int)


def apply_hu_windowing(arr, hu_min=HU_MIN, hu_max=HU_MAX):
    """Apply HU windowing and normalize to [-1, 1]."""
    arr = np.clip(arr, hu_min, hu_max)
    arr = (arr - hu_min) / (hu_max - hu_min)  # [0, 1]
    arr = arr * 2 - 1  # [-1, 1]
    return arr.astype(np.float32)


def extract_patch(arr, center, patch_size):
    """Extract a cubic patch centered at the given voxel coordinate.
    
    Args:
        arr: 3D numpy array (z, y, x)
        center: (x, y, z) voxel coordinates
        patch_size: Size of the cubic patch
    
    Returns:
        patch: Extracted patch or None if out of bounds
    """
    half = patch_size // 2
    cx, cy, cz = center
    
    # Compute bounds (note: arr is z, y, x)
    z_start, z_end = cz - half, cz + half
    y_start, y_end = cy - half, cy + half
    x_start, x_end = cx - half, cx + half
    
    # Check bounds
    if (z_start < 0 or z_end > arr.shape[0] or
        y_start < 0 or y_end > arr.shape[1] or
        x_start < 0 or x_end > arr.shape[2]):
        return None
    
    patch = arr[z_start:z_end, y_start:y_end, x_start:x_end]
    assert patch.shape == (patch_size, patch_size, patch_size), f"Invalid patch shape: {patch.shape}"
    return patch


def downsample_patch(patch, target_size):
    """Downsample a 3D patch using scipy zoom."""
    if patch is None:
        return None
    zoom_factor = target_size / patch.shape[0]
    return zoom(patch, zoom_factor, order=1)  # Linear interpolation


def process_scan(mhd_path, annotations_df, candidates_df, output_dir):
    """Process a single CT scan and extract patches for all candidates.
    
    Args:
        mhd_path: Path to .mhd file
        annotations_df: DataFrame with ground truth nodules
        candidates_df: DataFrame with all candidates
        output_dir: Directory to save patches
    
    Returns:
        results: List of dicts with patch metadata
        stats: Dict with extraction statistics
    """
    series_uid = mhd_path.stem
    
    # Get candidates for this scan
    scan_candidates = candidates_df[candidates_df['seriesuid'] == series_uid]
    
    if len(scan_candidates) == 0:
        return [], {'total': 0, 'success': 0, 'failures': {
            'nodule_out_of_bounds': 0, 'context_out_of_bounds': 0, 'load_error': 0
        }}
    
    # Initialize statistics
    stats = {
        'total': len(scan_candidates),
        'success': 0,
        'failures': {
            'nodule_out_of_bounds': 0,
            'context_out_of_bounds': 0,
            'load_error': 0
        }
    }
    
    # Load the scan
    try:
        arr, origin, spacing, direction = load_scan(mhd_path)
    except Exception as e:
        warn(f"Error loading {mhd_path}: {e}")
        stats['failures']['load_error'] = len(scan_candidates)
        return [], stats
    
    # Apply HU windowing
    arr_windowed = apply_hu_windowing(arr)
    
    # Get annotations for this scan (for hard negative identification)
    scan_annotations = annotations_df[annotations_df['seriesuid'] == series_uid]
    nodule_positions = []
    for _, row in scan_annotations.iterrows():
        world_coord = np.array([row['coordX'], row['coordY'], row['coordZ']])
        voxel_coord = world_to_voxel(world_coord, origin, spacing)
        nodule_positions.append({
            'voxel': voxel_coord,
            'diameter_mm': row['diameter_mm']
        })
    
    results = []
    
    for idx, row in scan_candidates.iterrows():
        world_coord = np.array([row['coordX'], row['coordY'], row['coordZ']])
        voxel_coord = world_to_voxel(world_coord, origin, spacing)
        label = row['class']
        
        # Extract nodule patch (64x64x64)
        nodule_patch = extract_patch(arr_windowed, voxel_coord, NODULE_PATCH_SIZE)
        if nodule_patch is None:
            stats['failures']['nodule_out_of_bounds'] += 1
            continue
        
        # Extract context patch (96x96x96 → 48x48x48)
        # FIX: Skip sample if context is out of bounds — don't create fake zero data
        context_patch = extract_patch(arr_windowed, voxel_coord, CONTEXT_PATCH_SIZE)
        if context_patch is None:
            stats['failures']['context_out_of_bounds'] += 1
            continue  # Don't create fake data with zero-padding
        
        context_patch = downsample_patch(context_patch, CONTEXT_DOWNSAMPLE)
        
        # Determine if this is a hard negative (close to a real nodule)
        # FIX: Use 1.5× diameter with average spacing (research-backed)
        is_hard_negative = False
        if label == 0 and len(nodule_positions) > 0:
            avg_spacing = np.mean(spacing)  # Average across x,y,z
            for nodule in nodule_positions:
                dist = np.linalg.norm(voxel_coord - nodule['voxel'])
                # Consider "hard" if within 2.0x nodule diameter (in voxels) — research-backed
                threshold_voxels = (nodule['diameter_mm'] * 2.0) / avg_spacing
                if dist < threshold_voxels:
                    is_hard_negative = True
                    break
        
        # Save patches
        patch_id = f"{series_uid}_{idx}"
        nodule_path = output_dir / "nodule_patches" / f"{patch_id}.npz"
        context_path = output_dir / "context_patches" / f"{patch_id}.npz"
        
        np.savez_compressed(nodule_path, patch=nodule_patch)
        np.savez_compressed(context_path, patch=context_patch)
        
        stats['success'] += 1
        results.append({
            'patch_id': patch_id,
            'series_uid': series_uid,
            'label': label,
            'is_hard_negative': is_hard_negative,
            'voxel_x': int(voxel_coord[0]),
            'voxel_y': int(voxel_coord[1]),
            'voxel_z': int(voxel_coord[2]),
            'nodule_path': str(nodule_path),
            'context_path': str(context_path)
        })
    
    # Log statistics if significant failures
    total_failures = stats['failures']['nodule_out_of_bounds'] + stats['failures']['context_out_of_bounds']
    if total_failures > 0:
        failure_rate = total_failures / stats['total'] * 100
        if failure_rate > 10:  # Warn if >10% failures
            warn(f"{series_uid}: {failure_rate:.1f}% extraction failures ({total_failures}/{stats['total']})")
    
    return results, stats


def balance_samples(metadata_df, pos_to_neg_ratio=7):
    """Balance positive and negative samples according to strategy.
    
    Strategy from roadmap:
    - Keep all positive samples
    - Hard negatives: 5x positives
    - Random negatives: 2x positives
    - Total ratio: 1:7
    """
    positives = metadata_df[metadata_df['label'] == 1]
    negatives = metadata_df[metadata_df['label'] == 0]
    hard_negatives = negatives[negatives['is_hard_negative'] == True]
    easy_negatives = negatives[negatives['is_hard_negative'] == False]
    
    n_pos = len(positives)
    n_hard = min(len(hard_negatives), n_pos * HARD_NEG_RATIO)
    n_easy = min(len(easy_negatives), n_pos * RANDOM_NEG_RATIO)
    
    # Sample negatives
    sampled_hard = hard_negatives.sample(n=n_hard, random_state=RANDOM_SEED) if n_hard > 0 else hard_negatives
    sampled_easy = easy_negatives.sample(n=n_easy, random_state=RANDOM_SEED) if n_easy > 0 else easy_negatives
    
    # Combine
    balanced_df = pd.concat([positives, sampled_hard, sampled_easy], ignore_index=True)
    balanced_df = balanced_df.sample(frac=1, random_state=RANDOM_SEED)  # Shuffle
    
    return balanced_df


def create_splits(metadata_df, scan_paths):
    """Create train/val/test splits based on available subsets.
    
    Split strategy (adaptive):
    - 10 subsets (full LUNA16):  Train subset0-7 (80%), Val subset8 (10%), Test subset9 (10%)
    - 5 subsets (partial):       Train subset0-2 (60%), Val subset3 (20%), Test subset4 (20%)
    - Other counts:             Train 60%, Val 20%, Test 20% by index
    """
    train_series = set()
    val_series = set()
    test_series = set()
    
    available_subsets = sorted(scan_paths.keys())  # e.g. ['subset0', 'subset1', ...]
    num_subsets = len(available_subsets)
    
    for subset_name, paths in scan_paths.items():
        for path in paths:
            series_uid = path.stem
            
            if num_subsets >= 10:
                # Full LUNA16 (10 subsets) — standard split
                if subset_name in ['subset0', 'subset1', 'subset2', 'subset3',
                                   'subset4', 'subset5', 'subset6', 'subset7']:
                    train_series.add(series_uid)
                elif subset_name == 'subset8':
                    val_series.add(series_uid)
                elif subset_name == 'subset9':
                    test_series.add(series_uid)
            elif num_subsets == 5:
                # Partial LUNA16 (5 subsets)
                if subset_name in ['subset0', 'subset1', 'subset2']:
                    train_series.add(series_uid)
                elif subset_name == 'subset3':
                    val_series.add(series_uid)
                elif subset_name == 'subset4':
                    test_series.add(series_uid)
            else:
                # Generic: 60/20/20 by index
                subset_idx = int(subset_name.replace('subset', ''))
                train_cutoff = int(num_subsets * 0.6)
                val_cutoff = int(num_subsets * 0.8)
                if subset_idx < train_cutoff:
                    train_series.add(series_uid)
                elif subset_idx < val_cutoff:
                    val_series.add(series_uid)
                else:
                    test_series.add(series_uid)
    
    train_df = metadata_df[metadata_df['series_uid'].isin(train_series)]
    val_df = metadata_df[metadata_df['series_uid'].isin(val_series)]
    test_df = metadata_df[metadata_df['series_uid'].isin(test_series)]
    
    return train_df, val_df, test_df


def load_checkpoint(checkpoint_path):
    """Load checkpoint file to resume processing."""
    if checkpoint_path.exists():
        with open(checkpoint_path, 'r') as f:
            return json.load(f)
    return {'processed_scans': [], 'all_results': []}


def save_checkpoint(checkpoint_path, processed_scans, all_results):
    """Save progress checkpoint."""
    checkpoint = {
        'processed_scans': list(processed_scans),
        'all_results': all_results
    }
    with open(checkpoint_path, 'w') as f:
        json.dump(checkpoint, f, indent=2)


# Global flag for graceful shutdown
shutdown_requested = False

def signal_handler(signum, frame):
    """Handle interrupt signals gracefully."""
    global shutdown_requested
    print(f"\n\n  {BOLD}{RED}⚠️  Shutdown requested. Saving checkpoint and exiting...{RESET}")
    shutdown_requested = True


def check_memory():
    """Return True if at least 10 GB of RAM is available."""
    mem = psutil.virtual_memory()
    available_gb = mem.available / (1024 ** 3)
    return available_gb >= 10


def process_scan_wrapper(args):
    """Top-level wrapper so ProcessPoolExecutor can pickle the call."""
    mhd_path, annotations_df, candidates_df, output_dir = args
    try:
        results, stats = process_scan(mhd_path, annotations_df, candidates_df, output_dir)
        return (mhd_path.stem, results, stats, None)
    except Exception as e:
        return (mhd_path.stem, [], {
            'total': 0,
            'success': 0,
            'failures': {
                'nodule_out_of_bounds': 0,
                'context_out_of_bounds': 0,
                'load_error': 1
            }
        }, str(e))


def main():
    """Main preprocessing pipeline with checkpoint/resume support."""
    global shutdown_requested
    
    # Register signal handlers for graceful shutdown
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    banner()
    log_file = setup_logging()
    print(f"  {DIM}Press Ctrl+C to safely stop and save progress{RESET}\n")
    
    # Create output directories
    (OUTPUT_DIR / "nodule_patches").mkdir(parents=True, exist_ok=True)
    (OUTPUT_DIR / "context_patches").mkdir(parents=True, exist_ok=True)
    (OUTPUT_DIR / "metadata").mkdir(parents=True, exist_ok=True)
    
    checkpoint_path = OUTPUT_DIR / "checkpoint.json"
    
    # Load checkpoint if exists
    checkpoint = load_checkpoint(checkpoint_path)
    processed_scans = set(checkpoint['processed_scans'])
    all_results = checkpoint['all_results']
    
    if processed_scans:
        info("Status", f"Resumed from {len(processed_scans)} scans processed")
        info("Status", f"{len(all_results)} patches already extracted\n")
    
    # Load annotations and candidates
    section("DATA PREPARATION")
    info("Status", "Loading annotations and candidates...")
    annotations_df = pd.read_csv(DATA_DIR / "annotations.csv")
    candidates_df = pd.read_csv(DATA_DIR / "candidates.csv")
    
    info("Annotations", f"{len(annotations_df)} nodules")
    info("Candidates", f"{len(candidates_df)} total")
    
    # Get scan paths
    scan_paths = get_scan_paths()
    total_scans = sum(len(paths) for paths in scan_paths.values())
    remaining = total_scans - len(processed_scans)
    info("Total scans", f"{total_scans} | Remaining: {remaining}")
    
    # Process each scan
    section("PATCH EXTRACTION")
    scans_processed_this_session = 0
    
    # Aggregate extraction statistics
    aggregate_stats = {
        'total_candidates': 0,
        'total_success': 0,
        'total_nodule_oob': 0,
        'total_context_oob': 0,
        'total_load_error': 0
    }
    
    # Check memory and auto-scale workers
    if not check_memory():
        warn(f"Low memory detected — reducing workers from {NUM_WORKERS} to {max(8, NUM_WORKERS // 2)}")
        n_workers = max(8, NUM_WORKERS // 2)
    else:
        n_workers = NUM_WORKERS
    info("Parallel workers", n_workers)

    for subset_name, paths in scan_paths.items():
        if shutdown_requested:
            break

        print(f"\nProcessing {subset_name}...")

        # Filter out already processed scans
        remaining_paths = [p for p in paths if p.stem not in processed_scans]
        skipped = len(paths) - len(remaining_paths)
        if skipped > 0:
            print(f"  (Skipping {skipped} already processed scans)")

        if len(remaining_paths) == 0:
            print(f"  All scans in {subset_name} already processed")
            continue

        # Build task args for parallel execution
        task_args = [(p, annotations_df, candidates_df, OUTPUT_DIR) for p in remaining_paths]
        print(f"  Processing {len(remaining_paths)} scans with {n_workers} CPU cores...")

        batch_results = []
        batch_stats = []
        completed = 0

        with ProcessPoolExecutor(max_workers=n_workers) as executor:
            futures = {executor.submit(process_scan_wrapper, a): a[0] for a in task_args}

            with tqdm(total=len(futures), desc=subset_name,
                      bar_format='{l_bar}%s{bar}%s{r_bar}' % (GREEN, RESET)) as pbar:

                for future in as_completed(futures):
                    if shutdown_requested:
                        executor.shutdown(wait=False, cancel_futures=True)
                        break

                    try:
                        series_uid, results, scan_stats, error = future.result()
                        if error:
                            warn(f"Error processing {series_uid}: {error}")

                        batch_results.extend(results)
                        batch_stats.append(scan_stats)
                        processed_scans.add(series_uid)
                        scans_processed_this_session += 1
                        completed += 1

                        # Checkpoint every ~1,000 patches in the batch
                        if len(batch_results) >= 1000:
                            all_results.extend(batch_results)
                            for s in batch_stats:
                                aggregate_stats['total_candidates'] += s['total']
                                aggregate_stats['total_success'] += s['success']
                                aggregate_stats['total_nodule_oob'] += s['failures']['nodule_out_of_bounds']
                                aggregate_stats['total_context_oob'] += s['failures']['context_out_of_bounds']
                                aggregate_stats['total_load_error'] += s['failures']['load_error']
                            save_checkpoint(checkpoint_path, processed_scans, all_results)
                            batch_results = []
                            batch_stats = []
                            gc.collect()

                    except Exception as e:
                        warn(f"Error getting result: {e}")
                    finally:
                        pbar.update(1)

        # Flush remaining batch results
        if batch_results:
            all_results.extend(batch_results)
            for s in batch_stats:
                aggregate_stats['total_candidates'] += s['total']
                aggregate_stats['total_success'] += s['success']
                aggregate_stats['total_nodule_oob'] += s['failures']['nodule_out_of_bounds']
                aggregate_stats['total_context_oob'] += s['failures']['context_out_of_bounds']
                aggregate_stats['total_load_error'] += s['failures']['load_error']

        gc.collect()
        print(f"  Completed {completed}/{len(remaining_paths)} scans in {subset_name}")
    
    # Final checkpoint save
    save_checkpoint(checkpoint_path, processed_scans, all_results)
    
    # Print extraction statistics
    section("EXTRACTION STATISTICS")
    info("Total candidates", f"{aggregate_stats['total_candidates']:,}")
    info("Successfully extracted", f"{aggregate_stats['total_success']:,}")
    info("Nodule out-of-bounds", f"{aggregate_stats['total_nodule_oob']:,}")
    info("Context out-of-bounds", f"{aggregate_stats['total_context_oob']:,}")
    info("Load errors", f"{aggregate_stats['total_load_error']:,}")
    if aggregate_stats['total_candidates'] > 0:
        total_rejected = (aggregate_stats['total_candidates'] - aggregate_stats['total_success'])
        rejection_rate = total_rejected / aggregate_stats['total_candidates'] * 100
        info("Rejection rate", f"{rejection_rate:.1f}%")
    
    if shutdown_requested:
        section("PREPROCESSING PAUSED")
        info("Progress", f"{len(processed_scans)}/{total_scans} scans processed")
        info("Patches extracted", f"{len(all_results)}")
        print(f"\n  {DIM}Run the script again to resume from checkpoint{RESET}")
        return
    
    # All scans processed - create final metadata
    section("METADATA GENERATION")
    
    metadata_df = pd.DataFrame(all_results)
    info("Total patches", f"{len(metadata_df)}")
    info("Positives", f"{(metadata_df['label'] == 1).sum()}")
    info("Negatives", f"{(metadata_df['label'] == 0).sum()}")
    info("Hard negatives", f"{metadata_df['is_hard_negative'].sum()}")
    
    # Create splits
    section("DATA SPLITS")
    info("Status", "Creating train/val/test splits...")
    train_df, val_df, test_df = create_splits(metadata_df, scan_paths)
    
    # FIX: Only balance TRAINING set — val/test must maintain natural distribution
    print(f"\n{BOLD}{BLUE}DATA BALANCING{RESET}")
    info("Strategy", "Balancing TRAINING ONLY (1:7 ratio)")
    info("Val/Test", "Keeping NATURAL distribution (~1:1350)")
    
    train_balanced = balance_samples(train_df, pos_to_neg_ratio=7)
    val_final = val_df    # NO balancing - natural distribution
    test_final = test_df  # NO balancing - natural distribution
    
    def _print_split(name, df, is_balanced=False):
        n_pos = (df['label'] == 1).sum()
        n_neg = (df['label'] == 0).sum()
        ratio = n_neg / max(n_pos, 1)
        status = "(balanced 1:7)" if is_balanced else "(natural ~1:1350)"
        
        print(f"\n  {BOLD}{name} {status}:{RESET} {len(df)} samples")
        print(f"    {GREEN}Positives:{RESET} {n_pos:,}")
        print(f"    {RED}Negatives:{RESET} {n_neg:,}")
        print(f"    Ratio: 1:{ratio:.0f}")
    
    print(f"\n  {BOLD}{BLUE}FINAL SPLIT STATISTICS{RESET}")
    _print_split("Train", train_balanced, is_balanced=True)
    _print_split("Validation", val_final, is_balanced=False)
    _print_split("Test", test_final, is_balanced=False)
    
    # Save metadata
    section("SAVING")
    info("Status", "Saving metadata CSVs...")
    train_balanced.to_csv(OUTPUT_DIR / "metadata" / "train_samples.csv", index=False)
    val_final.to_csv(OUTPUT_DIR / "metadata" / "val_samples.csv", index=False)
    test_final.to_csv(OUTPUT_DIR / "metadata" / "test_samples.csv", index=False)
    metadata_df.to_csv(OUTPUT_DIR / "metadata" / "all_samples.csv", index=False)
    
    # Save statistics
    stats = {
        'total_patches': len(metadata_df),
        'total_positives': int((metadata_df['label'] == 1).sum()),
        'total_negatives': int((metadata_df['label'] == 0).sum()),
        'train_samples': len(train_balanced),
        'val_samples': len(val_final),
        'test_samples': len(test_final),
        'extraction_stats': aggregate_stats,
        'patch_sizes': {
            'nodule': NODULE_PATCH_SIZE,
            'context': CONTEXT_DOWNSAMPLE
        },
        'hu_window': {
            'min': HU_MIN,
            'max': HU_MAX
        }
    }
    
    with open(OUTPUT_DIR / "statistics.json", 'w') as f:
        json.dump(stats, f, indent=2)
    
    # Remove checkpoint file on successful completion
    if checkpoint_path.exists():
        checkpoint_path.unlink()
    
    section("PREPROCESSING COMPLETE")
    print(f"\n  {DIM}Output directory:{RESET} {OUTPUT_DIR}/")
    print(f"    nodule_patches/:  {len(metadata_df)} .npz files")
    print(f"    context_patches/: {len(metadata_df)} .npz files")
    print(f"    metadata/:        train_samples.csv, val_samples.csv, test_samples.csv")
    print(f"    statistics.json")


if __name__ == "__main__":
    main()