src/robustness_analysis.py

"""
Robustness Analysis and Uncertainty Estimation for Brain Tumor Segmentation
"""

import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, Model
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
import json
import os
from sklearn.metrics import confusion_matrix, classification_report
import pandas as pd


class RobustnessAnalyzer:
    """
    Analyze model robustness to various perturbations and corruptions
    """
    
    def __init__(self, model, input_shape=(224, 224, 3)):
        """
        Initialize robustness analyzer
        
        Args:
            model: Trained segmentation model
            input_shape: Shape of input images
        """
        self.model = model
        self.input_shape = input_shape
        self.results = {}
    
    def add_gaussian_noise(self, image, std=0.1):
        """Add Gaussian noise to image"""
        noise = np.random.normal(0, std, image.shape)
        return np.clip(image + noise, 0, 1)
    
    def add_salt_pepper_noise(self, image, salt_prob=0.01, pepper_prob=0.01):
        """Add salt and pepper noise to image"""
        noisy_image = image.copy()
        
        # Salt noise
        salt_mask = np.random.random(image.shape) < salt_prob
        noisy_image[salt_mask] = 1
        
        # Pepper noise
        pepper_mask = np.random.random(image.shape) < pepper_prob
        noisy_image[pepper_mask] = 0
        
        return noisy_image
    
    def add_gaussian_blur(self, image, kernel_size=3):
        """Apply Gaussian blur to image"""
        from scipy.ndimage import gaussian_filter
        return gaussian_filter(image, sigma=kernel_size/3)
    
    def add_motion_blur(self, image, kernel_size=5, angle=0):
        """Apply motion blur to image"""
        from scipy.ndimage import convolve
        
        # Create motion blur kernel
        kernel = np.zeros((kernel_size, kernel_size))
        center = kernel_size // 2
        kernel[center, :] = 1 / kernel_size
        
        # Rotate kernel
        from scipy.ndimage import rotate
        kernel = rotate(kernel, angle, reshape=False)
        
        # Apply convolution
        blurred = convolve(image, kernel, mode='reflect')
        return blurred
    
    def change_brightness(self, image, factor=0.5):
        """Change image brightness"""
        return np.clip(image * factor, 0, 1)
    
    def change_contrast(self, image, factor=0.5):
        """Change image contrast"""
        mean = np.mean(image)
        return np.clip((image - mean) * factor + mean, 0, 1)
    
    def rotate_image(self, image, angle=10):
        """Rotate image by angle degrees"""
        from scipy.ndimage import rotate
        return rotate(image, angle, axes=(0, 1), reshape=False, mode='reflect')
    
    def scale_image(self, image, scale_factor=0.9):
        """Scale image by factor"""
        from scipy.ndimage import zoom
        h, w = image.shape[:2]
        scaled = zoom(image, (scale_factor, scale_factor, 1), order=1)
        
        # Crop or pad to original size
        if scaled.shape[0] > h:
            start = (scaled.shape[0] - h) // 2
            scaled = scaled[start:start+h, :, :]
        elif scaled.shape[0] < h:
            pad_h = (h - scaled.shape[0]) // 2
            scaled = np.pad(scaled, ((pad_h, h - scaled.shape[0] - pad_h), (0, 0), (0, 0)), mode='constant')
        
        if scaled.shape[1] > w:
            start = (scaled.shape[1] - w) // 2
            scaled = scaled[:, start:start+w, :]
        elif scaled.shape[1] < w:
            pad_w = (w - scaled.shape[1]) // 2
            scaled = np.pad(scaled, ((0, 0), (pad_w, w - scaled.shape[1] - pad_w), (0, 0)), mode='constant')
        
        return scaled
    
    def evaluate_robustness(self, X_test, y_test, corruption_type='gaussian_noise', 
                          corruption_levels=None, metric_fn=None):
        """
        Evaluate model robustness to a specific corruption type
        
        Args:
            X_test: Test images
            y_test: Ground truth masks
            corruption_type: Type of corruption to apply
            corruption_levels: List of corruption levels to test
            metric_fn: Function to compute metric (default: Dice coefficient)
            
        Returns:
            Dictionary of robustness metrics
        """
        if corruption_levels is None:
            corruption_levels = [0.01, 0.05, 0.1, 0.2, 0.3, 0.5]
        
        if metric_fn is None:
            def metric_fn(y_true, y_pred):
                intersection = np.sum(y_true * y_pred)
                union = np.sum(y_true) + np.sum(y_pred)
                return (2. * intersection + 1e-6) / (union + 1e-6)
        
        # Get baseline performance
        baseline_preds = self.model.predict(X_test)
        baseline_score = np.mean([metric_fn(y_true, y_pred) 
                                 for y_true, y_pred in zip(y_test, baseline_preds)])
        
        # Get corruption function
        corruption_fn = getattr(self, f'add_{corruption_type}')
        
        # Evaluate at each corruption level
        scores = []
        for level in corruption_levels:
            # Apply corruption
            corrupted_X = np.array([corruption_fn(x, level) for x in X_test])
            
            # Predict on corrupted images
            preds = self.model.predict(corrupted_X)
            
            # Compute metric
            level_scores = [metric_fn(y_true, y_pred) 
                          for y_true, y_pred in zip(y_test, preds)]
            scores.append(np.mean(level_scores))
        
        # Compute robustness metrics
        results = {
            'corruption_type': corruption_type,
            'baseline_score': float(baseline_score),
            'corruption_levels': corruption_levels,
            'scores': [float(s) for s in scores],
            'mean_corruption_score': float(np.mean(scores)),
            'robustness_index': float(np.mean(scores) / baseline_score) if baseline_score > 0 else 0,
            'performance_drop': float(baseline_score - np.mean(scores))
        }
        
        self.results[corruption_type] = results
        
        return results
    
    def evaluate_all_corruptions(self, X_test, y_test, corruption_types=None, **kwargs):
        """
        Evaluate model robustness to all corruption types
        
        Args:
            X_test: Test images
            y_test: Ground truth masks
            corruption_types: List of corruption types to test
            **kwargs: Additional arguments for evaluate_robustness
            
        Returns:
            Dictionary of all robustness results
        """
        if corruption_types is None:
            corruption_types = [
                'gaussian_noise',
                'salt_pepper_noise',
                'gaussian_blur',
                'motion_blur',
                'brightness',
                'contrast',
                'rotation',
                'scaling'
            ]
        
        all_results = {}
        for corruption_type in corruption_types:
            print(f"Evaluating robustness to {corruption_type}...")
            results = self.evaluate_robustness(X_test, y_test, corruption_type, **kwargs)
            all_results[corruption_type] = results
        
        return all_results
    
    def plot_robustness_results(self, results=None, save_path=None):
        """
        Plot robustness analysis results
        
        Args:
            results: Results dictionary (if None, uses self.results)
            save_path: Path to save plot
        """
        if results is None:
            results = self.results
        
        fig, axes = plt.subplots(2, 4, figsize=(20, 10))
        axes = axes.flatten()
        
        for idx, (corruption_type, result) in enumerate(results.items()):
            if idx >= 8:
                break
                
            ax = axes[idx]
            ax.plot(result['corruption_levels'], result['scores'], 'o-', linewidth=2)
            ax.axhline(y=result['baseline_score'], color='r', linestyle='--', alpha=0.5, label='Baseline')
            ax.set_title(f"{corruption_type.replace('_', ' ').title()}\n(Robustness Index: {result['robustness_index']:.3f})")
            ax.set_xlabel('Corruption Level')
            ax.set_ylabel('Dice Score')
            ax.legend()
            ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
    
    def save_results(self, results, save_dir='./robustness_results'):
        """Save robustness results to files"""
        os.makedirs(save_dir, exist_ok=True)
        
        # Save as JSON
        with open(os.path.join(save_dir, 'robustness_results.json'), 'w') as f:
            json.dump(results, f, indent=2)
        
        # Save as CSV
        rows = []
        for corruption_type, result in results.items():
            for level, score in zip(result['corruption_levels'], result['scores']):
                rows.append({
                    'corruption_type': corruption_type,
                    'corruption_level': level,
                    'dice_score': score,
                    'baseline_score': result['baseline_score'],
                    'robustness_index': result['robustness_index']
                })
        
        df = pd.DataFrame(rows)
        df.to_csv(os.path.join(save_dir, 'robustness_results.csv'), index=False)
        
        # Save plots
        self.plot_robustness_results(results, os.path.join(save_dir, 'robustness_plots.png'))


