algorithim "# -*- coding: utf-8 -*-
"""Complete Dental Disease Classification with YOLOv11 - Enhanced with Ablation Study"""

# Install required packages
!pip install ultralytics roboflow scikit-learn matplotlib seaborn torch torchvision -q
!pip install efficientnet-pytorch  # For attention components
!nvidia-smi
from IPython.display import clear_output
clear_output()

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from ultralytics import YOLO
from roboflow import Roboflow
from sklearn.metrics import (confusion_matrix, classification_report, 
                           accuracy_score, roc_curve, auc, precision_recall_curve)
from sklearn.preprocessing import label_binarize
from PIL import Image
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
import zipfile
from itertools import cycle
from efficientnet_pytorch import EfficientNet
from collections import defaultdict
##########################################################################################################################
# Optimized YOLOv11 Configuration
config = {
    'results_dir': '/content/results',
    'model_name': 'yolo11n-cls.pt',  # Using YOLOv11 Nano
    'epochs': 20,                   # Increased for better convergence
    'batch_size': 128,               # Increased batch size for stability
    'img_size': 224,                 # Standard for classification
    'lr': 0.001,                     # Lower initial learning rate
    'lr_schedule': 'cosine',         # Learning rate scheduling
    'momentum': 0.9,                 # For SGD optimizer
    'weight_decay': 1e-4,            # L2 regularization
    'augment': True,                 # Enable with optimized augmentations
    'augment_config': {              # Custom augmentation parameters
        'hsv_h': 0.015,             # Slight HSV augmentation
        'hsv_s': 0.7,
        'hsv_v': 0.4,
        'translate': 0.1,           # Small translations
        'scale': 0.2,              # Slight scaling
        'flipud': 0.0,              # No vertical flips
        'fliplr': 0.5,              # Horizontal flips
        'mosaic': 0.0,              # No mosaic for classification
        'mixup': 0.0                # No mixup by default
    },
    'seed': 42,
    'num_samples': 25,
    'font_sizes': {                  # Organized font sizes
        'small': 12,
        'medium': 14,
        'large': 16,
        'title': 18
    },
    'dpi': 300,
    'device': 'cuda' if torch.cuda.is_available() else 'cpu',
    'num_workers': min(8, os.cpu_count()),  # Optimal workers based on system
    'early_stopping': {              # Early stopping configuration
        'patience': 15,
        'delta': 0.001
    },
    'class_colors': ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', 
                    '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf'],
    'optimizer': 'AdamW',           # Better optimizer choice
    'label_smoothing': 0.1,         # Helps with regularization
    'cutmix_prob': 0.0,             # Disabled by default
    'ema': True,                    # Exponential Moving Average
    'ema_decay': 0.9999,
    'grad_clip': 1.0,               # Gradient clipping
    'warmup_epochs': 3,             # Learning rate warmup
    'save_best': True,              # Save best model only
    'eval_interval': 1              # Evaluate every epoch
}
##########################################################################################################################
# Create directories
os.makedirs(config['results_dir'], exist_ok=True)
for subdir in ['roc_curves', 'sample_images', 'confusion_matrices', 
               'classification_reports', 'ablation_study', 'pr_curves']:
    os.makedirs(os.path.join(config['results_dir'], subdir), exist_ok=True)
##########################################################################################################################

# ==============================================
# Novel Technique: Attention-Guided Multi-Scale Feature Fusion (AGMS-FF)
# ==============================================

