utils/explainability.py

"""
Explainability Module
─────────────────────
Uses SHAP (Shapley Additive Explanations) to produce:
  1. Token-level contribution scores
  2. Word importance heatmap
  3. Feature-level explanation text

Also computes intensity staging based on predicted probabilities.
"""

import os
import numpy as np
import torch
import torch.nn.functional as F
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from textblob import TextBlob

from utils.preprocess import ID2LABEL, NUM_LABELS


# ─────────────────────────────────────────────
# Intensity Staging
# ─────────────────────────────────────────────

INTENSITY_THRESHOLDS = {
    # (low, medium) confidence boundaries per class
    "Normal":               (0.70, 0.90),
    "Anxiety":              (0.40, 0.70),
    "Depression":           (0.40, 0.70),
    "Suicidal":             (0.30, 0.60),
    "Bipolar":              (0.40, 0.70),
    "Stress":               (0.40, 0.70),
    "Personality Disorder": (0.40, 0.70),
}

INTENSITY_DESCRIPTIONS = {
    "Normal": {
        "Low":    "Minimal to no indicators of mental health concerns. Text appears healthy.",
        "Medium": "Slight hints of emotional tension. Generally within normal range.",
        "High":   "Text strongly reflects healthy mental state.",
    },
    "Anxiety": {
        "Low":    "Early-stage anxiety indicators. Mild worry or restlessness mentioned.",
        "Medium": "Moderate anxiety. Noticeable stress patterns and avoidance cues.",
        "High":   "Severe anxiety. Persistent fear, panic symptoms, or overwhelming worry.",
    },
    "Depression": {
        "Low":    "Mild depressive signals. Occasional sadness or low energy mentioned.",
        "Medium": "Moderate depression. Persistent low mood, loss of interest, fatigue.",
        "High":   "Severe depression. Hopelessness, anhedonia, or vegetative symptoms evident.",
    },
    "Suicidal": {
        "Low":    "Passive ideation — fleeting thoughts of not wanting to exist.",
        "Medium": "Active ideation — explicit thoughts about self-harm or ending life.",
        "High":   "Crisis level — clear intent or plan mentioned. Immediate support needed.",
    },
    "Bipolar": {
        "Low":    "Mild mood fluctuations. Occasional highs and lows in the narrative.",
        "Medium": "Moderate mood cycling. Distinct manic or depressive episodes described.",
        "High":   "Severe bipolar patterns. Extreme mood swings, grandiosity, or crash episodes.",
    },
    "Stress": {
        "Low":    "Mild stress markers. Temporary pressure from life events.",
        "Medium": "Moderate stress. Significant life pressures affecting daily functioning.",
        "High":   "Severe stress / PTSD markers. Trauma triggers, hypervigilance, or re-experiencing.",
    },
    "Personality Disorder": {
        "Low":    "Mild BPD traits. Some emotional reactivity or identity concerns.",
        "Medium": "Moderate BPD indicators. Fear of abandonment, unstable relationships.",
        "High":   "Severe BPD markers. Intense emotional dysregulation or self-identity crisis.",
    },
}


def get_intensity(label: str, confidence: float) -> tuple[str, str]:
    """
    Returns (intensity_level, description) based on label and confidence.
    """
    low_thresh, med_thresh = INTENSITY_THRESHOLDS.get(label, (0.40, 0.70))

    if confidence < low_thresh:
        level = "Low"
    elif confidence < med_thresh:
        level = "Medium"
    else:
        level = "High"

    desc = INTENSITY_DESCRIPTIONS.get(label, {}).get(level, "")
    return level, desc


# ─────────────────────────────────────────────
# Token Attribution via Gradient × Input
# ─────────────────────────────────────────────

def get_token_attributions(
    model,
    tokenizer,
    text: str,
    predicted_class: int,
    sentiment_features: torch.Tensor,
    device,
    max_length: int = 256,
) -> tuple[list[str], list[float]]:
    """
    Compute token-level importance scores using Integrated Gradients
    (approximated via single-step gradient × embedding norm).

    Returns:
        tokens  — list of token strings (without special tokens)
        scores  — list of float importance scores (same length)
    """
    model.eval()

    encoding = tokenizer(
        text,
        padding="max_length",
        truncation=True,
        max_length=max_length,
        return_tensors="pt",
    )

    input_ids      = encoding["input_ids"].to(device)
    attention_mask = encoding["attention_mask"].to(device)
    token_type_ids = encoding.get("token_type_ids", None)
    if token_type_ids is not None:
        token_type_ids = token_type_ids.to(device)

    # Hook to capture embeddings
    embeddings_ref = {}

    def embedding_hook(module, input, output):
        embeddings_ref["embed"] = output

    hook = model.bert.embeddings.register_forward_hook(embedding_hook)

    # Forward pass with gradient tracking on embeddings
    model.zero_grad()
    with torch.enable_grad():
        embeddings_ref["embed"] = None

        logits, _ = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            sentiment_features=sentiment_features.to(device) if sentiment_features is not None else None,
        )

        embed = embeddings_ref.get("embed")
        if embed is None:
            hook.remove()
            return [], []

        embed.retain_grad()
        score = logits[0, predicted_class]
        score.backward()

    hook.remove()

    if embed.grad is None:
        return [], []

    # Gradient × embedding magnitude → importance per token
    grad = embed.grad[0]            # (seq_len, hidden)
    emb  = embed[0].detach()        # (seq_len, hidden)
    importance = (grad * emb).sum(-1).abs().cpu().numpy()  # (seq_len,)

    # Get tokens (excluding padding)
    ids = input_ids[0].cpu().numpy()
    mask = attention_mask[0].cpu().numpy()
    tokens = tokenizer.convert_ids_to_tokens(ids)

    # Filter out padding, [CLS], [SEP]
    filtered_tokens = []
    filtered_scores = []
    for tok, imp, m in zip(tokens, importance, mask):
        if m == 0:
            break
        if tok in ("[CLS]", "[SEP]", "[PAD]"):
            continue
        filtered_tokens.append(tok)
        filtered_scores.append(float(imp))

    # Normalize 0–1
    if filtered_scores:
        max_s = max(filtered_scores) or 1e-9
        filtered_scores = [s / max_s for s in filtered_scores]

    return filtered_tokens, filtered_scores