class UncertaintyEstimator:
    """
    Uncertainty estimation for segmentation models using Monte Carlo Dropout
    and Deep Ensembles
    """
    
    def __init__(self, model, num_samples=50, batch_size=32):
        """
        Initialize uncertainty estimator
        
        Args:
            model: Trained segmentation model with dropout
            num_samples: Number of Monte Carlo samples
            batch_size: Batch size for predictions
        """
        self.model = model
        self.num_samples = num_samples
        self.batch_size = batch_size
    
    def enable_dropout(self):
        """Enable dropout at inference time for MC Dropout"""
        for layer in self.model.layers:
            if isinstance(layer, layers.Dropout):
                layer.trainable = True
    
    def mc_dropout_predict(self, X, num_samples=None):
        """
        Monte Carlo Dropout prediction
        
        Args:
            X: Input images
            num_samples: Number of MC samples (if None, uses self.num_samples)
            
        Returns:
            mean_prediction, uncertainty, predictions_array
        """
        if num_samples is None:
            num_samples = self.num_samples
        
        # Enable dropout
        self.enable_dropout()
        
        # Get multiple predictions with dropout enabled
        predictions = []
        for _ in range(num_samples):
            pred = self.model.predict(X, verbose=0)
            predictions.append(pred)
        
        predictions = np.array(predictions)
        
        # Compute mean and uncertainty
        mean_pred = np.mean(predictions, axis=0)
        uncertainty = np.std(predictions, axis=0)
        
        return mean_pred, uncertainty, predictions
    
    def deep_ensemble_predict(self, models, X):
        """
        Deep ensemble prediction
        
        Args:
            models: List of trained models
            X: Input images
            
        Returns:
            mean_prediction, uncertainty, predictions_array
        """
        predictions = []
        for model in models:
            pred = model.predict(X, verbose=0)
            predictions.append(pred)
        
        predictions = np.array(predictions)
        
        # Compute mean and uncertainty
        mean_pred = np.mean(predictions, axis=0)
        uncertainty = np.std(predictions, axis=0)
        
        return mean_pred, uncertainty, predictions
    
    def compute_aleatoric_uncertainty(self, model, X, num_samples=10):
        """
        Estimate aleatoric uncertainty (data uncertainty)
        
        Args:
            model: Model that outputs both prediction and uncertainty
            X: Input images
            num_samples: Number of samples for estimation
            
        Returns:
            Aleatoric uncertainty map
        """
        # For models that output uncertainty directly
        predictions = []
        for _ in range(num_samples):
            pred = model.predict(X, verbose=0)
            if isinstance(pred, list) and len(pred) > 1:
                # Assume second output is uncertainty
                predictions.append(pred[1])
            else:
                predictions.append(pred)
        
        predictions = np.array(predictions)
        aleatoric_uncertainty = np.mean(predictions, axis=0)
        
        return aleatoric_uncertainty
    
    def compute_epistemic_uncertainty(self, models, X):
        """
        Estimate epistemic uncertainty (model uncertainty) using ensemble
        
        Args:
            models: List of trained models
            X: Input images
            
        Returns:
            Epistemic uncertainty map
        """
        predictions = []
        for model in models:
            pred = model.predict(X, verbose=0)
            predictions.append(pred)
        
        predictions = np.array(predictions)
        epistemic_uncertainty = np.var(predictions, axis=0)
        
        return epistemic_uncertainty
    
    def get_confidence_intervals(self, predictions_array, confidence_level=0.95):
        """
        Compute confidence intervals from prediction samples
        
        Args:
            predictions_array: Array of predictions (num_samples, height, width, channels)
            confidence_level: Confidence level for intervals
            
        Returns:
            lower_bound, upper_bound
        """
        alpha = 1 - confidence_level
        lower_percentile = alpha / 2 * 100
        upper_percentile = (1 - alpha / 2) * 100
        
        lower_bound = np.percentile(predictions_array, lower_percentile, axis=0)
        upper_bound = np.percentile(predictions_array, upper_percentile, axis=0)
        
        return lower_bound, upper_bound
    
    def visualize_uncertainty(self, image, mean_pred, uncertainty, threshold=0.5, 
                            save_path=None, cmap='viridis'):
        """
        Visualize prediction and uncertainty
        
        Args:
            image: Input image
            mean_pred: Mean prediction
            uncertainty: Uncertainty map
            threshold: Threshold for binary prediction
            save_path: Path to save visualization
            cmap: Colormap for uncertainty
        """
        fig, axes = plt.subplots(1, 4, figsize=(16, 4))
        
        # Input image
        axes[0].imshow(image)
        axes[0].set_title('Input Image')
        axes[0].axis('off')
        
        # Mean prediction
        binary_pred = (mean_pred >= threshold).astype(int)
        axes[1].imshow(binary_pred, cmap='gray')
        axes[1].set_title(f'Binary Prediction (threshold={threshold})')
        axes[1].axis('off')
        
        # Uncertainty map
        im = axes[2].imshow(uncertainty, cmap=cmap)
        axes[2].set_title('Uncertainty Map')
        axes[2].axis('off')
        plt.colorbar(im, ax=axes[2])
        
        # Overlay uncertainty on prediction
        axes[3].imshow(binary_pred, cmap='gray', alpha=0.7)
        im = axes[3].imshow(uncertainty, cmap=cmap, alpha=0.5)
        axes[3].set_title('Prediction with Uncertainty Overlay')
        axes[3].axis('off')
        plt.colorbar(im, ax=axes[3])
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
    
    def save_uncertainty_results(self, mean_pred, uncertainty, image, save_dir='./uncertainty_results'):
        """Save uncertainty estimation results"""
        os.makedirs(save_dir, exist_ok=True)
        
        # Save predictions
        np.save(os.path.join(save_dir, 'mean_prediction.npy'), mean_pred)
        np.save(os.path.join(save_dir, 'uncertainty.npy'), uncertainty)
        
        # Save visualization
        self.visualize_uncertainty(
            image, mean_pred, uncertainty,
            save_path=os.path.join(save_dir, 'uncertainty_visualization.png')
        )
        
        # Save summary statistics
        stats = {
            'mean_uncertainty': float(np.mean(uncertainty)),
            'max_uncertainty': float(np.max(uncertainty)),
            'min_uncertainty': float(np.min(uncertainty)),
            'std_uncertainty': float(np.std(uncertainty)),
            'high_uncertainty_pixels': float(np.mean(uncertainty > 0.5)),
            'medium_uncertainty_pixels': float(np.mean((uncertainty > 0.2) & (uncertainty <= 0.5))),
            'low_uncertainty_pixels': float(np.mean(uncertainty <= 0.2))
        }
        
        with open(os.path.join(save_dir, 'uncertainty_stats.json'), 'w') as f:
            json.dump(stats, f, indent=2)


