"""
XAI Tools for Face Recognition Explainability
Implements multiple explainability methods: Grad-CAM, SHAP, LIME, Integrated Gradients
"""

import torch
import torch.nn.functional as F
import numpy as np
from typing import Dict, Tuple, Optional, List
from pathlib import Path
import matplotlib.pyplot as plt
from PIL import Image
import shap
from lime import lime_image
from captum.attr import IntegratedGradients, LayerGradCam
from src.forensic_logger import ForensicLogger


class XAIAnalyzer:
    """
    XAI tools for explaining face recognition model decisions
    Supports Grad-CAM, SHAP, LIME, and Integrated Gradients
    """
    
    def __init__(self, model, device='cuda' if torch.cuda.is_available() else 'cpu', forensic_logger=None):
        """Initialize XAI analyzer"""
        self.model = model
        self.device = device
        self.logger = forensic_logger
        
        # find the right layer for Grad-CAM
        # this was tricky - needed to find last conv layer in InceptionResnetV1
        self.target_layer = self._find_target_layer()
        
        print("XAI Analyzer initialized")
        print(f"  Device: {device}")
        print(f"  Target layer for Grad-CAM: {self.target_layer}")
    
    def _find_target_layer(self):
        """Find last convolutional layer for Grad-CAM"""
        # for InceptionResnetV1, block8 works best
        target_name = None
        for name, module in self.model.model.named_modules():
            if 'block8' in name and 'conv2d' in name:
                target_name = name
        return target_name
    
    def gradcam(self, image, save_path=None):
        """
        Generate Grad-CAM heatmap
        Shows which regions of the face the model focuses on
        """
        if image.dim() == 3:
            image = image.unsqueeze(0)
        
        image = image.to(self.device)
        image.requires_grad = True
        
        # get the target layer
        target_module = dict(self.model.model.named_modules())[self.target_layer]
        
        # hooks to capture activations and gradients
        activations = []
        gradients = []
        
        def forward_hook(module, input, output):
            activations.append(output)
            output.retain_grad()
        
        def backward_hook(module, grad_input, grad_output):
            gradients.append(grad_output[0])
        
        # register hooks
        forward_handle = target_module.register_forward_hook(forward_hook)
        backward_handle = target_module.register_full_backward_hook(backward_hook)
        
        # forward pass
        self.model.model.zero_grad()
        embedding = self.model.model(image)
        target = embedding.norm()  # use embedding norm as target
        target.backward()
        
        # remove hooks
        forward_handle.remove()
        backward_handle.remove()
        
        # compute Grad-CAM
        grads = gradients[0]
        acts = activations[0]
        
        # global average pooling of gradients
        weights = torch.mean(grads, dim=(2, 3), keepdim=True)
        cam = torch.sum(weights * acts, dim=1, keepdim=True)
        cam = F.relu(cam)  # only positive influence
        
        # upsample to input size
        cam = F.interpolate(cam, size=(160, 160), mode='bilinear', align_corners=False)
        
        # normalize to [0, 1]
        cam = cam.squeeze().cpu().detach().numpy()
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
        
        # log if we have a logger
        if self.logger:
            import hashlib
            img_hash = hashlib.md5(image.cpu().detach().numpy().tobytes()).hexdigest()
            self.logger.log_explanation(
                method="GradCAM",
                model_name=self.model.config['model_name'],
                input_hash=img_hash,
                explanation_file=save_path or "not_saved"
            )
        
        if save_path:
            self._save_gradcam_visualization(image.squeeze().cpu().detach(), cam, save_path)
        
        return cam
    
    def _save_gradcam_visualization(self, image, heatmap, save_path):
        """Save Grad-CAM visualization with overlay"""
        # denormalize image back to [0, 1]
        img = image.permute(1, 2, 0).numpy()
        img = (img * 0.5) + 0.5
        img = np.clip(img, 0, 1)
        
        # create figure with 3 panels
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        axes[0].imshow(img)
        axes[0].set_title('Original Image')
        axes[0].axis('off')
        
        axes[1].imshow(heatmap, cmap='jet')
        axes[1].set_title('Grad-CAM Heatmap')
        axes[1].axis('off')
        
        axes[2].imshow(img)
        axes[2].imshow(heatmap, cmap='jet', alpha=0.5)
        axes[2].set_title('Overlay')
        axes[2].axis('off')
        
        plt.tight_layout()
        
        output_path = Path(save_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()
    
    def integrated_gradients(self, image, baseline=None, n_steps=50):
        """
        Compute Integrated Gradients attribution
        Path-based method that integrates gradients along path from baseline to input
        """
        if image.dim() == 3:
            image = image.unsqueeze(0)
        
        image = image.to(self.device)
        
        # use black image as baseline if not provided
        if baseline is None:
            baseline = torch.zeros_like(image)
        else:
            baseline = baseline.to(self.device)
        
        # wrapper to return scalar for captum
        def model_forward(x):
            embedding = self.model.model(x)
            return embedding.norm(dim=1, keepdim=True)
        
        ig = IntegratedGradients(model_forward)
        
        # compute attributions
        attributions = ig.attribute(image, baseline, n_steps=n_steps, return_convergence_delta=False)
        
        # aggregate across channels
        attr = attributions.squeeze().cpu().numpy()
        attr = np.abs(attr).sum(axis=0)
        
        # normalize
        attr = (attr - attr.min()) / (attr.max() - attr.min() + 1e-8)
        
        if self.logger:
            import hashlib
            img_hash = hashlib.md5(image.cpu().numpy().tobytes()).hexdigest()
            self.logger.log_explanation(
                method="IntegratedGradients",
                model_name=self.model.config['model_name'],
                input_hash=img_hash,
                explanation_file="not_saved"
            )
        
        return attr
    
    def shap_analysis(self, image, background_samples, save_path=None):
        """
        SHAP analysis using GradientExplainer
        Requires background samples to estimate feature importance
        """
        if image.dim() == 3:
            image = image.unsqueeze(0)
        
        image = image.to(self.device)
        background_samples = background_samples.to(self.device)
        
        # wrapper model that returns scalar
        # SHAP needs this for gradient computation
        class ModelWrapper(torch.nn.Module):
            def __init__(self, model):
                super().__init__()
                self.model = model
            
            def forward(self, x):
                embedding = self.model(x)
                return embedding.norm(dim=1, keepdim=True)
        
        wrapped_model = ModelWrapper(self.model.model)
        
        # GradientExplainer works better than DeepExplainer for FaceNet
        explainer = shap.GradientExplainer(wrapped_model, background_samples)
        
        # compute SHAP values
        shap_values = explainer.shap_values(image)
        
        if isinstance(shap_values, list):
            shap_values = shap_values[0]
        
        # aggregate across channels and normalize
        shap_map = np.abs(shap_values).squeeze().sum(axis=0)
        shap_map = (shap_map - shap_map.min()) / (shap_map.max() - shap_map.min() + 1e-8)
        
        if self.logger:
            import hashlib
            img_hash = hashlib.md5(image.cpu().numpy().tobytes()).hexdigest()
            self.logger.log_explanation(
                method="SHAP",
                model_name=self.model.config['model_name'],
                input_hash=img_hash,
                explanation_file=save_path or "not_saved"
            )
        
        if save_path:
            self._save_heatmap_visualization(image.squeeze().cpu(), shap_map, save_path, title="SHAP Attribution Map")
        
        return shap_map
    
    def lime_analysis(self, image, num_samples=1000, save_path=None):
        """
        LIME analysis
        Perturbs the image and sees how predictions change
        """
        if image.dim() == 3:
            image = image.unsqueeze(0)
        
        # denormalize for LIME (it expects [0, 1] images)
        img_denorm = (image * 0.5) + 0.5
        img_np = img_denorm.squeeze().permute(1, 2, 0).cpu().numpy()
        img_np = np.clip(img_np, 0, 1)
        
        def predict_fn(images):
            # LIME passes denormalized images, need to renormalize
            images_normalized = (images - 0.5) / 0.5
            images_tensor = torch.FloatTensor(images_normalized).permute(0, 3, 1, 2).to(self.device)
            
            with torch.no_grad():
                embeddings = self.model.model(images_tensor)
            
            # LIME expects class probabilities
            # we fake it by using embedding norm
            norms = embeddings.norm(dim=1, keepdim=True)
            probs = torch.cat([1 - norms, norms], dim=1).cpu().numpy()
            return probs
        
        explainer = lime_image.LimeImageExplainer()
        
        explanation = explainer.explain_instance(
            img_np,
            predict_fn,
            top_labels=1,
            hide_color=0,
            num_samples=num_samples
        )
        
        explained_label = explanation.top_labels[0]
        
        temp, mask = explanation.get_image_and_mask(
            explained_label,
            positive_only=False,
            hide_rest=False,
            num_features=10
        )
        
        if self.logger:
            import hashlib
            img_hash = hashlib.md5(image.cpu().numpy().tobytes()).hexdigest()
            self.logger.log_explanation(
                method="LIME",
                model_name=self.model.config['model_name'],
                input_hash=img_hash,
                explanation_file=save_path or "not_saved"
            )
        
        if save_path:
            self._save_lime_visualization(img_np, temp, mask, save_path)
        
        return mask.astype(np.float32)
    
    def _save_heatmap_visualization(self, image, heatmap, save_path, title="Attribution Map"):
        """Save generic heatmap visualization"""
        img = image.permute(1, 2, 0).numpy()
        img = (img * 0.5) + 0.5
        img = np.clip(img, 0, 1)
        
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        axes[0].imshow(img)
        axes[0].set_title('Original Image')
        axes[0].axis('off')
        
        axes[1].imshow(heatmap, cmap='hot')
        axes[1].set_title(title)
        axes[1].axis('off')
        
        axes[2].imshow(img)
        axes[2].imshow(heatmap, cmap='hot', alpha=0.5)
        axes[2].set_title('Overlay')
        axes[2].axis('off')
        
        plt.tight_layout()
        
        output_path = Path(save_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()
    
    def _save_lime_visualization(self, original_img, temp_img, mask, save_path):
        """Save LIME-specific visualization"""
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        axes[0].imshow(original_img)
        axes[0].set_title('Original Image')
        axes[0].axis('off')
        
        axes[1].imshow(mask, cmap='RdBu', vmin=-1, vmax=1)
        axes[1].set_title('LIME Importance Map')
        axes[1].axis('off')
        
        axes[2].imshow(temp_img)
        axes[2].set_title('LIME Explanation')
        axes[2].axis('off')
        
        plt.tight_layout()
        
        output_path = Path(save_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()


# quick test
if __name__ == "__main__":
    print("Testing XAI Tools...")
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")
    
    from src.forensic_logger import ForensicLogger
    logger = ForensicLogger(experiment_id="TEST_XAI")
    
    from src.model_utils import ForensicFaceNet
    model = ForensicFaceNet(pretrained='vggface2', device=device, forensic_logger=logger)
    
    xai = XAIAnalyzer(model=model, device=device, forensic_logger=logger)
    
    test_image = torch.randn(3, 160, 160)
    
    print("\nTesting Grad-CAM...")
    cam = xai.gradcam(test_image, save_path="results/figures/test_gradcam.png")
    print(f"  Grad-CAM shape: {cam.shape}")
    
    print("\nTesting Integrated Gradients...")
    ig_attr = xai.integrated_gradients(test_image, n_steps=25)
    print(f"  IG attribution shape: {ig_attr.shape}")
    
    print("\nSkipping SHAP and LIME (need background samples)")
    
    logger.verify_chain_integrity()
    print("\nTest completed successfully")