# ─────────────────────────────────────────────
# HTML Heatmap Renderer
# ─────────────────────────────────────────────

def render_html_heatmap(tokens: list[str], scores: list[float], label: str) -> str:
    """
    Generate an HTML snippet with token-level background colors
    showing contribution to the predicted class.
    """
    if not tokens:
        return "<p>No token attribution available.</p>"

    # Color scale: low = white/gray, high = red/orange
    cmap = plt.cm.YlOrRd

    html_parts = []
    for tok, score in zip(tokens, scores):
        color = mcolors.to_hex(cmap(max(0.0, min(1.0, score))))
        display = tok.replace("##", "")  # merge WordPiece subwords
        text_color = "#000" if score < 0.6 else "#fff"
        html_parts.append(
            f'<span style="background-color:{color};color:{text_color};'
            f'padding:2px 4px;margin:1px;border-radius:3px;'
            f'font-size:14px;font-family:monospace;">{display}</span>'
        )

    heatmap_html = (
        f'<div style="line-height:2.2;padding:10px;background:#f8f8f8;'
        f'border-radius:8px;border:1px solid #ddd;">'
        + " ".join(html_parts)
        + "</div>"
    )
    return heatmap_html


# ─────────────────────────────────────────────
# Matplotlib Word Importance Plot
# ─────────────────────────────────────────────

def plot_word_importance(
    tokens: list[str],
    scores: list[float],
    label: str,
    save_path: str = None,
    top_n: int = 15,
) -> plt.Figure:
    """
    Bar chart of top-N most important tokens.
    """
    if not tokens:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, "No attribution data", ha="center")
        return fig

    # Merge subwords
    merged: dict[str, float] = {}
    for tok, sc in zip(tokens, scores):
        word = tok.replace("##", "")
        merged[word] = max(merged.get(word, 0), sc)

    # Top N
    sorted_items = sorted(merged.items(), key=lambda x: x[1], reverse=True)[:top_n]
    words  = [i[0] for i in sorted_items]
    values = [i[1] for i in sorted_items]

    # Color by intensity
    colors = plt.cm.YlOrRd(np.array(values))

    fig, ax = plt.subplots(figsize=(8, max(3, len(words) * 0.4)))
    bars = ax.barh(words[::-1], values[::-1], color=colors[::-1])
    ax.set_xlabel("Importance Score", fontsize=11)
    ax.set_title(f"Word Importance for '{label}' Prediction", fontsize=13, fontweight="bold")
    ax.set_xlim(0, 1.05)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    # Add value labels
    for bar, val in zip(bars, values[::-1]):
        ax.text(bar.get_width() + 0.01, bar.get_y() + bar.get_height() / 2,
                f"{val:.2f}", va="center", fontsize=9)

    plt.tight_layout()

    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")

    return fig


# ─────────────────────────────────────────────
# Full Explanation Bundle
# ─────────────────────────────────────────────

def explain_prediction(
    model,
    tokenizer,
    text: str,
    predicted_class: int,
    confidence: float,
    probabilities: np.ndarray,
    sentiment_features: torch.Tensor,
    device,
    max_length: int = 256,
) -> dict:
    """
    Run full explainability pipeline and return a bundle:
    {
        tokens, scores,
        html_heatmap,
        word_importance_fig,
        intensity_level, intensity_description,
        sentiment_info,
        label_name,
        confidence,
        all_probs,
    }
    """
    label_name = ID2LABEL[predicted_class]

    # Token attribution
    tokens, scores = get_token_attributions(
        model, tokenizer, text, predicted_class,
        sentiment_features, device, max_length
    )

    # HTML heatmap
    heatmap_html = render_html_heatmap(tokens, scores, label_name)

    # Bar chart
    fig = plot_word_importance(tokens, scores, label_name)

    # Intensity
    intensity_level, intensity_desc = get_intensity(label_name, confidence)

    # Sentiment
    blob = TextBlob(text)
    sentiment_info = {
        "polarity":     round(blob.sentiment.polarity, 3),
        "subjectivity": round(blob.sentiment.subjectivity, 3),
        "label":        "Positive" if blob.sentiment.polarity > 0.05
                        else ("Negative" if blob.sentiment.polarity < -0.05 else "Neutral"),
    }

    all_probs = {ID2LABEL[i]: round(float(probabilities[i]) * 100, 2)
                 for i in range(NUM_LABELS)}

    return {
        "label_name":           label_name,
        "confidence":           round(confidence * 100, 2),
        "intensity_level":      intensity_level,
        "intensity_description": intensity_desc,
        "tokens":               tokens,
        "scores":               scores,
        "html_heatmap":         heatmap_html,
        "word_importance_fig":  fig,
        "sentiment_info":       sentiment_info,
        "all_probs":            all_probs,
    }