#!/usr/bin/env python3
"""
Entropy-Triggered Adaptive Depth (ETAD) MLP for CIFAR-10
- N-layer MLP with per-layer exit heads
- Learns entropy target profile per layer
- Skips computation when entropy drops below threshold
- Repeat batch iteration with fresh sample injection
- Single file: train + test
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
from collections import defaultdict

# ============================================================================
# CONFIGURATION
# ============================================================================
class Config:
    # Architecture
    INPUT_DIM = 32 * 32 * 3  # 3072
    HIDDEN_DIM = 150
    NUM_HIDDEN_LAYERS = 20
    OUTPUT_DIM = 10

    # Training
    BATCH_SIZE = 128
    EPOCHS = 30
    LR = 0.001
    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

    # Repeat batch iteration
    N_REPEATS = 10                  # Number of repeat iterations per epoch
    FRESH_RATIO_START = 1.0         # Starting fresh ratio (iteration 1)
    FRESH_RATIO_END = 0.3           # Ending fresh ratio (iteration N_REPEATS)

    # Entropy thresholds (learnable per layer)
    ENTROPY_TARGET_START = 1.8      # L0 target
    ENTROPY_TARGET_END = 0.2        # Final layer target
    MIN_DEPTH = 3                   # Minimum layers before allowing exit

    # Loss weights
    LAMBDA_ENTROPY = 0.1            # Entropy target loss weight
    LAMBDA_DEPTH = 0.01             # Penalize excessive depth

    # Logging
    LOG_INTERVAL = 100
    SEED = 42

# ============================================================================
# ENTROPY UTILS
# ============================================================================
def compute_entropy(logits, dim=-1, eps=1e-8):
    """Compute Shannon entropy of probability distribution"""
    probs = F.softmax(logits, dim=dim)
    return -torch.sum(probs * torch.log(probs + eps), dim=dim)

def batch_entropy(logits, eps=1e-8):
    """Mean entropy across batch"""
    return compute_entropy(logits, eps=eps).mean()

# ============================================================================
# MODEL: ETAD MLP WITH PER-LAYER EXIT HEADS
# ============================================================================
class ETADMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        # Input projection
        self.input_proj = nn.Linear(config.INPUT_DIM, config.HIDDEN_DIM)
        
        # Hidden layers
        self.hidden_layers = nn.ModuleList([
            nn.Linear(config.HIDDEN_DIM, config.HIDDEN_DIM)
            for _ in range(config.NUM_HIDDEN_LAYERS)
        ])
        
        # Per-layer exit heads (predict at any layer)
        self.exit_heads = nn.ModuleList([
            nn.Linear(config.HIDDEN_DIM, config.OUTPUT_DIM)
            for _ in range(config.NUM_HIDDEN_LAYERS + 1)  # +1 for after input_proj
        ])
        
        # Learnable entropy thresholds per layer
        # Initialized to decay from START to END
        self.entropy_thresholds = nn.Parameter(
            torch.linspace(
                config.ENTROPY_TARGET_START,
                config.ENTROPY_TARGET_END,
                config.NUM_HIDDEN_LAYERS + 1
            )
        )
        
        # Layer usage counters (for tracking, not trainable)
        self.layer_exit_counts = defaultdict(int)
        
    def forward(self, x, return_all=False, force_depth=None):
        """
        Args:
            x: Input tensor [B, 3072]
            return_all: If True, return all layer outputs (for analysis)
            force_depth: If set, exit at this layer regardless of entropy
            
        Returns:
            logits: Final predictions
            exit_layers: Which layer each sample exited at
            entropies: Entropy at each layer (if return_all)
        """
        batch_size = x.size(0)
        
        # Track exit layer for each sample (-1 = not exited yet)
        exit_layers = torch.full((batch_size,), -1, dtype=torch.long, device=x.device)
        final_logits = None
        all_entropies = []
        all_logits = []
        
        # Layer 0: After input projection
        h = F.relu(self.input_proj(x))
        logits = self.exit_heads[0](h)
        entropy = compute_entropy(logits)  # [B]
        all_entropies.append(entropy.clone())
        all_logits.append(logits.clone())
        
        # Check exit condition for layer 0
        if force_depth is None:
            should_exit = (entropy <= self.entropy_thresholds[0]) & (exit_layers == -1)
        else:
            should_exit = torch.zeros_like(should_exit, dtype=torch.bool) if 0 < force_depth else torch.ones_like(should_exit, dtype=torch.bool)
            
        exited = should_exit & (exit_layers == -1)
        exit_layers[exited] = 0
        if final_logits is None:
            final_logits = logits.clone()
        final_logits = torch.where(
            exited.unsqueeze(1),
            logits,
            final_logits
        )
        
        # Hidden layers
        for l, layer in enumerate(self.hidden_layers):
            # Skip computation for samples that already exited
            active_mask = exit_layers == -1
            if not active_mask.any():
                break
                
            h = F.relu(layer(h))
            logits = self.exit_heads[l + 1](h)
            entropy = compute_entropy(logits)
            all_entropies.append(entropy.clone())
            all_logits.append(logits.clone())
            
            # Check exit condition
            layer_idx = l + 1
            if force_depth is None:
                # Don't exit before MIN_DEPTH
                can_exit = layer_idx >= self.config.MIN_DEPTH
                should_exit = (entropy <= self.entropy_thresholds[layer_idx]) & can_exit
            else:
                should_exit = torch.zeros_like(entropy, dtype=torch.bool) if layer_idx < force_depth else torch.ones_like(entropy, dtype=torch.bool)
            
            exited = should_exit & active_mask
            exit_layers[exited] = layer_idx
            
            # Update final logits for exited samples
            final_logits = torch.where(
                exited.unsqueeze(1),
                logits,
                final_logits
            )
        
        # Force exit remaining samples at final layer
        remaining = exit_layers == -1
        exit_layers[remaining] = self.config.NUM_HIDDEN_LAYERS
        final_logits = torch.where(
            remaining.unsqueeze(1),
            logits,
            final_logits
        )
        
        # Update exit counts (for analysis)
        if self.training:
            for idx in exit_layers.cpu().numpy():
                self.layer_exit_counts[idx] += 1
        
        if return_all:
            return final_logits, exit_layers, all_entropies, all_logits
        
        return final_logits, exit_layers
    
    def get_entropy_profile(self):
        """Return current entropy thresholds (the learned profile)"""
        return self.entropy_thresholds.detach().cpu().numpy()
    
    def get_exit_distribution(self):
        """Return distribution of exit layers"""
        total = sum(self.layer_exit_counts.values())
        if total == 0:
            return {}
        return {k: v / total for k, v in sorted(self.layer_exit_counts.items())}
    
    def reset_exit_counts(self):
        """Reset exit counters for new evaluation period"""
        self.layer_exit_counts.clear()

# ============================================================================
# LOSS FUNCTIONS
# ============================================================================
class ETADLoss(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.ce = nn.CrossEntropyLoss()
        
    def forward(self, logits, targets, exit_layers, entropies=None):
        """
        Args:
            logits: Final predictions [B, 10]
            targets: Ground truth [B]
            exit_layers: Which layer each sample exited at [B]
            entropies: List of entropy tensors per layer (optional)
        """
        # 1. Task loss (cross-entropy)
        task_loss = self.ce(logits, targets)
        
        # 2. Entropy target loss (encourage entropy to match thresholds at exit)
        entropy_loss = 0
        if entropies is not None:
            for l, entropy_l in enumerate(entropies):
                # Only penalize layers where samples actually exited
                exited_mask = (exit_layers == l)
                if exited_mask.any():
                    target_entropy = self.config.ENTROPY_TARGET_START * \
                                    np.exp(-l / 3) + self.config.ENTROPY_TARGET_END
                    # Penalize if entropy at exit is too far from target
                    entropy_diff = (entropy_l[exited_mask] - target_entropy).abs()
                    entropy_loss += entropy_diff.mean()
            
            entropy_loss /= len(entropies)
        
        # 3. Depth penalty (encourage early exit when possible)
        avg_depth = exit_layers.float().mean()
        depth_penalty = avg_depth / self.config.NUM_HIDDEN_LAYERS
        
        # Total loss
        total_loss = (
            task_loss +
            self.config.LAMBDA_ENTROPY * entropy_loss +
            self.config.LAMBDA_DEPTH * depth_penalty
        )
        
        return total_loss, {
            'task_loss': task_loss.item(),
            'entropy_loss': entropy_loss.item() if entropies is not None else 0,
            'depth_penalty': depth_penalty.item(),
            'avg_depth': avg_depth.item()
        }

# ============================================================================
# REPEAT BATCH TRAINER WITH FRESH SAMPLE INJECTION
# ============================================================================
class RepeatBatchTrainer:
    """
    Manages repeat batch iteration with fresh sample injection.
    
    Each epoch consists of N_REPEAT iterations where:
    - Early iterations: mostly fresh samples (explore)
    - Later iterations: mostly hard samples from history (exploit)
    
    This prevents memorization while focusing on difficult examples.
    """
    def __init__(self, dataset_size, batch_size, n_repeats=10, 
                 fresh_ratio_start=1.0, fresh_ratio_end=0.3, seed=42):
        self.dataset_size = dataset_size
        self.batch_size = batch_size
        self.n_repeats = n_repeats
        self.fresh_ratio_start = fresh_ratio_start
        self.fresh_ratio_end = fresh_ratio_end
        
        # Track per-sample difficulty (EMA of loss)
        self.sample_losses = torch.zeros(dataset_size)
        self.sample_counts = torch.zeros(dataset_size, dtype=torch.long)
        
        # Track which samples have been seen this epoch
        self.seen_mask = torch.zeros(dataset_size, dtype=torch.bool)
        
        # RNG for reproducibility
        self.rng = np.random.RandomState(seed)
        
    def get_fresh_ratio(self, iteration):
        """Linear interpolation from fresh_ratio_start to fresh_ratio_end"""
        if self.n_repeats <= 1:
            return self.fresh_ratio_start
        t = iteration / (self.n_repeats - 1)
        return self.fresh_ratio_start * (1 - t) + self.fresh_ratio_end * t
    
    def get_next_indices(self, iteration):
        """
        Get batch indices mixing hard and fresh samples.
        
        Args:
            iteration: Current repeat iteration (0 to n_repeats-1)
            
        Returns:
            indices: Tensor of sample indices for this batch
        """
        fresh_ratio = self.get_fresh_ratio(iteration)
        n_fresh = int(self.batch_size * fresh_ratio)
        n_hard = self.batch_size - n_fresh
        
        # Fresh samples: not yet seen this epoch
        unseen_indices = (~self.seen_mask).nonzero(as_tuple=True)[0]
        if len(unseen_indices) < n_fresh:
            # Fallback: use all unseen + some random
            fresh_indices = unseen_indices
            n_extra = n_fresh - len(fresh_indices)
            extra_indices = torch.randperm(self.dataset_size)[:n_extra]
            fresh_indices = torch.cat([fresh_indices, extra_indices])
        else:
            perm = torch.randperm(len(unseen_indices))
            fresh_indices = unseen_indices[perm[:n_fresh]]
        
        # Mark fresh samples as seen
        self.seen_mask[fresh_indices] = True
        
        # Hard samples: highest historical loss, excluding fresh ones
        if n_hard > 0:
            # Mask out fresh samples
            candidate_losses = self.sample_losses.clone()
            candidate_losses[fresh_indices] = -1.0  # Exclude fresh
            
            # Get top-k hardest
            n_candidates = min(n_hard * 10, self.dataset_size)
            top_indices = torch.topk(candidate_losses, n_candidates).indices
            
            # Random selection from hard candidates (with loss-proportional weighting)
            hard_losses = candidate_losses[top_indices]
            hard_losses = torch.clamp(hard_losses, min=0.0)
            weights = hard_losses + 0.1  # Avoid zero probability
            weights = weights / weights.sum()
            
            hard_selection = self.rng.choice(
                top_indices.numpy(), 
                size=n_hard, 
                p=weights.numpy(),
                replace=True
            )
            hard_indices = torch.from_numpy(hard_selection)
        else:
            hard_indices = torch.empty(0, dtype=torch.long)
        
        # Combine and shuffle
        indices = torch.cat([fresh_indices, hard_indices])
        perm = torch.randperm(len(indices))
        return indices[perm]
    
    def update_sample_losses(self, indices, losses, alpha=0.3):
        """
        Update per-sample loss history with EMA.
        
        Args:
            indices: Sample indices that were in the batch
            losses: Per-sample loss values
            alpha: EMA update rate (higher = more weight to recent)
        """
        self.sample_losses[indices] = (
            alpha * losses + (1 - alpha) * self.sample_losses[indices]
        )
        self.sample_counts[indices] += 1
    
    def reset_epoch(self):
        """Reset seen mask for new epoch"""
        self.seen_mask.zero_()
    
    def get_difficulty_stats(self):
        """Return statistics about sample difficulty distribution"""
        seen = self.sample_counts > 0
        return {
            'mean_loss': self.sample_losses[seen].mean().item() if seen.any() else 0,
            'max_loss': self.sample_losses[seen].max().item() if seen.any() else 0,
            'n_seen': seen.sum().item(),
            'n_seen_multiple': (self.sample_counts[seen] > 1).sum().item(),
        }

# ============================================================================
# TRAINING & EVALUATION
# ============================================================================
def get_cifar10_loaders(config):
    """Load CIFAR-10 with proper preprocessing"""
    transform_train = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                           (0.2023, 0.1994, 0.2010)),
        transforms.Lambda(lambda x: x.view(-1))  # Flatten to 3072
    ])

    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                           (0.2023, 0.1994, 0.2010)),
        transforms.Lambda(lambda x: x.view(-1))
    ])

    train_dataset = datasets.CIFAR10(
        root='../data', train=True, download=True, transform=transform_train
    )
    test_dataset = datasets.CIFAR10(
        root='../data', train=False, download=True, transform=transform_test
    )

    train_loader = DataLoader(train_dataset, batch_size=config.BATCH_SIZE,
                             shuffle=True, num_workers=2)
    test_loader = DataLoader(test_dataset, batch_size=config.BATCH_SIZE,
                            shuffle=False, num_workers=2)

    return train_loader, test_loader, train_dataset

def train_epoch_with_repeats(model, train_dataset, optimizer, loss_fn, config, epoch):
    """
    Train one epoch using repeat batch iteration with fresh sample injection.
    
    Each epoch consists of N_REPEATS iterations over batches of size BATCH_SIZE.
    Early iterations use mostly fresh samples; later iterations focus on hard examples.
    """
    model.train()
    
    # Initialize trainer
    trainer = RepeatBatchTrainer(
        dataset_size=len(train_dataset),
        batch_size=config.BATCH_SIZE,
        n_repeats=config.N_REPEATS,
        fresh_ratio_start=config.FRESH_RATIO_START,
        fresh_ratio_end=config.FRESH_RATIO_END,
        seed=config.SEED + epoch  # Different seed per epoch
    )
    trainer.reset_epoch()
    
    total_loss = 0
    correct = 0
    total = 0
    total_batches = 0
    metrics = defaultdict(float)
    
    # Convert dataset to tensors for direct indexing
    all_data = torch.stack([train_dataset[i][0].view(-1) for i in range(len(train_dataset))])
    all_targets = torch.tensor([train_dataset[i][1] for i in range(len(train_dataset))])
    
    print(f"\n🔄 Epoch {epoch}: {config.N_REPEATS} repeat iterations")
    
    for iteration in range(config.N_REPEATS):
        # Get batch indices (mix of fresh and hard samples)
        indices = trainer.get_next_indices(iteration)
        
        # Fetch data
        data = all_data[indices].to(config.DEVICE)
        target = all_targets[indices].to(config.DEVICE)
        
        optimizer.zero_grad()
        
        # Forward with entropy tracking
        logits, exit_layers, entropies, _ = model(data, return_all=True)
        
        # Compute per-sample losses for tracking
        per_sample_losses = F.cross_entropy(logits, target, reduction='none')
        
        # Update sample difficulty history
        trainer.update_sample_losses(indices, per_sample_losses.detach().cpu())
        
        # Compute aggregate loss
        loss, batch_metrics = loss_fn(logits, target, exit_layers, entropies)
        
        loss.backward()
        optimizer.step()
        
        # Track metrics
        total_loss += loss.item()
        pred = logits.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)
        total_batches += 1
        
        for k, v in batch_metrics.items():
            metrics[k] += v
        
        # Print iteration progress
        fresh_ratio = trainer.get_fresh_ratio(iteration)
        diff_stats = trainer.get_difficulty_stats()
        
        if iteration % max(1, config.N_REPEATS // 5) == 0 or iteration == config.N_REPEATS - 1:
            print(f'  Iter {iteration:2d}/{config.N_REPEATS} '
                  f'(fresh={fresh_ratio:.0%}) | '
                  f'Loss: {loss.item():.4f} | '
                  f'Acc: {100.*correct/total:.2f}% | '
                  f'AvgDepth: {batch_metrics["avg_depth"]:.2f} | '
                  f'MeanDiff: {diff_stats["mean_loss"]:.3f}')
    
    # Reset exit counts for next epoch
    model.reset_exit_counts()
    
    return {
        'loss': total_loss / total_batches,
        'accuracy': 100. * correct / total,
        **{k: v / total_batches for k, v in metrics.items()},
        'difficulty_stats': trainer.get_difficulty_stats()
    }

def train_epoch(model, loader, optimizer, loss_fn, config, epoch):
    """Train one epoch (original version, kept for compatibility)"""
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    metrics = defaultdict(float)

    for batch_idx, (data, target) in enumerate(loader):
        data, target = data.to(config.DEVICE), target.to(config.DEVICE)

        optimizer.zero_grad()

        # Forward with entropy tracking
        logits, exit_layers, entropies, _ = model(data, return_all=True)

        # Compute loss
        loss, batch_metrics = loss_fn(logits, target, exit_layers, entropies)

        loss.backward()
        optimizer.step()

        # Track metrics
        total_loss += loss.item()
        pred = logits.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)

        for k, v in batch_metrics.items():
            metrics[k] += v

        if batch_idx % config.LOG_INTERVAL == 0:
            print(f'Epoch {epoch} [{batch_idx}/{len(loader)}] '
                  f'Loss: {loss.item():.4f} Acc: {100.*correct/total:.2f}% '
                  f'AvgDepth: {batch_metrics["avg_depth"]:.2f}')

    # Reset exit counts for next epoch
    model.reset_exit_counts()

    return {
        'loss': total_loss / len(loader),
        'accuracy': 100. * correct / total,
        **{k: v / len(loader) for k, v in metrics.items()}
    }

@torch.no_grad()
def evaluate(model, loader, loss_fn, config):
    """Evaluate model"""
    model.eval()
    total_loss = 0
    correct = 0
    total = 0
    metrics = defaultdict(float)
    exit_distribution = defaultdict(int)
    accuracy_per_layer = defaultdict(lambda: {'correct': 0, 'total': 0})
    
    for data, target in loader:
        data, target = data.to(config.DEVICE), target.to(config.DEVICE)
        
        logits, exit_layers, entropies, all_logits = model(data, return_all=True)
        
        loss, batch_metrics = loss_fn(logits, target, exit_layers, entropies)
        
        total_loss += loss.item()
        pred = logits.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)
        
        # Track exit distribution
        for idx in exit_layers.cpu().numpy():
            exit_distribution[idx] += 1
        
        # Track accuracy per exit layer
        for i, (l, p, t) in enumerate(zip(exit_layers, pred, target)):
            layer_idx = l.item()
            accuracy_per_layer[layer_idx]['total'] += 1
            if p.item() == t.item():
                accuracy_per_layer[layer_idx]['correct'] += 1
        
        for k, v in batch_metrics.items():
            metrics[k] += v
    
    model.reset_exit_counts()
    
    # Compute per-layer accuracy
    layer_accuracy = {}
    for l, stats in accuracy_per_layer.items():
        if stats['total'] > 0:
            layer_accuracy[l] = 100. * stats['correct'] / stats['total']
    
    return {
        'loss': total_loss / len(loader),
        'accuracy': 100. * correct / total,
        'exit_distribution': {k: v / total for k, v in sorted(exit_distribution.items())},
        'layer_accuracy': layer_accuracy,
        **{k: v / len(loader) for k, v in metrics.items()}
    }

# ============================================================================
# MAIN
# ============================================================================
def main(use_repeats=True):
    """
    Main training loop.
    
    Args:
        use_repeats: If True, use repeat batch iteration with fresh sample injection.
                    If False, use standard DataLoader iteration.
    """
    # Set seed
    torch.manual_seed(Config.SEED)
    np.random.seed(Config.SEED)

    print(f"🚀 ETAD MLP for CIFAR-10")
    print(f"   Device: {Config.DEVICE}")
    print(f"   Layers: {Config.NUM_HIDDEN_LAYERS} hidden + exit heads")
    print(f"   Min exit depth: {Config.MIN_DEPTH}")
    if use_repeats:
        print(f"   🔄 Repeat iterations: {Config.N_REPEATS}")
        print(f"   📊 Fresh ratio: {Config.FRESH_RATIO_START:.0%} → {Config.FRESH_RATIO_END:.0%}")
    print()

    # Load data
    train_loader, test_loader, train_dataset = get_cifar10_loaders(Config)

    # Initialize model
    model = ETADMLP(Config).to(Config.DEVICE)
    optimizer = optim.Adam(model.parameters(), lr=Config.LR)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
    loss_fn = ETADLoss(Config)

    # Training loop
    print("📈 Training...")
    print("=" * 80)

    best_acc = 0
    for epoch in range(1, Config.EPOCHS + 1):
        # Train
        if use_repeats:
            train_metrics = train_epoch_with_repeats(
                model, train_dataset, optimizer, loss_fn, Config, epoch
            )
        else:
            train_metrics = train_epoch(
                model, train_loader, optimizer, loss_fn, Config, epoch
            )

        # Evaluate
        test_metrics = evaluate(model, test_loader, loss_fn, Config)

        # Print summary
        print(f"\n{'='*80}")
        print(f"EPOCH {epoch:2d} SUMMARY")
        print(f"{'='*80}")
        print(f"Train Loss: {train_metrics['loss']:.4f} | Acc: {train_metrics['accuracy']:.2f}%")
        print(f"Test  Loss: {test_metrics['loss']:.4f} | Acc: {test_metrics['accuracy']:.2f}%")
        print(f"Avg Depth: {test_metrics['avg_depth']:.2f} / {Config.NUM_HIDDEN_LAYERS}")
        print(f"Entropy Profile: {model.get_entropy_profile().round(2)}")
        print(f"Exit Distribution: {test_metrics['exit_distribution']}")
        print(f"Layer Accuracy: {test_metrics['layer_accuracy']}")
        
        if use_repeats and 'difficulty_stats' in train_metrics:
            diff = train_metrics['difficulty_stats']
            print(f"Difficulty Stats: mean_loss={diff['mean_loss']:.3f} | "
                  f"n_seen={diff['n_seen']} | n_repeat={diff['n_seen_multiple']}")

        # Save best
        if test_metrics['accuracy'] > best_acc:
            best_acc = test_metrics['accuracy']
            torch.save({
                'model': model.state_dict(),
                'entropy_thresholds': model.get_entropy_profile(),
                'epoch': epoch,
            }, 'etad_mlp_best.pt')

        scheduler.step()

    print(f"\n{'='*80}")
    print(f"✅ Training Complete! Best Test Acc: {best_acc:.2f}%")
    print(f"{'='*80}")

    # Final analysis
    print("\n📊 FINAL ANALYSIS")
    print("=" * 80)

    # Show which layers are actually used
    exit_dist = test_metrics['exit_distribution']
    print(f"Exit Layer Distribution:")
    for layer, pct in exit_dist.items():
        bar = '█' * int(pct * 50)
        print(f"  Layer {layer:2d}: {pct*100:5.1f}% {bar}")

    # Show accuracy vs depth
    print(f"\nAccuracy by Exit Layer:")
    for layer, acc in sorted(test_metrics['layer_accuracy'].items()):
        print(f"  Layer {layer:2d}: {acc:.2f}%")

    # Compute efficiency gain
    avg_depth = test_metrics['avg_depth']
    efficiency = 1 - (avg_depth / Config.NUM_HIDDEN_LAYERS)
    print(f"\n⚡ Efficiency Gain: {efficiency*100:.1f}% fewer layers on average")
    print(f"   (Avg depth: {avg_depth:.2f} vs max {Config.NUM_HIDDEN_LAYERS})")

if __name__ == '__main__':
    main(use_repeats=True)