class ChannelAttention(nn.Module):
    """Channel attention module with learnable reduction ratio"""
    def __init__(self, in_channels, reduction_ratio=None):
        super(ChannelAttention, self).__init__()
        if reduction_ratio is None:
            reduction_ratio = max(1, in_channels // 16)  # Dynamic reduction
        
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction_ratio),
            nn.ReLU(inplace=True),
            nn.Linear(in_channels // reduction_ratio, in_channels),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        b, c, _, _ = x.size()
        avg_out = self.fc(self.avg_pool(x).view(b, c))
        max_out = self.fc(self.max_pool(x).view(b, c))
        out = avg_out + max_out
        return out.view(b, c, 1, 1)

class SpatialAttention(nn.Module):
    """Spatial attention module with configurable kernel size"""
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size % 2 == 1, "Kernel size must be odd"
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv(x)
        return self.sigmoid(x)

class AGMS_FF_Enhancer(nn.Module):
    """Enhanced version with dynamic scaling and residual learning"""
    def __init__(self, in_channels, reduction_ratio=None):
        super(AGMS_FF_Enhancer, self).__init__()
        
        # Multi-scale feature extractors with dynamic channels
        mid_channels = max(32, in_channels // 2)  # Ensure minimum channels
        
        self.conv1x1 = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, 1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True))
        
        self.conv3x3 = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, 3, padding=1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True))
        
        self.conv5x5 = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, 5, padding=2),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True))
        
        # Attention modules
        self.channel_att = ChannelAttention(mid_channels*3, reduction_ratio)
        self.spatial_att = SpatialAttention()
        
        # Feature fusion with residual connection
        self.fusion = nn.Sequential(
            nn.Conv2d(mid_channels*3, in_channels, 1),
            nn.BatchNorm2d(in_channels))
        
        self.relu = nn.ReLU(inplace=True)
        self._initialize_weights()
    
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        identity = x
        
        # Multi-scale feature extraction
        f1 = self.conv1x1(x)
        f3 = self.conv3x3(x)
        f5 = self.conv5x5(x)
        
        # Concatenate features
        fused = torch.cat([f1, f3, f5], dim=1)
        
        # Apply channel attention
        channel_weights = self.channel_att(fused)
        attended = fused * channel_weights
        
        # Apply spatial attention
        spatial_weights = self.spatial_att(attended)
        attended = attended * spatial_weights
        
        # Final fusion with residual
        out = self.fusion(attended)
        out += identity
        return self.relu(out)

class EnhancedYOLOClassifier:
    """Wrapper class that enhances YOLO with our AGMS-FF technique"""
    def __init__(self, base_model_path, num_blocks=1, device='cuda'):
        self.base_model = YOLO(base_model_path).to(device)
        self.device = device
        
        # Get the classification head
        cls_head = self.base_model.model.model[-1]
        
        # For YOLO models, we need to get the feature dimensions differently
        with torch.no_grad():
            test_input = torch.randn(1, 3, config['img_size'], config['img_size']).to(device)
            features = self.base_model.model.model[:-1](test_input)
            in_features = features.shape[1]  # Get channel dimension
            out_features = cls_head.linear.out_features if hasattr(cls_head, 'linear') else len(cls_head.prob)
        
        # Replace the classification head with our enhanced version
        enhanced_head = []
        for _ in range(num_blocks):
            enhanced_head.append(AGMS_FF_Enhancer(in_features))
        
        enhanced_head.extend([
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(in_features, out_features)
        ])
        
        self.enhanced_head = nn.Sequential(*enhanced_head).to(device)
        self.base_model.model.model[-1] = self.enhanced_head
    
    def train(self, data, epochs, imgsz, batch, lr0, project, name, exist_ok=True, **kwargs):
        """Train the enhanced model with YOLO-compatible arguments"""
        # Ensure augment is False for classification
        kwargs['augment'] = False
        return self.base_model.train(
            data=data,  # Note: using 'data' instead of 'data_path'
            epochs=epochs,
            imgsz=imgsz,
            batch=batch,
            lr0=lr0,
            project=project,
            name=name,
            exist_ok=exist_ok,
            **kwargs
        )
    
    def predict(self, source, verbose=False, **kwargs):
        """Make predictions with the enhanced model"""
        kwargs['augment'] = False
        return self.base_model.predict(source, verbose=verbose, **kwargs)

