"""SHAP explainer for Myanmar Ghost model.

Uses SHAP (SHapley Additive exPlanations) to explain
individual predictions and word importance.
"""

import logging
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import numpy as np
import shap
import torch
from tqdm import tqdm

logger = logging.getLogger(__name__)


class ThankingSHAPExplainer:
    """SHAP-based explainer for Myanmar text classification."""
    
    def __init__(
        self,
        model,
        tokenizer,
        background_size: int = 100,
        device: str = "cuda" if torch.cuda.is_available() else "cpu",
    ):
        """
        Args:
            model: PyTorch model or HuggingFace model
            tokenizer: Tokenizer for the model
            background_size: Number of background samples for SHAP
            device: Device to run on
        """
        self.model = model
        self.tokenizer = tokenizer
        self.background_size = background_size
        self.device = device
        
        self.model.to(device)
        self.model.eval()
        
        self.explainer = None
        self.background_data = None
    
    def _get_tokenizer(self):
        """Get the tokenizer, handling both HF and custom tokenizers."""
        if hasattr(self.tokenizer, "__call__"):
            return self.tokenizer
        return self.tokenizer.encode
    
    def _predict(self, texts: Union[List[str], np.ndarray]) -> np.ndarray:
        """Model prediction function for SHAP."""
        if isinstance(texts, np.ndarray):
            texts = texts.tolist()
        
        # Tokenize
        if hasattr(self.tokenizer, "batch_encode_plus"):
            encoding = self.tokenizer.batch_encode_plus(
                texts,
                padding=True,
                truncation=True,
                max_length=128,
                return_tensors="pt",
            )
            input_ids = encoding["input_ids"].to(self.device)
            attention_mask = encoding["attention_mask"].to(self.device)
        else:
            input_ids = torch.tensor(
                [self.tokenizer.encode(t) for t in texts]
            ).to(self.device)
            attention_mask = (input_ids != 0).long().to(self.device)
        
        with torch.no_grad():
            outputs = self.model(input_ids, attention_mask)
            
            if hasattr(outputs, "logits"):
                logits = outputs.logits
            else:
                logits = outputs
            
            probs = torch.softmax(logits, dim=-1).cpu().numpy()
        
        return probs
    
    def fit_background(
        self,
        background_texts: List[str],
    ) -> None:
        """Fit background distribution for SHAP.
        
        Args:
            background_texts: List of texts to use as background
        """
        logger.info(f"Fitting SHAP background with {len(background_texts)} samples")
        
        # Sample background if too large
        if len(background_texts) > self.background_size:
            indices = np.random.choice(
                len(background_texts),
                self.background_size,
                replace=False,
            )
            background_texts = [background_texts[i] for i in indices]
        
        self.background_data = background_texts
        
        # Create SHAP explainer
        self.explainer = shap.Explainer(
            self._predict,
            self.tokenizer,
            output_names=["negative", "neutral", "positive", "sarcastic"],
        )
        
        # Calculate background values
        logger.info("Computing SHAP values for background...")
        self.explainer(background_texts[:min(10, len(background_texts))])
        
        logger.info("Background fitting complete")
    
    def explain(
        self,
        text: str,
        num_samples: int = 100,
        output_names: Optional[List[str]] = None,
    ) -> shap.Explanation:
        """Explain a single text.
        
        Args:
            text: Myanmar text to explain
            num_samples: Number of Monte Carlo samples
            output_names: Names for output classes
        
        Returns:
            SHAP Explanation object
        """
        if self.explainer is None:
            logger.warning("No background data. Using default explainer.")
            self.explainer = shap.Explainer(
                self._predict,
                self.tokenizer,
                output_names=output_names or ["negative", "neutral", "positive", "sarcastic"],
            )
        
        logger.info(f"Explaining text: {text[:50]}...")
        
        explainer = shap.Explainer(
            self._predict,
            self.tokenizer,
            output_names=output_names,
        )
        
        shap_values = explainer([text])
        
        return shap_values
    
    def explain_batch(
        self,
        texts: List[str],
        output_names: Optional[List[str]] = None,
    ) -> List[shap.Explanation]:
        """Explain multiple texts.
        
        Args:
            texts: List of Myanmar texts
            output_names: Names for output classes
        
        Returns:
            List of SHAP Explanation objects
        """
        if output_names is None:
            output_names = ["negative", "neutral", "positive", "sarcastic"]
        
        explainer = shap.Explainer(
            self._predict,
            self.tokenizer,
            output_names=output_names,
        )
        
        explanations = []
        for text in tqdm(texts, desc="Explaining texts"):
            exp = explainer([text])
            explanations.append(exp)
        
        return explanations
    
    def get_word_importance(
        self,
        text: str,
        class_index: int = 2,  # positive by default
    ) -> List[Tuple[str, float]]:
        """Get word importance scores for a specific class.
        
        Args:
            text: Myanmar text
            class_index: Class index to explain
        
        Returns:
            List of (word, importance) tuples
        """
        explanation = self.explain(text)
        
        # Get tokens and their SHAP values
        tokens = self.tokenizer.tokenize(text)
        shap_vals = explanation.values[0, :, class_index]
        
        # Handle tokenization differences
        if len(shap_vals) < len(tokens):
            # Pad if needed
            shap_vals = np.pad(
                shap_vals,
                (0, len(tokens) - len(shap_vals)),
                constant_values=0,
            )
        elif len(shap_vals) > len(tokens):
            tokens = tokens + ["[PAD]"] * (len(shap_vals) - len(tokens))
        
        # Create word-score pairs
        word_importance = list(zip(tokens, shap_vals.tolist()))
        
        # Sort by absolute importance
        word_importance.sort(key=lambda x: abs(x[1]), reverse=True)
        
        return word_importance
    
    def visualize_text(
        self,
        explanation: shap.Explanation,
        output_path: Optional[str] = None,
    ) -> None:
        """Visualize SHAP explanation as text.
        
        Args:
            explanation: SHAP explanation
            output_path: Optional path to save visualization
        """
        text = explanation.data[0] if hasattr(explanation, "data") else ""
        
        output = f"\n{'='*60}\n"
        output += f"Text: {text}\n"
        output += f"{'='*60}\n"
        
        # Get top features for each class
        for i, class_name in enumerate(explanation.output_names):
            values = explanation.values[0, :, i]
            
            # Get top 5 words
            top_indices = np.argsort(np.abs(values))[-5:][::-1]
            
            output += f"\nClass: {class_name}\n"
            output += "-" * 40 + "\n"
            
            for idx in top_indices:
                if idx < len(text.split()):
                    word = text.split()[idx]
                    output += f"  {word}: {values[idx]:.4f}\n"
        
        print(output)
        
        if output_path:
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(output)
            logger.info(f"Visualization saved to {output_path}")


