redundancy.py

"""
Redundancy Detection for Neural Networks

Detects redundant filters, channels, and neurons that can be pruned or
compressed aggressively without loss of model capability.

Methods:
- Correlation analysis: Find highly correlated filters
- Activation sparsity: Find rarely-activated neurons
- Magnitude-based: Find small-magnitude weights
- Information-theoretic: Find low-entropy channels

High redundancy → Safe to prune/compress aggressively
Low redundancy → Keep to preserve model capacity
"""

import numpy as np
from typing import Dict, List, Tuple, Any, Optional, Set
import torch
import torch.nn as nn
from scipy.stats import pearsonr


class RedundancyDetector:
    """
    Detects redundant components in neural networks.
    
    Uses multiple heuristics to identify filters, channels, and neurons
    that contribute minimal unique information.
    """
    
    def __init__(
        self,
        correlation_threshold: float = 0.95,
        magnitude_threshold: float = 0.01,
        sparsity_threshold: float = 0.1
    ):
        """
        Initialize redundancy detector.
        
        Args:
            correlation_threshold: Correlation above which filters are redundant (default: 0.95)
            magnitude_threshold: Magnitude below which weights are redundant (default: 1%)
            sparsity_threshold: Activation frequency below which neurons are sparse (default: 10%)
        """
        self.correlation_threshold = correlation_threshold
        self.magnitude_threshold = magnitude_threshold
        self.sparsity_threshold = sparsity_threshold
        self.redundancy_info = {}
    
    def detect_redundant_channels(
        self,
        layer: nn.Module,
        method: str = 'correlation'
    ) -> Tuple[List[int], float]:
        """
        Detect redundant channels/filters in a convolutional layer.
        
        Args:
            layer: Conv2d layer
            method: Detection method ('correlation', 'magnitude', 'entropy')
        
        Returns:
            redundant_indices: List of redundant channel indices
            redundancy_score: Overall redundancy score (0-1)
        
        Example:
            >>> conv = nn.Conv2d(64, 64, 3)
            >>> redundant_idx, score = detector.detect_redundant_channels(conv)
            >>> print(f"Redundant channels: {redundant_idx}")
            >>> print(f"Redundancy score: {score:.2%}")
            Redundant channels: [5, 12, 34, 47]  # 4 out of 64
            Redundancy score: 6.25%  # Can prune these safely
        """
        if not isinstance(layer, nn.Conv2d):
            return [], 0.0
        
        weight = layer.weight.data.cpu().numpy()  # [out_ch, in_ch, kh, kw]
        out_channels = weight.shape[0]
        
        if method == 'correlation':
            return self._detect_by_correlation(weight)
        elif method == 'magnitude':
            return self._detect_by_magnitude(weight)
        elif method == 'entropy':
            return self._detect_by_entropy(weight)
        else:
            raise ValueError(f"Unknown method: {method}")
    
    def _detect_by_correlation(
        self,
        weight: np.ndarray
    ) -> Tuple[List[int], float]:
        """
        Detect redundant filters by correlation analysis.
        
        If two filters have correlation >0.95, they're learning the same features.
        One can be pruned without information loss.
        
        Args:
            weight: Conv2d weights [out_ch, in_ch, kh, kw]
        
        Returns:
            redundant_indices: Indices of redundant filters
            redundancy_score: Fraction of redundant filters
        """
        out_channels = weight.shape[0]
        
        # Flatten each filter
        filters = weight.reshape(out_channels, -1)
        
        # Compute pairwise correlations
        redundant = set()
        
        for i in range(out_channels):
            if i in redundant:
                continue
            
            for j in range(i + 1, out_channels):
                if j in redundant:
                    continue
                
                # Compute correlation
                try:
                    corr, _ = pearsonr(filters[i], filters[j])
                    
                    if abs(corr) > self.correlation_threshold:
                        # Mark j as redundant (keep i)
                        redundant.add(j)
                except:
                    # Handle constant filters
                    pass
        
        redundant_list = sorted(list(redundant))
        redundancy_score = len(redundant_list) / out_channels
        
        return redundant_list, redundancy_score
    
    def _detect_by_magnitude(
        self,
        weight: np.ndarray
    ) -> Tuple[List[int], float]:
        """
        Detect filters with very small magnitudes.
        
        Small-magnitude filters contribute little to output → Can prune.
        
        Args:
            weight: Conv2d weights [out_ch, in_ch, kh, kw]
        
        Returns:
            redundant_indices: Indices of low-magnitude filters
            redundancy_score: Fraction of low-magnitude filters
        """
        out_channels = weight.shape[0]
        
        # Compute L2 norm of each filter
        filter_norms = np.linalg.norm(
            weight.reshape(out_channels, -1),
            axis=1
        )
        
        # Normalize by max
        max_norm = filter_norms.max()
        if max_norm > 0:
            filter_norms_normalized = filter_norms / max_norm
        else:
            filter_norms_normalized = filter_norms
        
        # Find filters below threshold
        redundant = np.where(filter_norms_normalized < self.magnitude_threshold)[0]
        
        redundant_list = redundant.tolist()
        redundancy_score = len(redundant_list) / out_channels
        
        return redundant_list, redundancy_score
    
    def _detect_by_entropy(
        self,
        weight: np.ndarray
    ) -> Tuple[List[int], float]:
        """
        Detect filters with low entropy (redundant/repetitive patterns).
        
        Args:
            weight: Conv2d weights [out_ch, in_ch, kh, kw]
        
        Returns:
            redundant_indices: Indices of low-entropy filters
            redundancy_score: Fraction of low-entropy filters
        """
        from .entropy import compute_tensor_entropy
        
        out_channels = weight.shape[0]
        
        # Compute entropy of each filter
        entropies = []
        for i in range(out_channels):
            filter_weight = weight[i]
            entropy = compute_tensor_entropy(filter_weight)
            entropies.append(entropy)
        
        entropies = np.array(entropies)
        
        # Find filters with entropy below threshold (1.5 bits)
        redundant = np.where(entropies < 1.5)[0]
        
        redundant_list = redundant.tolist()
        redundancy_score = len(redundant_list) / out_channels
        
        return redundant_list, redundancy_score
    
    def analyze_activation_sparsity(
        self,
        model: nn.Module,
        test_inputs: torch.Tensor,
        layer_names: Optional[List[str]] = None
    ) -> Dict[str, Dict[str, Any]]:
        """
        Analyze activation sparsity across layers.
        
        Neurons that rarely activate (>90% zero) are redundant.
        
        Args:
            model: PyTorch model
            test_inputs: Test input batch
            layer_names: Specific layers to analyze (None = all Conv2d/Linear)
        
        Returns:
            sparsity_info: Dict mapping layer names to sparsity statistics
        
        Example:
            >>> model = SimpleCNN()
            >>> inputs = torch.randn(100, 3, 32, 32)  # 100 samples
            >>> sparsity = detector.analyze_activation_sparsity(model, inputs)
            >>> for name, info in sparsity.items():
            >>>     print(f"{name}: {info['sparsity']:.2%} sparse, "
            >>>           f"{len(info['dead_neurons'])} dead neurons")
        """
        model.eval()
        sparsity_info = {}
        
        # Hook to capture activations
        activation_storage = {}
        
        def capture_hook(name):
            def hook(module, input, output):
                if isinstance(output, torch.Tensor):
                    activation_storage[name] = output.detach().cpu()
            return hook
        
        # Register hooks
        hooks = []
        for name, module in model.named_modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                if layer_names is None or name in layer_names:
                    hook = module.register_forward_hook(capture_hook(name))
                    hooks.append(hook)
        
        # Run forward pass
        with torch.no_grad():
            _ = model(test_inputs)
        
        # Remove hooks
        for hook in hooks:
            hook.remove()
        
        # Analyze sparsity
        for name, activation in activation_storage.items():
            # Count zeros
            total_elements = activation.numel()
            zero_elements = (activation.abs() < 1e-6).sum().item()
            sparsity = zero_elements / total_elements
            
            # Identify dead neurons (always zero)
            if activation.ndim == 4:  # Conv2d [batch, channels, h, w]
                # Average over batch, h, w
                channel_activations = activation.mean(dim=[0, 2, 3])
                dead_neurons = (channel_activations.abs() < 1e-6).nonzero(as_tuple=True)[0].tolist()
            elif activation.ndim == 2:  # Linear [batch, features]
                # Average over batch
                neuron_activations = activation.mean(dim=0)
                dead_neurons = (neuron_activations.abs() < 1e-6).nonzero(as_tuple=True)[0].tolist()
            else:
                dead_neurons = []
            
            sparsity_info[name] = {
                'sparsity': sparsity,
                'dead_neurons': dead_neurons,
                'num_dead': len(dead_neurons),
                'is_redundant': sparsity > (1.0 - self.sparsity_threshold)
            }
        
        return sparsity_info
    
    def compute_redundancy_score(
        self,
        model: nn.Module,
        test_inputs: torch.Tensor = None
    ) -> Dict[str, float]:
        """
        Compute overall redundancy score for each layer.
        
        Combines correlation, magnitude, and entropy analysis.
        
        Args:
            model: PyTorch model
            test_inputs: Optional test inputs for activation analysis
        
        Returns:
            redundancy_scores: Dict mapping layer names to redundancy scores (0-1)
        
        Example:
            >>> model = SimpleCNN()
            >>> scores = detector.compute_redundancy_score(model)
            >>> for name, score in sorted(scores.items(), key=lambda x: -x[1]):
            >>>     print(f"{name}: {score:.2%} redundant")
            conv2: 15.3% redundant  # Can prune ~15% of filters
            conv1: 8.7% redundant
            fc:    3.2% redundant
        """
        redundancy_scores = {}
        
        print("\n🔍 Analyzing redundancy...")
        
        for name, module in model.named_modules():
            if isinstance(module, nn.Conv2d):
                print(f"   Analyzing {name}...", end=' ')
                
                # Combine multiple methods
                corr_redundant, corr_score = self.detect_redundant_channels(
                    module, method='correlation'
                )
                mag_redundant, mag_score = self.detect_redundant_channels(
                    module, method='magnitude'
                )
                ent_redundant, ent_score = self.detect_redundant_channels(
                    module, method='entropy'
                )
                
                # Average scores
                avg_score = (corr_score + mag_score + ent_score) / 3.0
                
                redundancy_scores[name] = avg_score
                
                print(f"✓ {avg_score:.2%} redundant")
        
        # Store for later access
        self.redundancy_info = redundancy_scores
        
        return redundancy_scores
    
    def get_pruning_mask(
        self,
        layer: nn.Module,
        method: str = 'correlation',
        prune_ratio: Optional[float] = None
    ) -> torch.Tensor:
        """
        Generate pruning mask for a layer.
        
        Args:
            layer: Layer to prune
            method: Detection method
            prune_ratio: Fraction to prune (None = auto-detect redundant)
        
        Returns:
            mask: Boolean mask [out_channels] (True = keep, False = prune)
        
        Example:
            >>> conv = nn.Conv2d(64, 64, 3)
            >>> mask = detector.get_pruning_mask(conv, prune_ratio=0.25)
            >>> # Apply mask to prune 25% of filters
            >>> pruned_conv = prune_conv_layer(conv, mask)
        """
        if not isinstance(layer, nn.Conv2d):
            return torch.ones(1, dtype=torch.bool)
        
        out_channels = layer.weight.shape[0]
        
        # Detect redundant channels
        redundant_indices, redundancy_score = self.detect_redundant_channels(
            layer, method=method
        )
        
        # Create mask
        mask = torch.ones(out_channels, dtype=torch.bool)
        
        if prune_ratio is not None:
            # Prune specified ratio
            num_to_prune = int(out_channels * prune_ratio)
            
            # Use magnitude-based ranking if not correlation
            if method != 'correlation':
                weight = layer.weight.data.cpu().numpy()
                filter_norms = np.linalg.norm(
                    weight.reshape(out_channels, -1),
                    axis=1
                )
                # Prune smallest magnitude filters
                prune_indices = np.argsort(filter_norms)[:num_to_prune]
            else:
                # Use detected redundant filters
                prune_indices = redundant_indices[:num_to_prune]
            
            mask[prune_indices] = False
        else:
            # Prune only detected redundant channels
            mask[redundant_indices] = False
        
        return mask