##########################################################################################################################
def run_ablation_study(data_path, test_path, class_names):
    """Comprehensive ablation study with visualization of all components"""
    results = {}
    variants = [
        ('baseline', 'Baseline YOLOv11', None),
        ('multi_scale', 'Multi-Scale Only', {'attention': False}),
        ('attention', 'Attention Only', {'multi_scale': False}),
        ('agms_ff', 'Full AGMS-FF', {})
    ]
    
    for variant, variant_name, params in variants:
        print(f"\nRunning {variant_name}...")
        model = create_model_variant(config['model_name'], variant, params)
        
        # Train model with YOLO-compatible arguments
        model.train(
            data=data_path,  # Using 'data' instead of 'data_path'
            epochs=config['epochs'],
            imgsz=config['img_size'],
            batch=config['batch_size'],
            lr0=config['lr'],
            project=config['results_dir'],
            name=f'dental_{variant}',
            exist_ok=True
        )
        
        # Evaluate model
        y_true, y_pred, y_score, img_paths = evaluate_model(model, test_path, class_names)
        
        if not y_true:
            print(f"Skipping {variant_name} - no valid predictions")
            continue
        
        # Calculate all metrics
        metrics = calculate_all_metrics(y_true, y_pred, y_score, class_names, variant_name)
        results[variant] = metrics
        
        # Generate all visualizations
        plot_confusion_matrix(y_true, y_pred, class_names, variant_name)
        plot_confusion_matrix(y_true, y_pred, class_names, variant_name, normalize=True)
        if len(class_names) > 1:
            plot_roc_curves(y_true, y_score, class_names, variant_name)
            plot_pr_curves(y_true, y_score, class_names, variant_name)
        plot_sample_predictions(y_true, y_pred, img_paths, class_names, 
                              model_name=variant_name)
    
    # Save and compare all results
    save_ablation_results(results, class_names)
    return results
##########################################################################################################################

# ==============================================
# Visualization and Evaluation Functions
# ==============================================

def plot_confusion_matrix(y_true, y_pred, class_names, model_name="Model", normalize=False):
    """Generate and save enhanced confusion matrix with multiple display options"""
    try:
        cm = confusion_matrix(y_true, y_pred, labels=class_names)
        
        plt.figure(figsize=(12, 10))
        sns.set(font_scale=1.2)
        
        if normalize:
            cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
            fmt = '.2f'
            title = f'Normalized Confusion Matrix - {model_name}'
        else:
            fmt = 'd'
            title = f'Confusion Matrix - {model_name}'
        
        # Access font sizes through 'font_sizes' dictionary
        sns.heatmap(cm, annot=True, fmt=fmt, cmap='Blues',
                   xticklabels=class_names, yticklabels=class_names,
                   cbar_kws={'label': 'Count' if not normalize else 'Proportion'},
                   annot_kws={"size": config['font_sizes']['medium']}) 
        
        plt.title(title, fontsize=config['font_sizes']['large']) # Access correct key
        plt.ylabel('True Label', fontsize=config['font_sizes']['medium']) # Access correct key
        plt.xlabel('Predicted Label', fontsize=config['font_sizes']['medium']) # Access correct key
        plt.xticks(rotation=45, ha='right', fontsize=config['font_sizes']['small']) # Access correct key
        plt.yticks(rotation=0, fontsize=config['font_sizes']['small']) # Access correct key
        
        # Save both normalized and regular versions
        suffix = '_normalized' if normalize else ''
        cm_path = os.path.join(config['results_dir'], 'confusion_matrices', 
                             f'cm_{model_name.lower().replace(" ", "_")}{suffix}.png')
        plt.savefig(cm_path, dpi=config['dpi'], bbox_inches='tight')
        plt.close()
        
        # Save confusion matrix as CSV
        cm_df = pd.DataFrame(cm, index=class_names, columns=class_names)
        cm_csv_path = os.path.join(config['results_dir'], 'confusion_matrices',
                                 f'cm_{model_name.lower().replace(" ", "_")}{suffix}.csv')
        cm_df.to_csv(cm_csv_path)
        
        return cm
    except Exception as e:
        print(f"Error generating confusion matrix: {str(e)}")
        return None

