import pandas as pd
import cv2
import numpy as np
from PIL import Image
import torch
import logging

logger = logging.getLogger(__name__)

class ExplainabilityEngine:
    """
    [LEGACY] Generates SHAP explanations for XGBoost predictions.
    """
    def __init__(self, model):
        self.model = model
        import shap
        self.explainer = shap.TreeExplainer(model)
        
    def generate_explanation(self, feature_df: pd.DataFrame) -> dict:
        """
        Returns feature importance contributions for a specific prediction.
        """
        shap_values = self.explainer.shap_values(feature_df)
        
        contributions = {}
        # shap_values could be a list if multi-class, or array if binary
        if isinstance(shap_values, list):
            sv = shap_values[1][0] # take positive class
        else:
            sv = shap_values[0]
            
        for i, col in enumerate(feature_df.columns):
            contributions[col] = float(sv[i])
            
        # Sort by absolute impact
        sorted_contrib = dict(sorted(contributions.items(), key=lambda item: abs(item[1]), reverse=True))
        
        return {
            "top_features": list(sorted_contrib.keys())[:3],
            "feature_contributions": sorted_contrib
        }

class AttentionRolloutExplainer:
    """
    Attention Rollout explainer that works with both ViT and SigLIP2 models.
    Generates heatmaps showing where the model focuses for its prediction.
    """
    def __init__(self, model_name: str):
        self.model_name = model_name
        self.model = None
        self.processor = None

    def _load(self):
        if self.model is not None:
            return
        try:
            # Try SigLIP2 first (our primary model)
            from transformers import AutoImageProcessor, SiglipForImageClassification
            self.processor = AutoImageProcessor.from_pretrained(self.model_name)
            self.model = SiglipForImageClassification.from_pretrained(
                self.model_name, output_attentions=True
            )
            self.model.eval()
            self._arch = "siglip"
            logger.info("AttentionRollout loaded SigLIP2 model: %s", self.model_name)
        except Exception:
            try:
                # Fallback to ViT
                from transformers import ViTForImageClassification, ViTImageProcessor
                self.processor = ViTImageProcessor.from_pretrained(self.model_name)
                self.model = ViTForImageClassification.from_pretrained(
                    self.model_name, output_attentions=True
                )
                self.model.eval()
                self._arch = "vit"
                logger.info("AttentionRollout loaded ViT model: %s", self.model_name)
            except Exception as e:
                logger.error("AttentionRollout could not load model %s: %s", self.model_name, e)
                raise

    def generate_attention_map(self, pil_image: Image.Image, save_path: str) -> str | None:
        try:
            self._load()
            inputs = self.processor(images=pil_image, return_tensors="pt")
            with torch.no_grad():
                outputs = self.model(**inputs)

            attentions = outputs.attentions
            if not attentions:
                logger.warning("Model returned no attention weights")
                return None

            # Attention rollout: multiply attention matrices across layers
            result = torch.eye(attentions[0].size(-1))
            for attention in attentions:
                attention_heads_fused = attention.squeeze(0).mean(dim=0)
                flat = attention_heads_fused + torch.eye(attention_heads_fused.size(0))
                flat = flat / flat.sum(dim=-1, keepdim=True)
                result = torch.matmul(flat, result)

            # Extract spatial attention
            # SigLIP2 may not have CLS token — use mean of all token-to-token attention
            num_tokens = result.size(0)
            if self._arch == "vit" and num_tokens > 1:
                # ViT has CLS token at index 0
                mask = result[0, 1:]
            else:
                # SigLIP2 / no CLS: average attention received by each spatial token
                mask = result.mean(dim=0)

            # Reshape to 2D grid
            num_patches = mask.size(0)
            grid_size = int(np.sqrt(num_patches))
            if grid_size * grid_size != num_patches:
                # Try to find closest square
                grid_size = int(np.ceil(np.sqrt(num_patches)))
                mask = mask[:grid_size * grid_size]  # truncate if needed

            mask = mask.reshape(grid_size, grid_size).numpy()
            if np.max(mask) > 0:
                mask = mask / np.max(mask)

            # Upsample and blend with original image
            cv_img = cv2.cvtColor(np.array(pil_image.resize((224, 224))), cv2.COLOR_RGB2BGR)
            h, w = cv_img.shape[:2]

            mask_resized = cv2.resize(mask, (w, h), interpolation=cv2.INTER_CUBIC)
            mask_resized = np.uint8(255 * mask_resized)
            heatmap = cv2.applyColorMap(mask_resized, cv2.COLORMAP_JET)

            blended = cv2.addWeighted(heatmap, 0.4, cv_img, 0.6, 0)
            cv2.imwrite(save_path, blended)

            return save_path
        except Exception as e:
            logger.error("Attention rollout failed: %s", e)
            return None