def detect_redundant_channels(
    layer: nn.Module,
    correlation_threshold: float = 0.95
) -> Tuple[List[int], float]:
    """
    Convenience function to detect redundant channels.
    
    Args:
        layer: Conv2d layer
        correlation_threshold: Correlation above which filters are redundant
    
    Returns:
        redundant_indices: List of redundant channel indices
        redundancy_score: Fraction of redundant channels
    """
    detector = RedundancyDetector(correlation_threshold=correlation_threshold)
    return detector.detect_redundant_channels(layer, method='correlation')


def create_redundancy_report(
    redundancy_scores: Dict[str, float],
    save_path: Optional[str] = None
) -> str:
    """
    Create human-readable redundancy analysis report.
    
    Args:
        redundancy_scores: Dict mapping layer names to redundancy scores
        save_path: Path to save report (optional)
    
    Returns:
        report: Formatted report string
    """
    report_lines = []
    report_lines.append("=" * 60)
    report_lines.append("REDUNDANCY ANALYSIS REPORT")
    report_lines.append("=" * 60)
    report_lines.append("")
    
    # Sort by redundancy (descending)
    sorted_layers = sorted(
        redundancy_scores.items(),
        key=lambda x: -x[1]
    )
    
    report_lines.append(f"{'Layer':<20} {'Redundancy':<12} {'Prune Potential'}")
    report_lines.append("-" * 60)
    
    for name, score in sorted_layers:
        if score > 0.3:
            potential = "High (>30%)"
        elif score > 0.15:
            potential = "Moderate (15-30%)"
        elif score > 0.05:
            potential = "Low (5-15%)"
        else:
            potential = "Minimal (<5%)"
        
        report_lines.append(
            f"{name:<20} {score:>8.2%}    {potential}"
        )
    
    report_lines.append("")
    report_lines.append("Interpretation:")
    report_lines.append("  High (>30%):      Significant redundancy → Aggressive pruning")
    report_lines.append("  Moderate (15-30%): Some redundancy → Standard pruning")
    report_lines.append("  Low (5-15%):       Minimal redundancy → Conservative pruning")
    report_lines.append("  Minimal (<5%):     No redundancy → Keep all filters")
    
    report = "\n".join(report_lines)
    
    if save_path:
        with open(save_path, 'w') as f:
            f.write(report)
        print(f"📄 Redundancy report saved to: {save_path}")
    
    return report


if __name__ == '__main__':
    print("=" * 60)
    print("Redundancy Detection")
    print("=" * 60)
    print("\nThis module implements redundancy detection:")
    print("  ✓ Correlation analysis: Find correlated filters")
    print("  ✓ Magnitude analysis: Find small-magnitude weights")
    print("  ✓ Entropy analysis: Find low-entropy channels")
    print("  ✓ Activation sparsity: Find rarely-activated neurons")
    print("\nRedundancy Levels:")
    print("  High (>30%):      Aggressive pruning safe")
    print("  Moderate (15-30%): Standard pruning recommended")
    print("  Low (5-15%):       Conservative pruning only")
    print("  Minimal (<5%):     Keep all filters")
    print("\nUsage:")
    print("  from passes.information_analysis import RedundancyDetector")
    print("  detector = RedundancyDetector()")
    print("  redundancy_scores = detector.compute_redundancy_score(model)")
    print("  report = create_redundancy_report(redundancy_scores)")