##########################################################################################################################
def plot_roc_curves(y_true, y_scores, class_names, model_name="Model"):
    """Plot enhanced ROC curves with per-class and macro-average"""
    try:
        y_true_bin = label_binarize([class_names.index(cls) for cls in y_true], 
                                   classes=range(len(class_names)))
        y_score = np.array(y_scores)
        
        # Compute ROC curve and ROC area for each class
        fpr, tpr, roc_auc = {}, {}, {}
        for i in range(len(class_names)):
            fpr[i], tpr[i], _ = roc_curve(y_true_bin[:, i], y_score[:, i])
            roc_auc[i] = auc(fpr[i], tpr[i])
        
        # Compute micro-average ROC curve and area
        fpr["micro"], tpr["micro"], _ = roc_curve(y_true_bin.ravel(), y_score.ravel())
        roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])
        
        # Compute macro-average ROC curve and area
        all_fpr = np.unique(np.concatenate([fpr[i] for i in range(len(class_names))]))
        mean_tpr = np.zeros_like(all_fpr)
        for i in range(len(class_names)):
            mean_tpr += np.interp(all_fpr, fpr[i], tpr[i])
        mean_tpr /= len(class_names)
        fpr["macro"] = all_fpr
        tpr["macro"] = mean_tpr
        roc_auc["macro"] = auc(fpr["macro"], tpr["macro"])
        
        # Plot all ROC curves
        plt.figure(figsize=(12, 10))
        plt.plot(fpr["micro"], tpr["micro"],
                label=f'micro-average (AUC = {roc_auc["micro"]:.2f})',
                color='deeppink', linestyle=':', linewidth=4)
        
        plt.plot(fpr["macro"], tpr["macro"],
                label=f'macro-average (AUC = {roc_auc["macro"]:.2f})',
                color='navy', linestyle=':', linewidth=4)
        
        colors = cycle(config['class_colors'][:len(class_names)])
        for i, color in zip(range(len(class_names)), colors):
            plt.plot(fpr[i], tpr[i], color=color, lw=2,
                    label=f'{class_names[i]} (AUC = {roc_auc[i]:.2f})')
        
        plt.plot([0, 1], [0, 1], 'k--', lw=2)
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate', fontsize=config['font_sizes']['medium'])
        plt.ylabel('True Positive Rate', fontsize=config['font_sizes']['medium'])
        plt.title(f'ROC Curves - {model_name}', fontsize=config['font_sizes']['large'])
        plt.legend(loc="lower right", fontsize=config['font_sizes']['small'])
        
        roc_path = os.path.join(config['results_dir'], 'roc_curves', 
                              f'roc_{model_name.lower().replace(" ", "_")}.png')
        plt.savefig(roc_path, dpi=config['dpi'], bbox_inches='tight')
        plt.close()
        
        # Save ROC data
        roc_data = {
            'fpr': fpr,
            'tpr': tpr,
            'roc_auc': roc_auc,
            'class_names': class_names
        }
        roc_pkl_path = os.path.join(config['results_dir'], 'roc_curves',
                                  f'roc_{model_name.lower().replace(" ", "_")}.pkl')
        pd.to_pickle(roc_data, roc_pkl_path)
        
        return roc_auc
    except Exception as e:
        print(f"Error generating ROC curves: {str(e)}")
        return None
##########################################################################################################################
def plot_pr_curves(y_true, y_scores, class_names, model_name="Model"):
    """Plot Precision-Recall curves for each class"""
    try:
        y_true_bin = label_binarize([class_names.index(cls) for cls in y_true], 
                                   classes=range(len(class_names)))
        y_score = np.array(y_scores)
        
        precision, recall, pr_auc = {}, {}, {}
        
        # Compute PR curve and area for each class
        for i in range(len(class_names)):
            precision[i], recall[i], _ = precision_recall_curve(y_true_bin[:, i], y_score[:, i])
            pr_auc[i] = auc(recall[i], precision[i])
        
        # Plot all PR curves
        plt.figure(figsize=(12, 10))
        colors = cycle(config['class_colors'][:len(class_names)])
        
        for i, color in zip(range(len(class_names)), colors):
            plt.plot(recall[i], precision[i], color=color, lw=2,
                    label=f'{class_names[i]} (AUC = {pr_auc[i]:.2f})')
        
        plt.xlabel('Recall', fontsize=config['font_sizes']['medium'])
        plt.ylabel('Precision', fontsize=config['font_sizes']['medium'])
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.title(f'Precision-Recall Curves - {model_name}', fontsize=config['font_sizes']['large'])
        plt.legend(loc="lower left", fontsize=config['font_sizes']['small'])
        
        pr_path = os.path.join(config['results_dir'], 'pr_curves', 
                             f'pr_{model_name.lower().replace(" ", "_")}.png')
        plt.savefig(pr_path, dpi=config['dpi'], bbox_inches='tight')
        plt.close()
        
        return pr_auc
    except Exception as e:
        print(f"Error generating PR curves: {str(e)}")
        return None
