src/models/neural_face_explainer.py

"""
Neural Face Explainer

This module provides explainable AI components for face morph detection,
visualizing which facial features contribute most to the detection decision.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import cv2
from scipy.ndimage import gaussian_filter
import os
import json

class GradCAMExplainer:
    """
    Gradient-weighted Class Activation Mapping for visualizing important
    regions in face morph detection
    """
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        # Register hooks
        self._register_hooks()
        
    def _register_hooks(self):
        def forward_hook(module, input, output):
            self.activations = output
            
        def backward_hook(module, grad_input, grad_output):
            self.gradients = grad_output[0]
        
        # Register hooks
        target_module = self._find_target_layer(self.model, self.target_layer)
        if target_module is None:
            raise ValueError(f"Target layer {self.target_layer} not found in model")
            
        self.forward_handle = target_module.register_forward_hook(forward_hook)
        self.backward_handle = target_module.register_full_backward_hook(backward_hook)
        
    def _find_target_layer(self, model, target_layer_name):
        """Find the target layer by name in the model"""
        # First try to get it directly
        if hasattr(model, target_layer_name):
            return getattr(model, target_layer_name)
            
        # Search recursively
        for name, module in model.named_modules():
            if name == target_layer_name:
                return module
                
        return None
        
    def remove_hooks(self):
        """Remove registered hooks"""
        self.forward_handle.remove()
        self.backward_handle.remove()
        
    def generate_cam(self, input_img, input_points=None, class_idx=None):
        """
        Generate a class activation map for the input
        
        Args:
            input_img: Input image tensor [1, C, H, W]
            input_points: 3D face points tensor [1, N, 3] (optional)
            class_idx: Target class index (default is None, which uses the predicted class)
            
        Returns:
            cam: Class activation map
            pred_class: Predicted class
            pred_conf: Prediction confidence
        """
        # Set model to eval mode
        self.model.eval()
        
        # Forward pass
        if input_points is not None:
            output = self.model(input_img, input_points)
        else:
            output = self.model(input_img)
            
        # If output is a dict (during training), get the fused_logits
        if isinstance(output, dict):
            logits = output['fused_logits']
        else:
            logits = output
            
        # Get predicted class if class_idx is None
        if class_idx is None:
            pred_class = logits.argmax(dim=1).item()
            class_idx = pred_class
        else:
            pred_class = class_idx
            
        # Get prediction confidence
        pred_conf = F.softmax(logits, dim=1)[0, class_idx].item()
        
        # Zero gradients
        self.model.zero_grad()
        
        # Target for backprop
        one_hot = torch.zeros_like(logits)
        one_hot[0, class_idx] = 1
        
        # Backward pass
        logits.backward(gradient=one_hot, retain_graph=True)
        
        # Get gradients and activations
        gradients = self.gradients.detach().cpu()
        activations = self.activations.detach().cpu()
        
        # Global average pooling of gradients
        weights = torch.mean(gradients, dim=[2, 3])
        
        # Weight activations by gradients
        cam = torch.zeros(activations.shape[2:], dtype=torch.float32)
        
        for i, w in enumerate(weights[0]):
            cam += w * activations[0, i, :, :]
            
        cam = F.relu(cam)  # Apply ReLU to focus on features that have positive influence
        
        # Normalize
        if cam.max() > 0:
            cam = cam / cam.max()
            
        return cam.numpy(), pred_class, pred_conf
        
    def overlay_cam(self, img, cam, alpha=0.5, cmap='jet'):
        """
        Overlay the CAM on the input image
        
        Args:
            img: Input image (numpy array, H x W x C)
            cam: Class activation map (numpy array, H x W)
            alpha: Overlay opacity
            cmap: Colormap for heatmap
            
        Returns:
            overlaid_img: Image with CAM overlay
        """
        # Resize CAM to match image dimensions
        cam_resized = cv2.resize(cam, (img.shape[1], img.shape[0]))
        
        # Apply colormap
        cmap = cm.get_cmap(cmap)
        cam_colored = cmap(cam_resized)[:, :, :3]  # Remove alpha channel
        
        # Convert to uint8
        cam_colored = (cam_colored * 255).astype(np.uint8)
        
        # Blend image and heatmap
        overlaid_img = cv2.addWeighted(
            img, 1 - alpha, 
            cam_colored, alpha, 
            0
        )
        
        return overlaid_img

class Face3DExplainer:
    """
    Visualize which 3D facial regions contribute most to morph detection
    """
    def __init__(self, model, target_3d_layer):
        self.model = model
        self.target_layer = target_3d_layer
        self.point_importances = None
        
        # Register hooks
        self._register_hooks()
        
    def _register_hooks(self):
        def forward_hook(module, input, output):
            self.activations = output
            
        def backward_hook(module, grad_input, grad_output):
            self.gradients = grad_output[0]
        
        # Register hooks
        target_module = self._find_target_layer(self.model, self.target_layer)
        if target_module is None:
            raise ValueError(f"Target layer {self.target_layer} not found in model")
            
        self.forward_handle = target_module.register_forward_hook(forward_hook)
        self.backward_handle = target_module.register_full_backward_hook(backward_hook)
        
    def _find_target_layer(self, model, target_layer_name):
        """Find the target layer by name in the model"""
        # First try to get it directly
        if hasattr(model, target_layer_name):
            return getattr(model, target_layer_name)
            
        # Search recursively
        for name, module in model.named_modules():
            if name == target_layer_name:
                return module
                
        return None
        
    def remove_hooks(self):
        """Remove registered hooks"""
        self.forward_handle.remove()
        self.backward_handle.remove()
        
    def generate_point_importances(self, input_img, input_points, class_idx=None):
        """
        Generate importance values for each 3D point
        
        Args:
            input_img: Input image tensor [1, C, H, W]
            input_points: 3D face points tensor [1, N, 3]
            class_idx: Target class index (default is None, which uses the predicted class)
            
        Returns:
            point_importances: Importance values for each point [N]
            points: 3D points [N, 3]
            pred_class: Predicted class
            pred_conf: Prediction confidence
        """
        # Set model to eval mode
        self.model.eval()
        
        # Forward pass
        output = self.model(input_img, input_points)
            
        # If output is a dict (during training), get the fused_logits
        if isinstance(output, dict):
            logits = output['fused_logits']
        else:
            logits = output
            
        # Get predicted class if class_idx is None
        if class_idx is None:
            pred_class = logits.argmax(dim=1).item()
            class_idx = pred_class
        else:
            pred_class = class_idx
            
        # Get prediction confidence
        pred_conf = F.softmax(logits, dim=1)[0, class_idx].item()
        
        # Zero gradients
        self.model.zero_grad()
        
        # Target for backprop
        one_hot = torch.zeros_like(logits)
        one_hot[0, class_idx] = 1
        
        # Backward pass
        logits.backward(gradient=one_hot, retain_graph=True)
        
        # Get gradients and activations for the 3D points
        gradients = self.gradients.detach().cpu().numpy()[0]  # [N, D]
        activations = self.activations.detach().cpu().numpy()[0]  # [N, D]
        
        # Calculate importance as the dot product of gradients and activations
        importances = np.sum(gradients * activations, axis=1)  # [N]
        
        # Normalize to 0-1 range
        importances = (importances - importances.min()) / (importances.max() - importances.min() + 1e-8)
        
        return importances, input_points[0].detach().cpu().numpy(), pred_class, pred_conf

    def visualize_3d_heatmap(self, points, importances, output_path=None, threshold=0.5):
        """
        Visualize 3D points colored by their importance
        
        Args:
            points: 3D points numpy array [N, 3]
            importances: Importance values [N]
            output_path: Path to save the visualization (optional)
            threshold: Importance threshold for highlighting
            
        Returns:
            fig: Matplotlib figure
        """
        # Create figure
        fig = plt.figure(figsize=(10, 10))
        ax = fig.add_subplot(111, projection='3d')
        
        # Color map based on importance
        colors = plt.cm.jet(importances)
        
        # Scatter plot
        scatter = ax.scatter(
            points[:, 0], points[:, 1], points[:, 2],
            c=importances,
            cmap='jet',
            s=50 * importances + 5,  # Size based on importance
            alpha=0.7
        )
        
        # Highlight high-importance points
        high_importance_idx = importances > threshold
        if np.any(high_importance_idx):
            ax.scatter(
                points[high_importance_idx, 0],
                points[high_importance_idx, 1],
                points[high_importance_idx, 2],
                color='red',
                s=100,
                edgecolors='white',
                linewidths=0.5
            )
        
        # Set equal aspect ratio
        ax.set_box_aspect([1, 1, 1])
        
        # Add colorbar
        cbar = plt.colorbar(scatter, ax=ax, shrink=0.7)
        cbar.set_label('Feature Importance')
        
        # Set labels
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_title('3D Face Morph Detection - Feature Importance')
        
        # Save if output path provided
        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            
        return fig
        
    def export_3d_visualization(self, points, importances, output_path, format='json'):
        """
        Export 3D visualization data in various formats
        
        Args:
            points: 3D points numpy array [N, 3]
            importances: Importance values [N]
            output_path: Path to save the data
            format: Output format ('json', 'ply', 'obj')
        """
        if format == 'json':
            # Export as JSON for web-based visualization
            data = {
                'points': points.tolist(),
                'importances': importances.tolist()
            }
            
            with open(output_path, 'w') as f:
                json.dump(data, f)
                
        elif format == 'ply':
            # Export as PLY file with vertex colors
            with open(output_path, 'w') as f:
                # Header
                f.write("ply\n")
                f.write("format ascii 1.0\n")
                f.write(f"element vertex {len(points)}\n")
                f.write("property float x\n")
                f.write("property float y\n")
                f.write("property float z\n")
                f.write("property uchar red\n")
                f.write("property uchar green\n")
                f.write("property uchar blue\n")
                f.write("end_header\n")
                
                # Convert importances to RGB colors using jet colormap
                colors = (plt.cm.jet(importances)[:, :3] * 255).astype(int)
                
                # Vertex data
                for i in range(len(points)):
                    f.write(f"{points[i, 0]} {points[i, 1]} {points[i, 2]} "
                           f"{colors[i, 0]} {colors[i, 1]} {colors[i, 2]}\n")
                           
        elif format == 'obj':
            # Export as OBJ file with vertex colors via MTL
            # This is a simplification as OBJ doesn't directly support vertex colors
            with open(output_path, 'w') as f:
                f.write("# Morph detection importance visualization\n")
                
                # Write vertices
                for i, p in enumerate(points):
                    f.write(f"v {p[0]} {p[1]} {p[2]}\n")
                    
                # Add some basic mesh structure (just points as vertices)
                f.write("# Point cloud only, no faces\n")
                
            # Create an MTL file with the same base name
            mtl_path = output_path.replace('.obj', '.mtl')
            with open(mtl_path, 'w') as f:
                f.write("# Morph detection importance materials\n")
                
                # Create a material for high importance regions
                f.write("newmtl high_importance\n")
                f.write("Ka 1.0 0.0 0.0\n")  # Ambient color (red)
                f.write("Kd 1.0 0.0 0.0\n")  # Diffuse color (red)
                f.write("Ks 1.0 1.0 1.0\n")  # Specular color (white)
                f.write("Ns 100.0\n")        # Specular exponent
                
                # Create a material for medium importance regions
                f.write("newmtl medium_importance\n")
                f.write("Ka 1.0 1.0 0.0\n")  # Ambient color (yellow)
                f.write("Kd 1.0 1.0 0.0\n")  # Diffuse color (yellow)
                f.write("Ks 1.0 1.0 1.0\n")  # Specular color (white)
                f.write("Ns 50.0\n")         # Specular exponent
                
                # Create a material for low importance regions
                f.write("newmtl low_importance\n")
                f.write("Ka 0.0 0.0 1.0\n")  # Ambient color (blue)
                f.write("Kd 0.0 0.0 1.0\n")  # Diffuse color (blue)
                f.write("Ks 1.0 1.0 1.0\n")  # Specular color (white)
                f.write("Ns 10.0\n")         # Specular exponent
        else:
            raise ValueError(f"Unsupported export format: {format}")

class FeatureAttributionExplainer:
    """
    Generate explanations for morph detection decisions based on facial feature attribution
    """
    def __init__(self, model, feature_names=None):
        self.model = model
        
        # Default facial feature names if none provided
        self.feature_names = feature_names or [
            'Eyes', 'Eyebrows', 'Nose', 'Mouth', 'Chin', 
            'Jawline', 'Cheeks', 'Forehead', 'Ears'
        ]
        
        # Feature region maps (normalized coordinates)
        self.feature_regions = {
            'Eyes': [(0.2, 0.3, 0.4, 0.15), (0.6, 0.3, 0.4, 0.15)],  # Left and right eye
            'Eyebrows': [(0.2, 0.25, 0.4, 0.1), (0.6, 0.25, 0.4, 0.1)],
            'Nose': [(0.4, 0.4, 0.2, 0.25)],
            'Mouth': [(0.35, 0.7, 0.3, 0.15)],
            'Chin': [(0.4, 0.85, 0.2, 0.15)],
            'Jawline': [(0.1, 0.65, 0.8, 0.3)],
            'Cheeks': [(0.15, 0.5, 0.25, 0.2), (0.6, 0.5, 0.25, 0.2)],
            'Forehead': [(0.3, 0.15, 0.4, 0.15)],
            'Ears': [(0.05, 0.4, 0.15, 0.25), (0.8, 0.4, 0.15, 0.25)]
        }
        
    def _get_region_mask(self, feature_name, img_shape):
        """Generate a binary mask for a facial feature region"""
        h, w = img_shape
        mask = np.zeros((h, w), dtype=np.float32)
        
        for region in self.feature_regions[feature_name]:
            # Region is (center_x, center_y, width, height) in normalized coordinates
            cx, cy, rw, rh = region
            
            # Convert to pixel coordinates
            x1 = max(0, int((cx - rw/2) * w))
            y1 = max(0, int((cy - rh/2) * h))
            x2 = min(w, int((cx + rw/2) * w))
            y2 = min(h, int((cy + rh/2) * h))
            
            # Create soft mask with Gaussian falloff
            y, x = np.ogrid[0:h, 0:w]
            center_x = (x1 + x2) / 2
            center_y = (y1 + y2) / 2
            
            # Create distance map from center
            dist_map = np.sqrt((x - center_x)**2 + (y - center_y)**2)
            
            # Convert distance to a soft mask
            max_dist = np.sqrt((x2 - x1)**2 + (y2 - y1)**2) / 2
            soft_mask = np.maximum(0, 1 - dist_map / max_dist)
            
            # Add to the main mask
            mask = np.maximum(mask, soft_mask)
            
        return mask
    
    def analyze_feature_importance(self, input_img, input_points=None, class_idx=None):
        """
        Analyze which facial features contribute most to the morph detection decision
        
        Args:
            input_img: Input image tensor [1, C, H, W]
            input_points: 3D face points tensor [1, N, 3] (optional)
            class_idx: Target class index (default is None, which uses the predicted class)
            
        Returns:
            feature_importances: Dictionary mapping feature names to importance scores
            pred_class: Predicted class
            pred_conf: Prediction confidence
        """
        # Set model to eval mode
        self.model.eval()
        
        # Get original prediction
        with torch.no_grad():
            if input_points is not None:
                output = self.model(input_img, input_points)
            else:
                output = self.model(input_img)
                
            # If output is a dict (during training), get the fused_logits
            if isinstance(output, dict):
                logits = output['fused_logits']
            else:
                logits = output
                
            # Get predicted class if class_idx is None
            if class_idx is None:
                pred_class = logits.argmax(dim=1).item()
                class_idx = pred_class
            else:
                pred_class = class_idx
                
            # Get prediction confidence
            orig_conf = F.softmax(logits, dim=1)[0, class_idx].item()
        
        # Convert image to numpy for masking
        img_np = input_img[0].detach().cpu().numpy()
        c, h, w = img_np.shape
        
        # Dictionary to store feature importances
        feature_importances = {}
        
        # For each facial feature, mask it and observe change in prediction
        for feature in self.feature_names:
            # Create mask for this feature
            mask = self._get_region_mask(feature, (h, w))
            
            # Apply mask to each channel
            masked_img = img_np.copy()
            for i in range(c):
                # Invert the mask to keep everything except this feature
                inv_mask = 1 - mask
                
                # Calculate mean pixel value for this feature region
                feature_mean = np.sum(img_np[i] * mask) / (np.sum(mask) + 1e-8)
                
                # Replace feature with its mean value
                masked_img[i] = img_np[i] * inv_mask + feature_mean * mask
            
            # Convert back to tensor
            masked_tensor = torch.tensor(masked_img).unsqueeze(0).to(input_img.device)
            
            # Get prediction for masked image
            with torch.no_grad():
                if input_points is not None:
                    masked_output = self.model(masked_tensor, input_points)
                else:
                    masked_output = self.model(masked_tensor)
                    
                # Get logits
                if isinstance(masked_output, dict):
                    masked_logits = masked_output['fused_logits']
                else:
                    masked_logits = masked_output
                    
                # Get confidence for target class
                masked_conf = F.softmax(masked_logits, dim=1)[0, class_idx].item()
            
            # Calculate importance as the difference in confidence
            importance = orig_conf - masked_conf
            feature_importances[feature] = importance
            
        # Normalize importances to sum to 1
        total_importance = sum(abs(imp) for imp in feature_importances.values())
        if total_importance > 0:
            for feature in feature_importances:
                feature_importances[feature] /= total_importance
            
        return feature_importances, pred_class, orig_conf
    
    def visualize_feature_importances(self, feature_importances, output_path=None):
        """
        Visualize feature importances as a bar chart
        
        Args:
            feature_importances: Dictionary mapping feature names to importance scores
            output_path: Path to save the visualization (optional)
            
        Returns:
            fig: Matplotlib figure
        """
        # Sort features by importance
        sorted_features = sorted(
            feature_importances.items(),
            key=lambda x: abs(x[1]),
            reverse=True
        )
        
        # Extract feature names and values
        features = [f[0] for f in sorted_features]
        values = [f[1] for f in sorted_features]
        
        # Create figure
        fig, ax = plt.subplots(figsize=(10, 6))
        
        # Create colormap based on values
        colors = ['red' if v < 0 else 'green' for v in values]
        
        # Create bar chart
        bars = ax.barh(features, [abs(v) for v in values], color=colors)
        
        # Add value labels
        for i, bar in enumerate(bars):
            width = bar.get_width()
            label = f"{values[i]:.3f}"
            ax.text(
                width + 0.01,
                bar.get_y() + bar.get_height()/2,
                label,
                ha='left',
                va='center'
            )
        
        # Set labels and title
        ax.set_xlabel('Feature Importance')
        ax.set_title('Facial Feature Importance for Morph Detection')
        
        # Add legend
        ax.text(
            0.01, -0.15,
            "Green: Feature increases likelihood of being classified as morphed\n"
            "Red: Feature decreases likelihood of being classified as morphed",
            transform=ax.transAxes,
            ha='left'
        )
        
        # Adjust layout
        plt.tight_layout()
        
        # Save if output path provided
        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            
        return fig

    def generate_feature_heatmap(self, input_img, feature_importances, output_path=None):
        """
        Generate a heatmap highlighting important facial regions
        
        Args:
            input_img: Input image (numpy array, H x W x C or tensor)
            feature_importances: Dictionary mapping feature names to importance scores
            output_path: Path to save the visualization (optional)
            
        Returns:
            overlaid_img: Image with importance heatmap overlay
        """
        # Convert image to numpy if it's a tensor
        if torch.is_tensor(input_img):
            if input_img.dim() == 4:  # [1, C, H, W]
                img_np = input_img[0].detach().cpu().numpy().transpose(1, 2, 0)
            else:  # [C, H, W]
                img_np = input_img.detach().cpu().numpy().transpose(1, 2, 0)
        else:
            img_np = input_img
            
        # Ensure image is in range 0-255 and uint8
        if img_np.max() <= 1.0:
            img_np = (img_np * 255).astype(np.uint8)
        
        h, w = img_np.shape[:2]
        
        # Create importance heatmap
        heatmap = np.zeros((h, w), dtype=np.float32)
        
        # Add each feature's contribution to the heatmap
        for feature, importance in feature_importances.items():
            mask = self._get_region_mask(feature, (h, w))
            heatmap += mask * importance
        
        # Normalize heatmap to 0-1 range
        heatmap_min, heatmap_max = heatmap.min(), heatmap.max()
        if heatmap_max > heatmap_min:
            heatmap = (heatmap - heatmap_min) / (heatmap_max - heatmap_min)
        
        # Apply Gaussian blur to smooth the heatmap
        heatmap = gaussian_filter(heatmap, sigma=5)
        
        # Convert heatmap to colormap
        heatmap_colored = plt.cm.jet(heatmap)[:, :, :3]
        heatmap_colored = (heatmap_colored * 255).astype(np.uint8)
        
        # Overlay heatmap on image
        alpha = 0.5
        overlaid_img = cv2.addWeighted(
            img_np, 1 - alpha,
            heatmap_colored, alpha,
            0
        )
        
        # Save if output path provided
        if output_path:
            plt.imsave(output_path, overlaid_img)
            
        return overlaid_img

class ModelExplainer:
    """
    Comprehensive model explainer that combines multiple explanation techniques
    """
    def __init__(self, model, device='cuda'):
        self.model = model
        self.device = device
        
        # Initialize explainers
        self.grad_cam = GradCAMExplainer(model, 'image_encoder')
        self.face_3d = Face3DExplainer(model, 'transformer_3d')
        self.feature_attr = FeatureAttributionExplainer(model)
        
    def explain(self, img, points=None, output_dir=None, class_idx=None):
        """
        Generate comprehensive explanations for a detection decision
        
        Args:
            img: Input image tensor [1, C, H, W]
            points: 3D face points tensor [1, N, 3] (optional)
            output_dir: Directory to save visualizations (optional)
            class_idx: Target class index (default is None, which uses the predicted class)
            
        Returns:
            explanation: Dictionary containing all explanation elements
        """
        # Create output directory if provided
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
            
        # Move inputs to correct device
        img = img.to(self.device)
        if points is not None:
            points = points.to(self.device)
            
        # Run model to get base prediction
        self.model.eval()
        with torch.no_grad():
            output = self.model(img, points) if points is not None else self.model(img)
            
            # If output is a dict (during training), get the fused_logits
            if isinstance(output, dict):
                logits = output['fused_logits']
            else:
                logits = output
                
            # Get predicted class if class_idx is None
            if class_idx is None:
                pred_class = logits.argmax(dim=1).item()
                class_idx = pred_class
            else:
                pred_class = class_idx
                
            # Get prediction confidence
            pred_conf = F.softmax(logits, dim=1)[0, class_idx].item()
            
        # Get class name
        class_name = "Morphed" if pred_class == 1 else "Real"
            
        # Initialize explanation dictionary
        explanation = {
            'prediction': {
                'class_idx': pred_class,
                'class_name': class_name,
                'confidence': pred_conf
            },
            'explanations': {}
        }
        
        # 1. Generate 2D heatmap using Grad-CAM
        cam, _, _ = self.grad_cam.generate_cam(img, points, class_idx)
        
        # Convert image to numpy for visualization
        if img.dim() == 4:  # [1, C, H, W]
            img_np = img[0].detach().cpu().numpy().transpose(1, 2, 0)
        else:  # [C, H, W]
            img_np = img.detach().cpu().numpy().transpose(1, 2, 0)
            
        # Ensure image is in range 0-255 and uint8
        if img_np.max() <= 1.0:
            img_np = (img_np * 255).astype(np.uint8)
            
        # Resize CAM to match image dimensions
        cam_resized = cv2.resize(cam, (img_np.shape[1], img_np.shape[0]))
        
        # Overlay CAM on image
        overlaid_cam = self.grad_cam.overlay_cam(img_np, cam_resized)
        
        explanation['explanations']['2d_heatmap'] = {
            'cam': cam_resized.tolist(),
            'visualization': overlaid_cam.tolist() if output_dir else None
        }
        
        # Save 2D visualization if output directory provided
        if output_dir:
            cam_path = os.path.join(output_dir, '2d_heatmap.png')
            cv2.imwrite(cam_path, cv2.cvtColor(overlaid_cam, cv2.COLOR_RGB2BGR))
            explanation['explanations']['2d_heatmap']['file_path'] = cam_path
            
        # 2. Generate 3D point importances if points provided
        if points is not None:
            point_imp, point_coords, _, _ = self.face_3d.generate_point_importances(
                img, points, class_idx
            )
            
            # Generate 3D visualization
            fig = self.face_3d.visualize_3d_heatmap(point_coords, point_imp)
            
            explanation['explanations']['3d_heatmap'] = {
                'point_importances': point_imp.tolist(),
                'point_coordinates': point_coords.tolist()
            }
            
            # Save 3D visualization if output directory provided
            if output_dir:
                vis_3d_path = os.path.join(output_dir, '3d_heatmap.png')
                plt.savefig(vis_3d_path, dpi=300, bbox_inches='tight')
                plt.close(fig)
                
                # Export 3D data for interactive visualization
                json_path = os.path.join(output_dir, '3d_data.json')
                self.face_3d.export_3d_visualization(point_coords, point_imp, json_path)
                
                explanation['explanations']['3d_heatmap']['file_paths'] = {
                    'image': vis_3d_path,
                    'json': json_path
                }
                
        # 3. Generate facial feature attributions
        feature_imp, _, _ = self.feature_attr.analyze_feature_importance(
            img, points, class_idx
        )
        
        # Generate feature importance visualization
        fig = self.feature_attr.visualize_feature_importances(feature_imp)
        
        # Generate feature heatmap
        feat_heatmap = self.feature_attr.generate_feature_heatmap(img_np, feature_imp)
        
        explanation['explanations']['feature_attribution'] = {
            'feature_importances': feature_imp,
            'visualization': feat_heatmap.tolist() if output_dir else None
        }
        
        # Save feature attribution visualizations if output directory provided
        if output_dir:
            feat_bar_path = os.path.join(output_dir, 'feature_importance.png')
            plt.savefig(feat_bar_path, dpi=300, bbox_inches='tight')
            plt.close(fig)
            
            feat_heatmap_path = os.path.join(output_dir, 'feature_heatmap.png')
            cv2.imwrite(feat_heatmap_path, cv2.cvtColor(feat_heatmap, cv2.COLOR_RGB2BGR))
            
            explanation['explanations']['feature_attribution']['file_paths'] = {
                'bar_chart': feat_bar_path,
                'heatmap': feat_heatmap_path
            }
            
        # 4. Generate explanation text
        explanation['explanation_text'] = self._generate_explanation_text(
            pred_class, pred_conf, feature_imp
        )
        
        return explanation
    
    def _generate_explanation_text(self, pred_class, confidence, feature_importances):
        """Generate human-readable explanation text"""
        class_name = "morphed" if pred_class == 1 else "real"
        
        # Sort features by importance
        sorted_features = sorted(
            feature_importances.items(),
            key=lambda x: abs(x[1]),
            reverse=True
        )
        
        # Get top 3 most important features
        top_features = sorted_features[:3]
        
        # Generate explanation
        text = f"The face is classified as {class_name} with {confidence:.1%} confidence. "
        
        if pred_class == 1:  # Morphed
            text += "The key indicators of morphing are: "
        else:  # Real
            text += "The key indicators of authenticity are: "
            
        for i, (feature, importance) in enumerate(top_features):
            if i > 0:
                text += ", " if i < len(top_features) - 1 else " and "
                
            if importance > 0:
                text += f"the {feature.lower()}" 
                if pred_class == 1:
                    text += " (which appears artificially manipulated)"
                else:
                    text += " (which appears natural)"
            else:
                text += f"the {feature.lower()}"
                if pred_class == 1:
                    text += " (which lacks natural variation)"
                else:
                    text += " (which shows natural characteristics)"
                    
        text += "."
        
        return text