class ThankingSHAPValues:
    """Compute SHAP values for thanking expression analysis."""
    
    def __init__(
        self,
        model,
        tokenizer,
        class_names: List[str] = None,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.class_names = class_names or [
            "genuine", "sarcastic", "complaining", "neutral"
        ]
    
    def compute_feature_importance(
        self,
        texts: List[str],
        feature_names: List[str],
    ) -> Dict[str, float]:
        """Compute SHAP-based feature importance.
        
        Args:
            texts: List of texts
            feature_names: Names of features to analyze
        
        Returns:
            Dictionary of feature importance scores
        """
        import shap
        
        def predict_proba(texts: List[str]) -> np.ndarray:
            encoding = self.tokenizer.batch_encode_plus(
                texts,
                padding=True,
                truncation=True,
                max_length=128,
                return_tensors="pt",
            )
            
            with torch.no_grad():
                outputs = self.model(
                    encoding["input_ids"],
                    encoding["attention_mask"],
                )
                probs = torch.softmax(outputs.logits, dim=-1)
                return probs.numpy()
        
        # Create simple background
        background = texts[:min(20, len(texts))]
        
        explainer = shap.Explainer(predict_proba, background)
        shap_values = explainer(texts[:5])  # Sample for speed
        
        # Aggregate importance
        importance = {}
        for i, feature in enumerate(feature_names):
            importance[feature] = np.mean(
                np.abs(shap_values.values[:, :, i])
            )
        
        return importance
    
    def analyze_sentence(
        self,
        text: str,
    ) -> Dict[str, Any]:
        """Analyze a single sentence with SHAP.
        
        Args:
            text: Myanmar text
        
        Returns:
            Analysis results including word importance
        """
        import shap
        
        encoding = self.tokenizer(
            text,
            padding=True,
            truncation=True,
            max_length=128,
            return_tensors="pt",
        )
        
        explainer = shap.Explainer(
            lambda x: self._predict_batch(x),
            self.tokenizer,
        )
        
        explanation = explainer([text])
        
        # Extract word-level importance
        tokens = self.tokenizer.convert_ids_to_tokens(
            encoding["input_ids"][0]
        )
        
        result = {
            "text": text,
            "tokens": tokens,
            "prediction": explanation.output_names[
                np.argmax(explanation.values[0].mean(axis=1))
            ],
            "shap_values": explanation.values[0].tolist(),
        }
        
        return result
    
    def _predict_batch(self, texts: List[str]) -> np.ndarray:
        """Batch prediction for SHAP."""
        encoding = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=128,
            return_tensors="pt",
        )
        
        with torch.no_grad():
            outputs = self.model(
                encoding["input_ids"],
                encoding["attention_mask"],
            )
            probs = torch.softmax(outputs.logits, dim=-1)
            return probs.numpy()


def create_shap_explainer(
    model,
    tokenizer,
    background_texts: Optional[List[str]] = None,
) -> ThankingSHAPExplainer:
    """Factory function to create SHAP explainer."""
    explainer = ThankingSHAPExplainer(model, tokenizer)
    
    if background_texts:
        explainer.fit_background(background_texts)
    
    return explainer


if __name__ == "__main__":
    print("ThankingSHAPExplainer loaded")
    print("Use create_shap_explainer() to create an explainer")