##########################################################################################################################
def plot_sample_predictions(y_true, y_pred, image_paths, class_names, n_samples=25, model_name="Model"):
    """Sample predictions visualization focusing on 10% incorrect predictions"""
    try:
        # Organize samples by class and prediction status
        samples_dict = defaultdict(list)
        for i, (true, pred) in enumerate(zip(y_true, y_pred)):
            samples_dict[(true, pred == true)].append(i)
        
        # Calculate total incorrect samples available
        incorrect_indices = []
        for class_name in class_names:
            incorrect = [i for i in samples_dict.get((class_name, False), []) 
                        if y_pred[i] in class_names]  # Only show valid classes
            incorrect_indices.extend(incorrect)
        
        # Calculate 10% of incorrect samples
        n_incorrect = max(1, int(0.1 * len(incorrect_indices)))
        n_correct = n_samples - n_incorrect
        
        # Select samples
        selected = []
        
        # First select correct samples (distributed across classes)
        correct_samples_per_class = max(1, n_correct // len(class_names))
        for class_name in class_names:
            correct = samples_dict.get((class_name, True), [])
            selected.extend(random.sample(correct, min(correct_samples_per_class, len(correct))))
        
        # Then add incorrect samples (10% of available incorrect)
        selected.extend(random.sample(incorrect_indices, min(n_incorrect, len(incorrect_indices))))
        
        # If we still need more samples (due to rounding), add random correct ones
        if len(selected) < n_samples:
            remaining = n_samples - len(selected)
            all_correct = [i for i in range(len(y_true)) if y_true[i] == y_pred[i]]
            selected.extend(random.sample(all_correct, min(remaining, len(all_correct))))
        
        selected = selected[:n_samples]
        
        # Create plot with improved layout
        n_cols = 5
        n_rows = int(np.ceil(n_samples / n_cols))
        fig, axes = plt.subplots(n_rows, n_cols, figsize=(20, 4*n_rows))
        if n_rows == 1:
            axes = axes.reshape(1, -1)  # Ensure 2D array for consistency
        
        for i, idx in enumerate(selected):
            row, col = i // n_cols, i % n_cols
            try:
                img = Image.open(image_paths[idx]).convert('RGB').resize((200, 200))
                axes[row, col].imshow(img)
                
                # Enhanced title formatting
                title = f"True: {y_true[idx]}\nPred: {y_pred[idx]}"
                color = 'green' if y_true[idx] == y_pred[idx] else 'red'
                axes[row, col].set_title(title, color=color, fontsize=16, pad=10)
                axes[row, col].axis('off')
                
                # Add border to highlight incorrect predictions
                if y_true[idx] != y_pred[idx]:
                    for spine in axes[row, col].spines.values():
                        spine.set_edgecolor('red')
                        spine.set_linewidth(3)
            except Exception as e:
                print(f"Error displaying image {image_paths[idx]}: {str(e)}")
                axes[row, col].axis('off')
        
        # Turn off empty axes
        for j in range(i+1, n_rows*n_cols) if 'i' in locals() else range(n_rows*n_cols):
            row, col = j // n_cols, j % n_cols
            axes[row, col].axis('off')
        
        plt.suptitle(f'Sample Predictions - {model_name}\n(Green=Correct, Red=Incorrect)', 
                    fontsize=18, y=1.02)
        plt.tight_layout()
        
        # Save plot
        sample_path = os.path.join(config['results_dir'], 'sample_images', 
                                 f'samples_{model_name.lower().replace(" ", "_")}.png')
        plt.savefig(sample_path, dpi=config['dpi'], bbox_inches='tight')
        plt.close()
        
        # Save sample info with probabilities
        sample_info = []
        for idx in selected:
            sample_info.append({
                'image': os.path.basename(image_paths[idx]),
                'true_label': y_true[idx],
                'pred_label': y_pred[idx],
                'correct': y_true[idx] == y_pred[idx]
            })
        
        sample_df = pd.DataFrame(sample_info)
        sample_csv_path = os.path.join(config['results_dir'], 'sample_images',
                                     f'samples_{model_name.lower().replace(" ", "_")}.csv')
        sample_df.to_csv(sample_csv_path, index=False)
        
        return selected
    except Exception as e:
        print(f"Error generating sample predictions: {str(e)}")
        return None

# ==============================================
# Ablation Study Functions
# ==============================================

def create_model_variant(base_model_path, variant, params=None):
    """Create different model variants for ablation study"""
    if variant == 'baseline':
        return YOLO(base_model_path).to(config['device'])
    
    enhanced_model = EnhancedYOLOClassifier(base_model_path, device=config['device'])
    
    if variant == 'multi_scale':
        # Disable attention modules
        for module in enhanced_model.enhanced_head.modules():
            if isinstance(module, AGMS_FF_Enhancer):
                module.channel_att = nn.Identity()
                module.spatial_att = nn.Identity()
    elif variant == 'attention':
        # Disable multi-scale convolutions
        for module in enhanced_model.enhanced_head.modules():
            if isinstance(module, AGMS_FF_Enhancer):
                module.conv3x3 = nn.Identity()
                module.conv5x5 = nn.Identity()
    
    return enhanced_model


def save_ablation_results(results, class_names):
    """Save and visualize comprehensive ablation study results"""
    # Create comparison table
    comparison_data = []
    for variant, metrics in results.items():
        row = {
            'Model': metrics['model'],
            'Accuracy': metrics['accuracy'],
            'Macro Precision': metrics['classification_report']['macro avg']['precision'],
            'Macro Recall': metrics['classification_report']['macro avg']['recall'],
            'Macro F1': metrics['classification_report']['macro avg']['f1-score']
        }
        
        if 'macro_avg_auc' in metrics:
            row['Macro AUC'] = metrics['macro_avg_auc']
            row['Macro PR AUC'] = metrics.get('macro_avg_pr', np.nan)
        
        comparison_data.append(row)
    
    df = pd.DataFrame(comparison_data)
    
    # Save table
    table_path = os.path.join(config['results_dir'], 'ablation_study', 'comparison_table.csv')
    df.to_csv(table_path, index=False)
    
    # Plot metric comparison
    plt.figure(figsize=(15, 10))
    metrics_to_plot = ['Accuracy', 'Macro Precision', 'Macro Recall', 'Macro F1']
    if 'Macro AUC' in df.columns:
        metrics_to_plot.extend(['Macro AUC', 'Macro PR AUC'])
    
    for i, metric in enumerate(metrics_to_plot):
        plt.subplot(2, 3, i+1)
        ax = sns.barplot(
            x=metric, 
            y='Model', 
            data=df.sort_values(metric, ascending=False),
            hue='Model' if i == 0 else None,  # Add hue to first plot only
            palette=config['class_colors'][:len(df)],
            legend=False
        )
        plt.title(metric, fontsize=config['font_sizes']['medium'])
        plt.xlabel('')
        
        # Add value annotations
        for p in ax.patches:
            width = p.get_width()
            plt.text(width * 1.01, p.get_y() + p.get_height()/2.,
                   f'{width:.3f}',
                   ha='left', va='center', fontsize=12)
    
    plt.suptitle('Ablation Study Results Comparison', fontsize=config['font_sizes']['large'], y=1.02)
    plt.tight_layout()
    
    plot_path = os.path.join(config['results_dir'], 'ablation_study', 'metrics_comparison.png')
    plt.savefig(plot_path, dpi=config['dpi'], bbox_inches='tight')
    plt.close()
    
    # Save detailed reports and visualizations for each model
    for variant, metrics in results.items():
        variant_dir = os.path.join(config['results_dir'], 'ablation_study', variant)
        os.makedirs(variant_dir, exist_ok=True)
        
        # Save detailed metrics
        report_path = os.path.join(variant_dir, 'detailed_report.txt')
        with open(report_path, 'w') as f:
            f.write(f"=== {metrics['model']} Detailed Report ===\n\n")
            f.write(f"Accuracy: {metrics['accuracy']:.4f}\n")
            
            if 'macro_avg_auc' in metrics:
                f.write(f"Macro Average ROC AUC: {metrics['macro_avg_auc']:.4f}\n")
                f.write(f"Macro Average PR AUC: {metrics['macro_avg_pr']:.4f}\n")
            
            f.write("\nClassification Report:\n")
            # Use the saved classification report from metrics
            f.write(classification_report(
                metrics['y_true'], 
                metrics['y_pred'], 
                target_names=class_names, 
                digits=4
            ))
            
            if 'roc_auc' in metrics:
                f.write("\nClass-wise ROC AUC Scores:\n")
                for cls, auc_score in metrics['roc_auc'].items():
                    f.write(f"{cls}: {auc_score:.4f}\n")
                
                f.write("\nClass-wise PR AUC Scores:\n")
                for cls, pr_score in metrics['pr_auc'].items():
                    f.write(f"{cls}: {pr_score:.4f}\n")
    
    return df
###########################################################################################
def calculate_all_metrics(y_true, y_pred, y_score, class_names, model_name):
    """Calculate comprehensive metrics for model evaluation"""
    metrics = {
        'model': model_name,
        'accuracy': accuracy_score(y_true, y_pred),
        'class_names': class_names,
        'classification_report': classification_report(y_true, y_pred, 
                                                     target_names=class_names, 
                                                     output_dict=True),
        'confusion_matrix': confusion_matrix(y_true, y_pred, labels=class_names),
        'y_true': y_true,  # Save true labels
        'y_pred': y_pred   # Save predicted labels
    }
    
    if len(class_names) > 1:
        y_true_bin = label_binarize([class_names.index(cls) for cls in y_true], 
                                   classes=range(len(class_names)))
        y_score_arr = np.array(y_score)
        
        # ROC AUC metrics
        roc_auc = {}
        for i in range(len(class_names)):
            fpr, tpr, _ = roc_curve(y_true_bin[:, i], y_score_arr[:, i])
            roc_auc[class_names[i]] = auc(fpr, tpr)
        
        metrics['roc_auc'] = roc_auc
        metrics['macro_avg_auc'] = np.mean(list(roc_auc.values()))
        
        # Precision-Recall metrics
        pr_auc = {}
        for i in range(len(class_names)):
            precision, recall, _ = precision_recall_curve(y_true_bin[:, i], y_score_arr[:, i])
            pr_auc[class_names[i]] = auc(recall, precision)
        
        metrics['pr_auc'] = pr_auc
        metrics['macro_avg_pr'] = np.mean(list(pr_auc.values()))
    
    return metrics
# ==============================================
# Core Pipeline Functions
# ==============================================

def load_dataset():
    """Load dataset from Roboflow with validation"""
    print("Downloading dental disease dataset...")
    try:
        rf = Roboflow(api_key="xs542BAQQ4nGbfcWQdcq")
        project = rf.workspace("rakshitha-m-s-nefdj").project("dentaldisease-pbral")
        dataset = project.version(3).download("folder")
        
        # Validate dataset structure
        required_folders = ['train', 'test', 'valid']
        for folder in required_folders:
            if not os.path.exists(os.path.join(dataset.location, folder)):
                raise ValueError(f"Missing required folder: {folder}")
        
        print("Dataset downloaded and validated successfully.")
        return dataset.location
    except Exception as e:
        print(f"Error loading dataset: {str(e)}")
        raise

def evaluate_model(model, test_path, class_names):
    """Comprehensive model evaluation with error handling"""
    y_true, y_pred, y_score, img_paths = [], [], [], []
    
    for class_name in class_names:
        class_dir = os.path.join(test_path, class_name)
        if not os.path.exists(class_dir):
            print(f"Warning: Missing test directory for class {class_name}")
            continue
            
        for img_file in os.listdir(class_dir):
            if img_file.lower().endswith(('.png', '.jpg', '.jpeg')):
                img_path = os.path.join(class_dir, img_file)
                try:
                    result = model.predict(img_path, augment=False, verbose=False)[0]  # Disable augmentation
                    
                    if hasattr(result, 'probs'):
                        pred_idx = result.probs.top1
                        y_true.append(class_name)
                        y_pred.append(class_names[pred_idx])
                        y_score.append(result.probs.data.tolist())
                        img_paths.append(img_path)
                except Exception as e:
                    print(f"Error processing image {img_path}: {str(e)}")
    
    if not y_true:
        raise ValueError("No valid predictions were generated during evaluation")
    
    return y_true, y_pred, y_score, img_paths

def zip_results():
    """Package all results into a downloadable zip file"""
    try:
        zip_path = os.path.join('/content', 'dental_classification_results.zip')
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, dirs, files in os.walk(config['results_dir']):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, config['results_dir'])
                    zipf.write(file_path, arcname)
        
        # Calculate and print zip size
        zip_size = os.path.getsize(zip_path) / (1024 * 1024)  # in MB
        print(f"Results package created: {zip_path} ({zip_size:.2f} MB)")
        
        from google.colab import files
        files.download(zip_path)
        return zip_path
    except Exception as e:
        print(f"Error creating zip file: {str(e)}")
        return None

