src/interpretability/advanced_xai.py

"""
Advanced Explainable AI (XAI) Module for MorphGuard
Implements state-of-the-art interpretability techniques beyond basic gradient saliency
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import cv2
from typing import Dict, List, Tuple, Optional, Any, Union
import logging
from dataclasses import dataclass
from datetime import datetime
import matplotlib.pyplot as plt
import seaborn as sns
from captum.attr import (
    IntegratedGradients,
    GradientShap,
    DeepLift,
    LayerConductance,
    LayerActivation,
    LayerGradientXActivation,
    Saliency,
    InputXGradient,
    GuidedBackprop,
    GuidedGradCam,
    Occlusion,
    KernelShap,
    Lime,
    ShapleyValueSampling
)
from captum.attr._utils.visualization import visualize_image_attr
from captum.concept import TCAV
from captum.robust import PGD, FGSM
import shap
import lime
from lime import lime_image
import json

logger = logging.getLogger(__name__)

@dataclass
class ExplanationResult:
    """Comprehensive explanation result with multiple XAI techniques"""
    method: str
    attribution_map: np.ndarray
    confidence: float
    interpretation_score: float
    feature_importance: Dict[str, float]
    concept_activations: Optional[Dict[str, float]]
    counterfactual_regions: Optional[List[Tuple[int, int, int, int]]]
    textual_explanation: str
    timestamp: datetime
    processing_time_ms: float

@dataclass
class ConceptActivation:
    """TCAV-based concept activation"""
    concept_name: str
    activation_strength: float
    statistical_significance: float
    directional_derivative: float

class AdvancedXAIEngine:
    """
    Advanced Explainable AI engine implementing multiple state-of-the-art techniques
    """
    
    def __init__(self, model: nn.Module, device: str = 'cuda'):
        self.model = model
        self.device = device
        self.model.eval()
        
        # Initialize different attribution methods
        self.integrated_gradients = IntegratedGradients(self.model)
        self.gradient_shap = GradientShap(self.model)
        self.deep_lift = DeepLift(self.model)
        self.saliency = Saliency(self.model)
        self.input_x_gradient = InputXGradient(self.model)
        self.guided_backprop = GuidedBackprop(self.model)
        self.occlusion = Occlusion(self.model)
        
        # SHAP explainer
        self.shap_explainer = None
        
        # LIME explainer
        self.lime_explainer = lime_image.LimeImageExplainer()
        
        # Concept vectors for TCAV
        self.concept_vectors = {}
        
        logger.info("Advanced XAI Engine initialized with multiple attribution methods")
    
    def explain_prediction(
        self, 
        image: torch.Tensor, 
        target_class: int = 1,
        methods: Optional[List[str]] = None,
        baseline_strategy: str = 'zero'
    ) -> Dict[str, ExplanationResult]:
        """
        Generate comprehensive explanations using multiple XAI techniques
        
        Args:
            image: Input image tensor [1, C, H, W]
            target_class: Target class for explanation (1 for morph)
            methods: List of methods to use. If None, uses all available
            baseline_strategy: Baseline strategy for integrated gradients
            
        Returns:
            Dictionary of explanation results by method name
        """
        if methods is None:
            methods = [
                'integrated_gradients', 'gradient_shap', 'deep_lift',
                'saliency', 'guided_backprop', 'occlusion', 'shap', 'lime'
            ]
        
        explanations = {}
        image = image.to(self.device)
        
        # Get model prediction for confidence
        with torch.no_grad():
            prediction = self.model(image)
            confidence = float(prediction.sigmoid().cpu())
        
        for method in methods:
            start_time = time.time()
            try:
                explanation = self._generate_explanation(
                    image, target_class, method, baseline_strategy, confidence
                )
                explanation.processing_time_ms = (time.time() - start_time) * 1000
                explanations[method] = explanation
                logger.debug(f"Generated {method} explanation in {explanation.processing_time_ms:.2f}ms")
            except Exception as e:
                logger.error(f"Failed to generate {method} explanation: {e}")
        
        return explanations
    
    def _generate_explanation(
        self,
        image: torch.Tensor,
        target_class: int,
        method: str,
        baseline_strategy: str,
        confidence: float
    ) -> ExplanationResult:
        """Generate explanation using specific method"""
        
        if method == 'integrated_gradients':
            return self._integrated_gradients_explanation(
                image, target_class, baseline_strategy, confidence
            )
        elif method == 'gradient_shap':
            return self._gradient_shap_explanation(image, target_class, confidence)
        elif method == 'deep_lift':
            return self._deep_lift_explanation(image, target_class, confidence)
        elif method == 'saliency':
            return self._saliency_explanation(image, target_class, confidence)
        elif method == 'guided_backprop':
            return self._guided_backprop_explanation(image, target_class, confidence)
        elif method == 'occlusion':
            return self._occlusion_explanation(image, target_class, confidence)
        elif method == 'shap':
            return self._shap_explanation(image, target_class, confidence)
        elif method == 'lime':
            return self._lime_explanation(image, target_class, confidence)
        else:
            raise ValueError(f"Unknown explanation method: {method}")
    
    def _integrated_gradients_explanation(
        self, image: torch.Tensor, target_class: int, baseline_strategy: str, confidence: float
    ) -> ExplanationResult:
        """Enhanced Integrated Gradients with multiple baselines"""
        
        # Generate multiple baselines
        baselines = self._generate_baselines(image, baseline_strategy)
        attributions_list = []
        
        for baseline in baselines:
            attr = self.integrated_gradients.attribute(
                image, 
                baselines=baseline,
                target=target_class,
                n_steps=100,
                internal_batch_size=1
            )
            attributions_list.append(attr)
        
        # Ensemble attributions
        final_attribution = torch.mean(torch.stack(attributions_list), dim=0)
        attribution_map = self._process_attribution(final_attribution)
        
        # Calculate interpretation score
        interpretation_score = self._calculate_interpretation_score(
            final_attribution, image
        )
        
        # Extract feature importance
        feature_importance = self._extract_feature_importance(final_attribution)
        
        # Generate textual explanation
        textual_explanation = self._generate_textual_explanation(
            attribution_map, feature_importance, confidence, "Integrated Gradients"
        )
        
        return ExplanationResult(
            method="Integrated Gradients",
            attribution_map=attribution_map,
            confidence=confidence,
            interpretation_score=interpretation_score,
            feature_importance=feature_importance,
            concept_activations=None,
            counterfactual_regions=None,
            textual_explanation=textual_explanation,
            timestamp=datetime.now()
        )
    
    def _gradient_shap_explanation(
        self, image: torch.Tensor, target_class: int, confidence: float
    ) -> ExplanationResult:
        """Gradient SHAP explanation with noise baseline"""
        
        # Generate noise baseline
        noise_baseline = torch.randn_like(image) * 0.1
        
        attribution = self.gradient_shap.attribute(
            image,
            baselines=noise_baseline,
            target=target_class,
            n_samples=50,
            stdevs=0.1
        )
        
        attribution_map = self._process_attribution(attribution)
        interpretation_score = self._calculate_interpretation_score(attribution, image)
        feature_importance = self._extract_feature_importance(attribution)
        
        textual_explanation = self._generate_textual_explanation(
            attribution_map, feature_importance, confidence, "Gradient SHAP"
        )
        
        return ExplanationResult(
            method="Gradient SHAP",
            attribution_map=attribution_map,
            confidence=confidence,
            interpretation_score=interpretation_score,
            feature_importance=feature_importance,
            concept_activations=None,
            counterfactual_regions=None,
            textual_explanation=textual_explanation,
            timestamp=datetime.now()
        )
    
    def _shap_explanation(
        self, image: torch.Tensor, target_class: int, confidence: float
    ) -> ExplanationResult:
        """SHAP explanation using Deep Explainer"""
        
        # Initialize SHAP explainer if not done
        if self.shap_explainer is None:
            # Create background dataset from random samples
            background = torch.randn(10, *image.shape[1:]).to(self.device)
            self.shap_explainer = shap.DeepExplainer(self.model, background)
        
        # Generate SHAP values
        shap_values = self.shap_explainer.shap_values(image)
        
        if isinstance(shap_values, list):
            shap_values = shap_values[target_class]
        
        attribution_map = self._process_attribution(torch.tensor(shap_values))
        interpretation_score = np.mean(np.abs(shap_values))
        feature_importance = self._extract_feature_importance(torch.tensor(shap_values))
        
        textual_explanation = self._generate_textual_explanation(
            attribution_map, feature_importance, confidence, "SHAP"
        )
        
        return ExplanationResult(
            method="SHAP",
            attribution_map=attribution_map,
            confidence=confidence,
            interpretation_score=float(interpretation_score),
            feature_importance=feature_importance,
            concept_activations=None,
            counterfactual_regions=None,
            textual_explanation=textual_explanation,
            timestamp=datetime.now()
        )
    
    def _lime_explanation(
        self, image: torch.Tensor, target_class: int, confidence: float
    ) -> ExplanationResult:
        """LIME explanation using superpixel segmentation"""
        
        # Convert to numpy for LIME
        img_np = image.squeeze().cpu().numpy().transpose(1, 2, 0)
        if img_np.min() < 0:  # Denormalize if needed
            img_np = (img_np + 1) / 2
        img_np = (img_np * 255).astype(np.uint8)
        
        def predict_fn(images):
            """Prediction function for LIME"""
            batch = []
            for img in images:
                # Normalize and convert to tensor
                img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0
                img_tensor = (img_tensor - 0.5) / 0.5  # Normalize to [-1, 1]
                batch.append(img_tensor)
            
            batch_tensor = torch.stack(batch).to(self.device)
            with torch.no_grad():
                predictions = self.model(batch_tensor).sigmoid().cpu().numpy()
            return predictions
        
        # Generate LIME explanation
        explanation = self.lime_explainer.explain_instance(
            img_np,
            predict_fn,
            top_labels=2,
            hide_color=0,
            num_samples=1000
        )
        
        # Get mask for target class
        temp, mask = explanation.get_image_and_mask(
            target_class, positive_only=False, num_features=10, hide_rest=False
        )
        
        attribution_map = mask.astype(np.float32)
        interpretation_score = np.std(mask)  # Variability as interpretation score
        
        # Feature importance from LIME segments
        feature_importance = {}
        for i, (feature_id, weight) in enumerate(explanation.local_exp[target_class]):
            feature_importance[f"superpixel_{feature_id}"] = float(weight)
        
        textual_explanation = self._generate_textual_explanation(
            attribution_map, feature_importance, confidence, "LIME"
        )
        
        return ExplanationResult(
            method="LIME",
            attribution_map=attribution_map,
            confidence=confidence,
            interpretation_score=interpretation_score,
            feature_importance=feature_importance,
            concept_activations=None,
            counterfactual_regions=None,
            textual_explanation=textual_explanation,
            timestamp=datetime.now()
        )
    
    def analyze_concepts(
        self, 
        image: torch.Tensor, 
        concept_dataset: Dict[str, torch.Tensor]
    ) -> Dict[str, ConceptActivation]:
        """
        Perform Testing with Concept Activation Vectors (TCAV) analysis
        
        Args:
            image: Input image tensor
            concept_dataset: Dictionary of concept name to concept examples
            
        Returns:
            Dictionary of concept activations
        """
        concept_activations = {}
        
        # This is a simplified TCAV implementation
        # In practice, you'd need to train linear classifiers for each concept
        for concept_name, concept_examples in concept_dataset.items():
            try:
                # Extract activations from a specific layer
                layer_name = 'features'  # Adjust based on your model architecture
                
                # Get activations for input image
                image_activation = self._get_layer_activation(image, layer_name)
                
                # Get activations for concept examples
                concept_activations_list = []
                for concept_img in concept_examples:
                    concept_activation = self._get_layer_activation(concept_img, layer_name)
                    concept_activations_list.append(concept_activation)
                
                # Calculate concept activation strength (simplified)
                concept_mean = torch.mean(torch.stack(concept_activations_list), dim=0)
                activation_strength = torch.cosine_similarity(
                    image_activation.flatten(), 
                    concept_mean.flatten(), 
                    dim=0
                ).item()
                
                concept_activations[concept_name] = ConceptActivation(
                    concept_name=concept_name,
                    activation_strength=activation_strength,
                    statistical_significance=abs(activation_strength),  # Simplified
                    directional_derivative=activation_strength
                )
                
            except Exception as e:
                logger.error(f"Failed to analyze concept {concept_name}: {e}")
        
        return concept_activations
    
    def generate_counterfactuals(
        self, 
        image: torch.Tensor, 
        target_change: float = 0.5
    ) -> List[Tuple[np.ndarray, float, List[Tuple[int, int, int, int]]]]:
        """
        Generate counterfactual explanations by finding minimal changes
        
        Args:
            image: Input image tensor
            target_change: Desired change in prediction confidence
            
        Returns:
            List of (modified_image, new_confidence, changed_regions)
        """
        counterfactuals = []
        
        # Get original prediction
        with torch.no_grad():
            original_pred = self.model(image).sigmoid().item()
        
        target_pred = max(0.0, min(1.0, original_pred + target_change))
        
        # Use gradient-based optimization to find minimal changes
        modified_image = image.clone().requires_grad_(True)
        optimizer = torch.optim.Adam([modified_image], lr=0.01)
        
        for iteration in range(100):
            optimizer.zero_grad()
            
            prediction = self.model(modified_image).sigmoid()
            
            # Loss: difference from target + L2 regularization
            loss = (prediction - target_pred).pow(2) + 0.1 * (modified_image - image).pow(2).sum()
            
            loss.backward()
            optimizer.step()
            
            # Clamp to valid image range
            with torch.no_grad():
                modified_image.clamp_(-1, 1)
            
            if iteration % 20 == 0:
                current_pred = prediction.item()
                if abs(current_pred - target_pred) < 0.05:
                    break
        
        # Find changed regions
        diff = torch.abs(modified_image - image).squeeze()
        diff_np = diff.cpu().numpy()
        
        # Threshold to find significant changes
        threshold = np.percentile(diff_np.flatten(), 95)
        changed_mask = diff_np > threshold
        
        # Find bounding boxes of changed regions
        changed_regions = self._find_changed_regions(changed_mask)
        
        final_pred = self.model(modified_image).sigmoid().item()
        modified_image_np = modified_image.squeeze().cpu().numpy().transpose(1, 2, 0)
        
        counterfactuals.append((modified_image_np, final_pred, changed_regions))
        
        return counterfactuals
    
    def _generate_baselines(self, image: torch.Tensor, strategy: str) -> List[torch.Tensor]:
        """Generate multiple baselines for robust attribution"""
        baselines = []
        
        if strategy == 'zero':
            baselines.append(torch.zeros_like(image))
        elif strategy == 'mean':
            baselines.append(torch.full_like(image, image.mean()))
        elif strategy == 'blur':
            # Gaussian blur baseline
            blurred = F.avg_pool2d(image, kernel_size=21, stride=1, padding=10)
            baselines.append(blurred)
        elif strategy == 'noise':
            baselines.append(torch.randn_like(image) * 0.1)
        elif strategy == 'all':
            # Use all baseline strategies
            baselines.extend(self._generate_baselines(image, 'zero'))
            baselines.extend(self._generate_baselines(image, 'mean'))
            baselines.extend(self._generate_baselines(image, 'blur'))
            baselines.extend(self._generate_baselines(image, 'noise'))
        
        return baselines
    
    def _process_attribution(self, attribution: torch.Tensor) -> np.ndarray:
        """Process attribution tensor to visualization-ready format"""
        attr = attribution.squeeze().cpu().numpy()
        
        # Handle multi-channel attributions
        if len(attr.shape) == 3:
            attr = np.mean(attr, axis=0)  # Average across channels
        
        # Normalize
        attr_max = np.max(np.abs(attr))
        if attr_max > 0:
            attr = attr / attr_max
        
        return attr
    
    def _calculate_interpretation_score(
        self, attribution: torch.Tensor, image: torch.Tensor
    ) -> float:
        """Calculate interpretation quality score"""
        # Coherence: spatial continuity of attributions
        attr_np = self._process_attribution(attribution)
        
        # Calculate gradient magnitude (smoothness)
        grad_x = np.gradient(attr_np, axis=1)
        grad_y = np.gradient(attr_np, axis=0)
        gradient_magnitude = np.sqrt(grad_x**2 + grad_y**2)
        coherence = 1.0 / (1.0 + np.mean(gradient_magnitude))
        
        # Sparsity: concentration of important regions
        attr_abs = np.abs(attr_np)
        sparsity = 1.0 - (np.count_nonzero(attr_abs > 0.1) / attr_abs.size)
        
        # Combined interpretation score
        interpretation_score = 0.6 * coherence + 0.4 * sparsity
        
        return float(interpretation_score)
    
    def _extract_feature_importance(self, attribution: torch.Tensor) -> Dict[str, float]:
        """Extract regional feature importance"""
        attr_np = self._process_attribution(attribution)
        h, w = attr_np.shape
        
        # Divide into regions and calculate importance
        regions = {
            'top_left': attr_np[:h//2, :w//2],
            'top_right': attr_np[:h//2, w//2:],
            'bottom_left': attr_np[h//2:, :w//2],
            'bottom_right': attr_np[h//2:, w//2:],
            'center': attr_np[h//4:3*h//4, w//4:3*w//4],
            'edges': np.concatenate([
                attr_np[0, :], attr_np[-1, :], 
                attr_np[:, 0], attr_np[:, -1]
            ])
        }
        
        feature_importance = {}
        for region_name, region_attr in regions.items():
            feature_importance[region_name] = float(np.mean(np.abs(region_attr)))
        
        return feature_importance
    
    def _generate_textual_explanation(
        self, 
        attribution_map: np.ndarray,
        feature_importance: Dict[str, float],
        confidence: float,
        method: str
    ) -> str:
        """Generate human-readable textual explanation"""
        
        # Find most important regions
        sorted_regions = sorted(
            feature_importance.items(), 
            key=lambda x: x[1], 
            reverse=True
        )
        
        top_regions = [region for region, _ in sorted_regions[:3]]
        
        confidence_text = "high" if confidence > 0.7 else "medium" if confidence > 0.3 else "low"
        
        prediction_text = "morphed" if confidence > 0.5 else "authentic"
        
        explanation = f"Using {method}, the model predicts this image is {prediction_text} " \
                     f"with {confidence_text} confidence ({confidence:.2f}). " \
                     f"The most influential regions are: {', '.join(top_regions)}. "
        
        # Add specific insights based on attribution patterns
        if feature_importance.get('edges', 0) > 0.3:
            explanation += "Edge artifacts suggest potential morphing. "
        
        if feature_importance.get('center', 0) > 0.5:
            explanation += "Central facial features show strong morphing indicators. "
        
        return explanation
    
    def _get_layer_activation(self, image: torch.Tensor, layer_name: str) -> torch.Tensor:
        """Extract activation from specific layer"""
        activations = {}
        
        def hook_fn(module, input, output):
            activations[layer_name] = output
        
        # Register hook (simplified - you'd need to adapt this to your model)
        handle = None
        for name, module in self.model.named_modules():
            if layer_name in name:
                handle = module.register_forward_hook(hook_fn)
                break
        
        if handle is None:
            raise ValueError(f"Layer {layer_name} not found in model")
        
        with torch.no_grad():
            _ = self.model(image)
        
        handle.remove()
        
        return activations.get(layer_name, torch.tensor([]))
    
    def _find_changed_regions(self, mask: np.ndarray) -> List[Tuple[int, int, int, int]]:
        """Find bounding boxes of changed regions"""
        # Convert to uint8 for OpenCV
        mask_uint8 = (mask * 255).astype(np.uint8)
        
        # Find contours
        contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        regions = []
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            if w > 10 and h > 10:  # Filter small regions
                regions.append((x, y, x + w, y + h))
        
        return regions
    
    def visualize_explanations(
        self, 
        image: torch.Tensor,
        explanations: Dict[str, ExplanationResult],
        save_path: Optional[str] = None
    ) -> np.ndarray:
        """Create comprehensive visualization of all explanations"""
        
        num_methods = len(explanations)
        fig, axes = plt.subplots(2, max(4, num_methods//2), figsize=(20, 10))
        axes = axes.flatten()
        
        # Original image
        img_np = image.squeeze().cpu().numpy().transpose(1, 2, 0)
        if img_np.min() < 0:
            img_np = (img_np + 1) / 2
        
        axes[0].imshow(img_np)
        axes[0].set_title("Original Image")
        axes[0].axis('off')
        
        # Plot each explanation
        for i, (method, explanation) in enumerate(explanations.items(), 1):
            if i < len(axes):
                im = axes[i].imshow(explanation.attribution_map, cmap='RdBu_r', vmin=-1, vmax=1)
                axes[i].set_title(f"{method}\nScore: {explanation.interpretation_score:.3f}")
                axes[i].axis('off')
                plt.colorbar(im, ax=axes[i], fraction=0.046, pad=0.04)
        
        # Hide unused subplots
        for i in range(len(explanations) + 1, len(axes)):
            axes[i].axis('off')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        # Convert to numpy array for return
        fig.canvas.draw()
        buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        buf = buf.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        
        plt.close(fig)
        return buf


class ConfidenceCalibrator:
    """
    Confidence calibration for improved model reliability
    """
    
    def __init__(self):
        self.temperature = 1.0
        self.calibration_data = []
        
    def calibrate(self, logits: torch.Tensor, labels: torch.Tensor) -> float:
        """
        Temperature scaling calibration
        
        Args:
            logits: Model logits
            labels: True labels
            
        Returns:
            Optimal temperature parameter
        """
        # Find optimal temperature using validation set
        temperature = nn.Parameter(torch.ones(1) * 1.5)
        optimizer = torch.optim.LBFGS([temperature], lr=0.01, max_iter=50)
        
        def eval_loss():
            optimizer.zero_grad()
            loss = F.cross_entropy(logits / temperature, labels)
            loss.backward()
            return loss
        
        optimizer.step(eval_loss)
        
        self.temperature = temperature.item()
        return self.temperature
    
    def apply_calibration(self, logits: torch.Tensor) -> torch.Tensor:
        """Apply temperature scaling to logits"""
        return logits / self.temperature


if __name__ == "__main__":
    # Example usage
    logging.basicConfig(level=logging.INFO)
    
    # This would be your actual model
    model = torch.nn.Sequential(
        torch.nn.Conv2d(3, 64, 3),
        torch.nn.ReLU(),
        torch.nn.AdaptiveAvgPool2d(1),
        torch.nn.Flatten(),
        torch.nn.Linear(64, 1)
    )
    
    xai_engine = AdvancedXAIEngine(model)
    
    # Example image
    image = torch.randn(1, 3, 224, 224)
    
    # Generate explanations
    explanations = xai_engine.explain_prediction(image)
    
    for method, explanation in explanations.items():
        print(f"{method}: {explanation.textual_explanation}")