class MulticlassSegmentationModel:
    """
    Multiclass segmentation model for different tumor types
    """
    
    def __init__(self, input_shape=(224, 224, 3), num_classes=4, base_filters=64, dropout_rate=0.2):
        """
        Initialize multiclass segmentation model
        
        Args:
            input_shape: Shape of input images
            num_classes: Number of segmentation classes (including background)
            base_filters: Number of base filters
            dropout_rate: Dropout rate
        """
        self.input_shape = input_shape
        self.num_classes = num_classes
        self.base_filters = base_filters
        self.dropout_rate = dropout_rate
        self.model = None
        
        # Class names (can be customized)
        self.class_names = ['background', 'glioma', 'meningioma', 'pituitary']
    
    def build_model(self, use_attention=False):
        """
        Build multiclass U-Net model
        
        Args:
            use_attention: Whether to use attention gates
            
        Returns:
            Compiled model
        """
        inputs = layers.Input(shape=self.input_shape, name='image_input')
        
        # Normalize input
        x = layers.Rescaling(1.0 / 255)(inputs)
        
        # Encoder
        filters = self.base_filters
        skip_connections = []
        
        for i in range(4):
            # Convolutional block
            x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                            kernel_initializer='he_normal')(x)
            x = layers.BatchNormalization()(x)
            x = layers.Dropout(self.dropout_rate)(x)
            x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                            kernel_initializer='he_normal')(x)
            x = layers.BatchNormalization()(x)
            
            skip_connections.append(x)
            x = layers.MaxPooling2D(pool_size=(2, 2))(x)
            filters *= 2
        
        # Bottleneck
        x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                        kernel_initializer='he_normal')(x)
        x = layers.BatchNormalization()(x)
        x = layers.Dropout(self.dropout_rate)(x)
        x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                        kernel_initializer='he_normal')(x)
        x = layers.BatchNormalization()(x)
        
        # Decoder
        filters //= 2
        
        for i in range(4):
            # Upsampling
            x = layers.Conv2DTranspose(filters, (2, 2), strides=(2, 2), padding='same')(x)
            
            # Apply attention gate if enabled
            if use_attention:
                # Attention mechanism
                skip = skip_connections.pop()
                attention_weights = layers.Conv2D(filters, 1, padding='same', use_bias=False)(skip)
                attention_weights = layers.Activation('sigmoid')(attention_weights)
                x = layers.Concatenate()([x, layers.Multiply()([skip, attention_weights])])
            else:
                # Simple concatenation
                x = layers.Concatenate()([x, skip_connections.pop()])
            
            # Convolutional block
            x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                            kernel_initializer='he_normal')(x)
            x = layers.BatchNormalization()(x)
            x = layers.Dropout(self.dropout_rate)(x)
            x = layers.Conv2D(filters, 3, activation='relu', padding='same', 
                            kernel_initializer='he_normal')(x)
            x = layers.BatchNormalization()(x)
            
            filters //= 2
        
        # Output layer with softmax for multiclass
        outputs = layers.Conv2D(self.num_classes, (1, 1), activation='softmax', padding='same')(x)
        
        self.model = Model(inputs=[inputs], outputs=[outputs], name='multiclass_unet')
        
        return self.model
    
    def compile_model(self, learning_rate=1e-4):
        """
        Compile model with appropriate loss and metrics
        
        Args:
            learning_rate: Learning rate for optimizer
        """
        if self.model is None:
            self.build_model()
        
        # Categorical crossentropy for multiclass
        self.model.compile(
            optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
            loss='categorical_crossentropy',
            metrics=[
                'accuracy',
                tf.keras.metrics.MeanIoU(num_classes=self.num_classes),
                self.dice_coefficient_multiclass
            ]
        )
        
        return self.model
    
    def dice_coefficient_multiclass(self, y_true, y_pred, smooth=1e-6):
        """
        Dice coefficient for multiclass segmentation
        
        Args:
            y_true: Ground truth one-hot encoded masks
            y_pred: Predicted probabilities
            smooth: Smoothing factor
            
        Returns:
            Mean Dice coefficient across classes
        """
        # Get class predictions
        y_true_classes = tf.argmax(y_true, axis=-1)
        y_pred_classes = tf.argmax(y_pred, axis=-1)
        
        # One-hot encode predictions
        y_true_onehot = tf.one_hot(y_true_classes, depth=self.num_classes)
        y_pred_onehot = tf.one_hot(y_pred_classes, depth=self.num_classes)
        
        # Compute Dice for each class
        dice_scores = []
        for i in range(self.num_classes):
            intersection = tf.reduce_sum(y_true_onehot[..., i] * y_pred_onehot[..., i])
            union = tf.reduce_sum(y_true_onehot[..., i]) + tf.reduce_sum(y_pred_onehot[..., i])
            dice = (2. * intersection + smooth) / (union + smooth)
            dice_scores.append(dice)
        
        return tf.reduce_mean(dice_scores)
    
    def prepare_multiclass_masks(self, masks, num_classes=None):
        """
        Convert integer masks to one-hot encoded masks
        
        Args:
            masks: Integer masks with values 0 to num_classes-1
            num_classes: Number of classes (if None, uses self.num_classes)
            
        Returns:
            One-hot encoded masks
        """
        if num_classes is None:
            num_classes = self.num_classes
        
        # Convert to one-hot
        one_hot = tf.one_hot(masks.astype(int), depth=num_classes)
        
        return one_hot.numpy()
    
    def train(self, X_train, y_train, X_val, y_val, epochs=100, batch_size=16, callbacks=None):
        """
        Train the multiclass model
        
        Args:
            X_train: Training images
            y_train: Training masks (integer or one-hot)
            X_val: Validation images
            y_val: Validation masks
            epochs: Number of training epochs
            batch_size: Batch size
            callbacks: List of Keras callbacks
            
        Returns:
            Training history
        """
        # Build and compile model
        self.compile_model()
        
        # Convert masks to one-hot if needed
        if len(y_train.shape) == 3 or (len(y_train.shape) == 4 and y_train.shape[-1] != self.num_classes):
            y_train = self.prepare_multiclass_masks(y_train)
            y_val = self.prepare_multiclass_masks(y_val)
        
        # Default callbacks
        if callbacks is None:
            callbacks = []
        
        callbacks.extend([
            tf.keras.callbacks.EarlyStopping(
                monitor='val_loss',
                patience=15,
                restore_best_weights=True
            ),
            tf.keras.callbacks.ReduceLROnPlateau(
                monitor='val_loss',
                factor=0.5,
                patience=5,
                min_lr=1e-7
            )
        ])
        
        # Train
        history = self.model.fit(
            X_train,
            y_train,
            validation_data=(X_val, y_val),
            epochs=epochs,
            batch_size=batch_size,
            callbacks=callbacks,
            verbose=1
        )
        
        return history
    
    def predict(self, X):
        """
        Predict segmentation masks
        
        Args:
            X: Input images
            
        Returns:
            Predicted masks (integer format)
        """
        # Get predictions
        predictions = self.model.predict(X)
        
        # Convert to integer masks
        predicted_masks = np.argmax(predictions, axis=-1)
        
        return predicted_masks
    
    def predict_proba(self, X):
        """
        Predict class probabilities
        
        Args:
            X: Input images
            
        Returns:
            Probability maps for each class
        """
        return self.model.predict(X)
    
    def evaluate_multiclass(self, X_test, y_test, class_names=None):
        """
        Evaluate multiclass segmentation with per-class metrics
        
        Args:
            X_test: Test images
            y_test: Ground truth masks (integer format)
            class_names: List of class names
            
        Returns:
            Dictionary of evaluation metrics
        """
        if class_names is None:
            class_names = self.class_names
        
        # Get predictions
        y_pred = self.predict(X_test)
        
        # Flatten for metric calculation
        y_true_flat = y_test.flatten()
        y_pred_flat = y_pred.flatten()
        
        # Compute per-class metrics
        metrics = {}
        
        # Overall metrics
        overall_accuracy = np.mean(y_true_flat == y_pred_flat)
        
        # Per-class metrics
        class_metrics = {}
        for class_idx, class_name in enumerate(class_names):
            # Binary mask for this class
            true_binary = (y_true_flat == class_idx)
            pred_binary = (y_pred_flat == class_idx)
            
            # Compute metrics
            intersection = np.sum(true_binary & pred_binary)
            union = np.sum(true_binary) + np.sum(pred_binary) - intersection
            
            if np.sum(true_binary) > 0:
                recall = intersection / np.sum(true_binary)
            else:
                recall = 0
            
            if np.sum(pred_binary) > 0:
                precision = intersection / np.sum(pred_binary)
            else:
                precision = 0
            
            if precision + recall > 0:
                f1 = 2 * precision * recall / (precision + recall)
            else:
                f1 = 0
            
            if union > 0:
                iou = intersection / union
            else:
                iou = 0
            
            dice = (2 * intersection + 1e-6) / (np.sum(true_binary) + np.sum(pred_binary) + 1e-6)
            
            class_metrics[class_name] = {
                'precision': float(precision),
                'recall': float(recall),
                'f1_score': float(f1),
                'iou': float(iou),
                'dice': float(dice),
                'support': int(np.sum(true_binary))
            }
        
        # Compute mean metrics
        mean_metrics = {}
        for metric in ['precision', 'recall', 'f1_score', 'iou', 'dice']:
            mean_metrics[f'mean_{metric}'] = float(np.mean([cm[metric] for cm in class_metrics.values()]))
        
        metrics['overall_accuracy'] = float(overall_accuracy)
        metrics['class_metrics'] = class_metrics
        metrics['mean_metrics'] = mean_metrics
        
        # Confusion matrix
        cm = confusion_matrix(y_true_flat, y_pred_flat, labels=list(range(len(class_names))))
        metrics['confusion_matrix'] = cm.tolist()
        
        return metrics
    
    def visualize_multiclass_prediction(self, image, true_mask, pred_mask, class_names=None, save_path=None):
        """
        Visualize multiclass segmentation results
        
        Args:
            image: Input image
            true_mask: Ground truth mask
            pred_mask: Predicted mask
            class_names: List of class names
            save_path: Path to save visualization
        """
        if class_names is None:
            class_names = self.class_names
        
        # Create color map for classes
        colors = plt.cm.Set3(np.linspace(0, 1, len(class_names)))
        
        fig, axes = plt.subplots(1, 4, figsize=(16, 4))
        
        # Input image
        axes[0].imshow(image)
        axes[0].set_title('Input Image')
        axes[0].axis('off')
        
        # Ground truth
        axes[1].imshow(true_mask, cmap='tab10', vmin=0, vmax=len(class_names)-1)
        axes[1].set_title('Ground Truth')
        axes[1].axis('off')
        
        # Prediction
        axes[2].imshow(pred_mask, cmap='tab10', vmin=0, vmax=len(class_names)-1)
        axes[2].set_title('Prediction')
        axes[2].axis('off')
        
        # Overlay
        axes[3].imshow(image)
        axes[3].imshow(pred_mask, cmap='tab10', vmin=0, vmax=len(class_names)-1, alpha=0.5)
        axes[3].set_title('Prediction Overlay')
        axes[3].axis('off')
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()


def create_multiclass_dataset_from_binary(binary_images, binary_masks, tumor_types):
    """
    Create multiclass dataset from binary segmentation data
    
    Args:
        binary_images: List of binary images
        binary_masks: List of binary masks
        tumor_types: List of tumor type labels for each image
        
    Returns:
        Multiclass images and masks
    """
    # Create mapping from tumor type to class index
    tumor_type_to_class = {tumor_type: idx+1 for idx, tumor_type in enumerate(set(tumor_types))}
    
    # Create multiclass masks
    multiclass_masks = []
    for mask, tumor_type in zip(binary_masks, tumor_types):
        # Start with background (class 0)
        multiclass_mask = np.zeros_like(mask, dtype=np.int32)
        
        # Set tumor region to appropriate class
        class_idx = tumor_type_to_class.get(tumor_type, 0)
        multiclass_mask[mask > 0.5] = class_idx
        
        multiclass_masks.append(multiclass_mask)
    
    return np.array(binary_images), np.array(multiclass_masks)