def save_full_report(results, ablation_results, class_names, y_true, y_pred, accuracy):
    """Save comprehensive final report"""
    full_report_path = os.path.join(config['results_dir'], 'full_report.txt')
    with open(full_report_path, 'w') as f:
        f.write("=== COMPREHENSIVE DENTAL DISEASE CLASSIFICATION REPORT ===\n\n")
        f.write(f"Dataset: {len(class_names)} classes - {', '.join(class_names)}\n")
        f.write(f"Total Test Samples: {len(y_true)}\n")
        f.write(f"Final Model Accuracy: {accuracy:.4f}\n\n")
        
        f.write("\n=== FINAL MODEL PERFORMANCE ===\n")
        f.write(classification_report(y_true, y_pred, target_names=class_names, digits=4))
        
        f.write("\n=== ABLATION STUDY SUMMARY ===\n")
        best_model = max(ablation_results.values(), key=lambda x: x['accuracy'])
        f.write(f"\nBest Performing Model: {best_model['model']} (Accuracy: {best_model['accuracy']:.4f})\n")
        
        f.write("\nModel Comparison:\n")
        comparison_df = pd.DataFrame([{
            'Model': metrics['model'],
            'Accuracy': metrics['accuracy'],
            'Macro F1': metrics['classification_report']['macro avg']['f1-score'],
            'Macro AUC': metrics.get('macro_avg_auc', 'N/A'),
            'Macro PR AUC': metrics.get('macro_avg_pr', 'N/A')
        } for metrics in ablation_results.values()])
        f.write(comparison_df.to_string(index=False))
        
        f.write("\n\n=== TRAINING CONFIGURATION ===\n")
        for key, value in config.items():
            if key not in ['class_colors']:  # Skip large objects
                f.write(f"{key}: {value}\n")

