"""Quantitative evaluation metrics for model interpretations.

This module provides metrics to evaluate the quality of explanation methods:
- Faithfulness: How well does the explanation reflect model behavior?
- Sensitivity: Does the explanation respond to meaningful changes?
- Localization: Does the explanation correctly identify important regions?
- Sanity checks: Does the method behave reasonably?
"""

from typing import Optional, Union, Dict, List
import numpy as np
import torch
import torch.nn as nn
from PIL import Image

from autotimm.core.logging import logger


class ExplanationMetrics:
    """
    Compute quantitative metrics for explanation quality.

    Args:
        model: The model being explained
        explainer: The interpretation method to evaluate
        use_cuda: Whether to use CUDA if available

    Example:
        >>> from autotimm import ImageClassifier
        >>> from autotimm.interpretation import GradCAM, ExplanationMetrics
        >>>
        >>> model = ImageClassifier(backbone="resnet50", num_classes=10)
        >>> explainer = GradCAM(model)
        >>> metrics = ExplanationMetrics(model, explainer)
        >>>
        >>> # Evaluate faithfulness
        >>> deletion_score = metrics.deletion(image, steps=50)
        >>> insertion_score = metrics.insertion(image, steps=50)
    """

    def __init__(
        self,
        model: nn.Module,
        explainer,
        use_cuda: bool = True,
    ):
        self.model = model
        self.model.eval()
        self.explainer = explainer
        self.device = torch.device(
            "cuda" if use_cuda and torch.cuda.is_available() else "cpu"
        )
        self.model.to(self.device)

    def deletion(
        self,
        image: Union[Image.Image, np.ndarray, torch.Tensor],
        target_class: Optional[int] = None,
        steps: int = 50,
        baseline: str = "blur",
    ) -> Dict[str, float]:
        """
        Deletion metric: progressively delete most important pixels.

        Measures how prediction confidence drops as most important pixels
        (according to the explanation) are removed.

        Args:
            image: Input image
            target_class: Target class (None = predicted class)
            steps: Number of deletion steps
            baseline: What to replace deleted pixels with ('blur', 'black', 'mean')

        Returns:
            Dictionary with:
                - auc: Area under the curve (lower = better explanation)
                - final_drop: Final prediction drop (higher = better)
                - scores: List of prediction scores at each step

        Example:
            >>> result = metrics.deletion(image, steps=50)
            >>> print(f"Deletion AUC: {result['auc']:.3f}")
            >>> print(f"Final drop: {result['final_drop']:.2%}")
        """
        # Get original prediction
        input_tensor = self._preprocess_image(image)
        _, original_score, target_class = self._get_prediction(
            input_tensor, target_class
        )

        # Get explanation
        heatmap = self.explainer.explain(image, target_class=target_class)

        # Resize heatmap to image size if needed
        if heatmap.shape != input_tensor.shape[2:]:
            heatmap = self._resize_heatmap(heatmap, input_tensor.shape[2:])

        # Create baseline
        baseline_tensor = self._create_baseline(input_tensor, baseline)

        # Get pixel importance ranking
        pixel_order = np.argsort(heatmap.flatten())[::-1]  # Most important first

        # Progressively delete pixels
        scores = [original_score]
        modified = input_tensor.clone()

        pixels_per_step = len(pixel_order) // steps

        for step in range(1, steps + 1):
            # Delete pixels
            pixels_to_delete = pixel_order[: step * pixels_per_step]

            # Convert flat indices to 2D
            rows = pixels_to_delete // heatmap.shape[1]
            cols = pixels_to_delete % heatmap.shape[1]

            # Replace with baseline
            for r, c in zip(rows, cols):
                modified[:, :, r, c] = baseline_tensor[:, :, r, c]

            # Get new prediction
            with torch.no_grad():
                output = self.model(modified)
                if isinstance(output, dict):
                    output = output.get(
                        "logits", output.get("output", list(output.values())[0])
                    )
                probs = torch.softmax(output, dim=1)
                score = probs[0, target_class].item()

            scores.append(score)

        # Compute metrics
        try:
            auc = np.trapezoid(scores, dx=1.0 / steps) / original_score
        except AttributeError:
            # Fallback for older numpy versions
            auc = np.trapz(scores, dx=1.0 / steps) / original_score
        final_drop = (original_score - scores[-1]) / original_score

        return {
            "auc": auc,
            "final_drop": final_drop,
            "scores": scores,
            "original_score": original_score,
        }

    def insertion(
        self,
        image: Union[Image.Image, np.ndarray, torch.Tensor],
        target_class: Optional[int] = None,
        steps: int = 50,
        baseline: str = "blur",
    ) -> Dict[str, float]:
        """
        Insertion metric: progressively insert most important pixels.

        Measures how prediction confidence rises as most important pixels
        (according to the explanation) are added to a baseline.

        Args:
            image: Input image
            target_class: Target class (None = predicted class)
            steps: Number of insertion steps
            baseline: Initial baseline ('blur', 'black', 'mean')

        Returns:
            Dictionary with:
                - auc: Area under the curve (higher = better explanation)
                - final_rise: Final prediction rise (higher = better)
                - scores: List of prediction scores at each step

        Example:
            >>> result = metrics.insertion(image, steps=50)
            >>> print(f"Insertion AUC: {result['auc']:.3f}")
            >>> print(f"Final rise: {result['final_rise']:.2%}")
        """
        # Get original prediction
        input_tensor = self._preprocess_image(image)
        _, original_score, target_class = self._get_prediction(
            input_tensor, target_class
        )

        # Get explanation
        heatmap = self.explainer.explain(image, target_class=target_class)

        # Resize heatmap to image size if needed
        if heatmap.shape != input_tensor.shape[2:]:
            heatmap = self._resize_heatmap(heatmap, input_tensor.shape[2:])

        # Create baseline
        baseline_tensor = self._create_baseline(input_tensor, baseline)

        # Get baseline prediction
        with torch.no_grad():
            output = self.model(baseline_tensor)
            if isinstance(output, dict):
                output = output.get(
                    "logits", output.get("output", list(output.values())[0])
                )
            probs = torch.softmax(output, dim=1)
            baseline_score = probs[0, target_class].item()

        # Get pixel importance ranking
        pixel_order = np.argsort(heatmap.flatten())[::-1]  # Most important first

        # Progressively insert pixels
        scores = [baseline_score]
        modified = baseline_tensor.clone()

        pixels_per_step = len(pixel_order) // steps

        for step in range(1, steps + 1):
            # Insert pixels
            pixels_to_insert = pixel_order[: step * pixels_per_step]

            # Convert flat indices to 2D
            rows = pixels_to_insert // heatmap.shape[1]
            cols = pixels_to_insert % heatmap.shape[1]

            # Replace with original pixels
            for r, c in zip(rows, cols):
                modified[:, :, r, c] = input_tensor[:, :, r, c]

            # Get new prediction
            with torch.no_grad():
                output = self.model(modified)
                if isinstance(output, dict):
                    output = output.get(
                        "logits", output.get("output", list(output.values())[0])
                    )
                probs = torch.softmax(output, dim=1)
                score = probs[0, target_class].item()

            scores.append(score)

        # Compute metrics
        try:
            auc = np.trapezoid(scores, dx=1.0 / steps) / original_score
        except AttributeError:
            # Fallback for older numpy versions
            auc = np.trapz(scores, dx=1.0 / steps) / original_score

        # Compute final rise - handle edge case where baseline >= original
        denominator = original_score - baseline_score
        if abs(denominator) < 1e-6:
            # Baseline and original are essentially the same
            final_rise = 0.0
        else:
            # Use absolute value to avoid negative denominators
            final_rise = (scores[-1] - baseline_score) / abs(denominator)

        return {
            "auc": auc,
            "final_rise": final_rise,
            "scores": scores,
            "original_score": original_score,
            "baseline_score": baseline_score,
        }

    def sensitivity_n(
        self,
        image: Union[Image.Image, np.ndarray, torch.Tensor],
        target_class: Optional[int] = None,
        n_samples: int = 50,
        noise_level: float = 0.15,
    ) -> Dict[str, float]:
        """
        Sensitivity-n metric: explanation stability under input perturbations.

        Measures how much the explanation changes when small noise is added
        to the input. Lower values indicate more stable explanations.

        Args:
            image: Input image
            target_class: Target class
            n_samples: Number of noisy samples
            noise_level: Standard deviation of Gaussian noise

        Returns:
            Dictionary with:
                - sensitivity: Average explanation change
                - std: Standard deviation of changes
                - max_change: Maximum change observed

        Example:
            >>> result = metrics.sensitivity_n(image, n_samples=30)
            >>> print(f"Sensitivity: {result['sensitivity']:.4f}")
        """
        input_tensor = self._preprocess_image(image)
        _, _, target_class = self._get_prediction(input_tensor, target_class)

        # Get original explanation
        original_heatmap = self.explainer.explain(image, target_class=target_class)

        # Generate noisy versions
        changes = []
        for _ in range(n_samples):
            # Add noise
            noise = torch.randn_like(input_tensor) * noise_level
            noisy_input = input_tensor + noise
            noisy_input = torch.clamp(noisy_input, 0, 1)

            # Get explanation for noisy input
            noisy_heatmap = self.explainer.explain(
                noisy_input, target_class=target_class
            )

            # Compute change
            change = np.abs(noisy_heatmap - original_heatmap).mean()
            changes.append(change)

        return {
            "sensitivity": np.mean(changes),
            "std": np.std(changes),
            "max_change": np.max(changes),
            "changes": changes,
        }

    def model_parameter_randomization_test(
        self,
        image: Union[Image.Image, np.ndarray, torch.Tensor],
        target_class: Optional[int] = None,
    ) -> Dict[str, float]:
        """
        Sanity check: explanation should change significantly if model is randomized.

        Tests whether the explanation method is sensitive to model parameters.
        A good explanation method should produce very different results when
        model weights are randomized.

        Args:
            image: Input image
            target_class: Target class

        Returns:
            Dictionary with:
                - correlation: Correlation between original and randomized explanations
                  (lower = better, method is sensitive to model)
                - change: Mean absolute difference
                - passes: True if correlation < 0.5 (reasonable threshold)

        Example:
            >>> result = metrics.model_parameter_randomization_test(image)
            >>> print(f"Passes sanity check: {result['passes']}")
            >>> print(f"Correlation: {result['correlation']:.3f}")
        """
        input_tensor = self._preprocess_image(image)
        _, _, target_class = self._get_prediction(input_tensor, target_class)

        # Get original explanation
        original_heatmap = self.explainer.explain(image, target_class=target_class)

        # Save original model state
        original_state = {
            name: param.clone() for name, param in self.model.named_parameters()
        }

        # Randomize model parameters
        with torch.no_grad():
            for param in self.model.parameters():
                if param.requires_grad:
                    param.data = torch.randn_like(param.data)

        # Get explanation with randomized model
        randomized_heatmap = self.explainer.explain(image, target_class=target_class)

        # Restore original parameters
        with torch.no_grad():
            for name, param in self.model.named_parameters():
                param.data = original_state[name].data

        # Compute correlation
        correlation = np.corrcoef(
            original_heatmap.flatten(), randomized_heatmap.flatten()
        )[0, 1]

        # Compute mean absolute change
        change = np.abs(randomized_heatmap - original_heatmap).mean()

        # Test passes if correlation is low (< 0.5)
        passes = correlation < 0.5

        return {
            "correlation": correlation,
            "change": change,
            "passes": passes,
        }

    def data_randomization_test(
        self,
        image: Union[Image.Image, np.ndarray, torch.Tensor],
        target_class: Optional[int] = None,
    ) -> Dict[str, float]:
        """
        Sanity check: explanation should change significantly if input is randomized.

        Tests whether the explanation is sensitive to input data.
        A good explanation method should produce very different results when
        labels are randomized (i.e., explaining wrong class).

        Args:
            image: Input image
            target_class: Target class

        Returns:
            Dictionary with:
                - correlation: Correlation with random label explanation
                - change: Mean absolute difference
                - passes: True if correlation < 0.5

        Example:
            >>> result = metrics.data_randomization_test(image)
            >>> print(f"Passes sanity check: {result['passes']}")
        """
        input_tensor = self._preprocess_image(image)
        _, _, target_class = self._get_prediction(input_tensor, target_class)

        # Get original explanation
        original_heatmap = self.explainer.explain(image, target_class=target_class)

        # Get explanation for random class (different from target)
        num_classes = self._get_num_classes()
        random_classes = [c for c in range(num_classes) if c != target_class]
        if random_classes:
            random_class = np.random.choice(random_classes)
            random_heatmap = self.explainer.explain(image, target_class=random_class)
        else:
            # If only one class, use same class (test will fail, which is correct)
            random_heatmap = original_heatmap

        # Compute correlation
        correlation = np.corrcoef(original_heatmap.flatten(), random_heatmap.flatten())[
            0, 1
        ]

        # Compute change
        change = np.abs(random_heatmap - original_heatmap).mean()

        # Test passes if correlation is low
        passes = correlation < 0.5

        return {
            "correlation": correlation,
            "change": change,
            "passes": passes,
        }

    def pointing_game(
        self,
        image: Union[Image.Image, np.ndarray, torch.Tensor],
        bbox: List[float],
        target_class: Optional[int] = None,
    ) -> Dict[str, bool]:
        """
        Pointing game: does the max attention fall within the object bbox?

        Used for object detection/localization. Tests whether the most
        important pixel (according to explanation) falls within the
        ground-truth bounding box.

        Args:
            image: Input image
            bbox: Ground truth bounding box [x1, y1, x2, y2]
            target_class: Target class

        Returns:
            Dictionary with:
                - hit: True if max attention is inside bbox
                - max_location: (row, col) of maximum attention
                - bbox: Input bbox

        Example:
            >>> bbox = [100, 150, 300, 400]  # Object location
            >>> result = metrics.pointing_game(image, bbox)
            >>> print(f"Hit: {result['hit']}")
        """
        # Get explanation
        heatmap = self.explainer.explain(image, target_class=target_class)

        # Find max location
        max_idx = np.argmax(heatmap)
        max_row = max_idx // heatmap.shape[1]
        max_col = max_idx % heatmap.shape[1]

        # Convert to original image coordinates
        input_tensor = self._preprocess_image(image)
        h_ratio = input_tensor.shape[2] / heatmap.shape[0]
        w_ratio = input_tensor.shape[3] / heatmap.shape[1]

        max_y = max_row * h_ratio
        max_x = max_col * w_ratio

        # Check if inside bbox
        x1, y1, x2, y2 = bbox
        hit = (x1 <= max_x <= x2) and (y1 <= max_y <= y2)

        return {
            "hit": hit,
            "max_location": (max_y, max_x),
            "bbox": bbox,
        }

    def evaluate_all(
        self,
        image: Union[Image.Image, np.ndarray, torch.Tensor],
        target_class: Optional[int] = None,
        bbox: Optional[List[float]] = None,
    ) -> Dict[str, Dict]:
        """
        Run all applicable metrics on an image.

        Args:
            image: Input image
            target_class: Target class
            bbox: Optional bounding box for pointing game

        Returns:
            Dictionary with all metric results

        Example:
            >>> results = metrics.evaluate_all(image)
            >>> print(f"Deletion AUC: {results['deletion']['auc']:.3f}")
            >>> print(f"Insertion AUC: {results['insertion']['auc']:.3f}")
            >>> print(f"Sensitivity: {results['sensitivity']['sensitivity']:.4f}")
        """
        results = {}

        # Faithfulness metrics
        logger.info("Computing deletion metric...")
        results["deletion"] = self.deletion(image, target_class=target_class)

        logger.info("Computing insertion metric...")
        results["insertion"] = self.insertion(image, target_class=target_class)

        # Sensitivity
        logger.info("Computing sensitivity...")
        results["sensitivity"] = self.sensitivity_n(image, target_class=target_class)

        # Sanity checks
        logger.info("Running model parameter randomization test...")
        results["param_randomization"] = self.model_parameter_randomization_test(
            image, target_class=target_class
        )

        logger.info("Running data randomization test...")
        results["data_randomization"] = self.data_randomization_test(
            image, target_class=target_class
        )

        # Pointing game (if bbox provided)
        if bbox is not None:
            results["pointing_game"] = self.pointing_game(
                image, bbox, target_class=target_class
            )

        return results

    # Helper methods

    def _preprocess_image(
        self, image: Union[Image.Image, np.ndarray, torch.Tensor]
    ) -> torch.Tensor:
        """Preprocess image to tensor."""
        if isinstance(image, torch.Tensor):
            if image.ndim == 3:
                image = image.unsqueeze(0)
            return image.to(self.device)

        if isinstance(image, np.ndarray):
            if image.dtype in [np.float32, np.float64]:
                if image.max() <= 1.0:
                    image = (image * 255).astype(np.uint8)
            image = Image.fromarray(image)

        # Convert PIL to tensor
        import torchvision.transforms as T

        transform = T.Compose([T.ToTensor()])
        tensor = transform(image).unsqueeze(0)
        return tensor.to(self.device)

    def _get_prediction(
        self, input_tensor: torch.Tensor, target_class: Optional[int] = None
    ):
        """Get model prediction."""
        with torch.no_grad():
            output = self.model(input_tensor)
            if isinstance(output, dict):
                output = output.get(
                    "logits", output.get("output", list(output.values())[0])
                )
            probs = torch.softmax(output, dim=1)

            if target_class is None:
                target_class = output.argmax(dim=1).item()

            pred_class = output.argmax(dim=1).item()
            score = probs[0, target_class].item()

        return pred_class, score, target_class

    def _create_baseline(
        self, input_tensor: torch.Tensor, baseline: str
    ) -> torch.Tensor:
        """Create baseline tensor."""
        if baseline == "black":
            return torch.zeros_like(input_tensor)
        elif baseline == "mean":
            mean = input_tensor.mean()
            return torch.full_like(input_tensor, mean)
        elif baseline == "blur":
            import torchvision.transforms.functional as TF

            # Apply Gaussian blur
            blurred = TF.gaussian_blur(input_tensor, kernel_size=51, sigma=20.0)
            return blurred
        else:
            raise ValueError(f"Unknown baseline: {baseline}")

    def _resize_heatmap(self, heatmap: np.ndarray, target_size: tuple) -> np.ndarray:
        """Resize heatmap to target size."""
        import cv2

        return cv2.resize(heatmap, (target_size[1], target_size[0]))

    def _get_num_classes(self) -> int:
        """Get number of output classes."""
        try:
            # Try to infer from model
            dummy_input = torch.randn(1, 3, 224, 224).to(self.device)
            with torch.no_grad():
                output = self.model(dummy_input)
                if isinstance(output, dict):
                    output = output.get(
                        "logits", output.get("output", list(output.values())[0])
                    )
                return output.shape[1]
        except Exception:
            # Default to 1000 (ImageNet)
            return 1000


__all__ = [
    "ExplanationMetrics",
]