# ==============================================
# Main Execution Pipeline
# ==============================================

if __name__ == "__main__":
    # Set random seeds for reproducibility
    random.seed(config['seed'])
    np.random.seed(config['seed'])
    torch.manual_seed(config['seed'])
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(config['seed'])
    
    # Load and validate dataset
    data_path = load_dataset()
    test_path = os.path.join(data_path, 'test')
    class_names = sorted([d for d in os.listdir(test_path) 
                         if os.path.isdir(os.path.join(test_path, d))])
    
    print(f"\nDataset contains {len(class_names)} classes: {', '.join(class_names)}")
    
    # Run ablation study
    print("\nStarting comprehensive ablation study...")
    ablation_results = run_ablation_study(data_path, test_path, class_names)
    
    # Train final enhanced model
    print("\nTraining final AGMS-FF enhanced model...")
    final_model = EnhancedYOLOClassifier(config['model_name'], device=config['device'])
    final_model.train(
        data=data_path,
        epochs=config['epochs'],
        imgsz=config['img_size'],
        batch=config['batch_size'],
        lr0=config['lr'],
        augment=False,  # Explicitly set to False
        seed=config['seed'],
        project=config['results_dir'],
        name='final_model',
        exist_ok=True
    )
    
    # Evaluate final model
    print("\nEvaluating final model...")
    y_true, y_pred, y_score, img_paths = evaluate_model(final_model.base_model, test_path, class_names)
    accuracy = accuracy_score(y_true, y_pred)
    
    # Generate all visualizations for final model
    print("\nGenerating final model visualizations...")
    plot_confusion_matrix(y_true, y_pred, class_names, "AGMS-FF Enhanced YOLOv11")
    plot_confusion_matrix(y_true, y_pred, class_names, "AGMS-FF Enhanced YOLOv11", normalize=True)
    
    if len(class_names) > 1:
        roc_auc = plot_roc_curves(y_true, y_score, class_names, "AGMS-FF Enhanced YOLOv11")
        pr_auc = plot_pr_curves(y_true, y_score, class_names, "AGMS-FF Enhanced YOLOv11")
    
    sample_indices = plot_sample_predictions(y_true, y_pred, img_paths, class_names,
                                           model_name="AGMS-FF Enhanced YOLOv11")
    
    # Save comprehensive report
    save_full_report(ablation_results, ablation_results, class_names, y_true, y_pred, accuracy)
    
    # Package all results
    print("\nPackaging all results...")
    zip_path = zip_results()
    if zip_path:
        print(f"\nAll results packaged and ready for download: {zip_path}")
    else:
        print(f"\nResults available in: {config['results_dir']}")
    
    print("\n=== EXPERIMENT COMPLETE ===") compact algorithm