src/classify/explain.py

"""SHAP-based top-token explanations for the linear models (§3.4).

Why only Logistic Regression and only top-token bar charts?
-----------------------------------------------------------
- README §3 deliverables list ``lr_shap_classaction.png`` and
  ``lr_shap_casetype.png`` — both are LR.
- For a sparse linear model on TF-IDF features, SHAP's ``LinearExplainer``
  reduces to the model's own (mean-centred) coefficients × feature value,
  so we get exact, fast attributions without sampling. This is the
  right tool — kernel/tree explainers would be slow and stochastic here.
- Bar chart of the top-k positive- and negative-impact tokens is what
  slide 8 needs: a story like "the words 'class' and 'representative'
  push the model toward 'class action sought'".

For the multi-class case-type plot we tile per-class bar charts so each
of the 5 case-type groups gets its own most-distinctive vocabulary.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline

from src.utils import RESULTS_DIR, get_logger

logger = get_logger(__name__)


@dataclass(frozen=True)
class TokenImpact:
    """One row of the bar chart: a token + its average signed SHAP value."""

    token: str
    mean_shap: float


def _ensure_linear(pipe: Pipeline) -> tuple[TfidfVectorizer, LogisticRegression]:
    """Defensive unpack — only LR + TF-IDF pipelines are supported here."""
    vectorizer = pipe.named_steps.get("tfidf")
    classifier = pipe.named_steps.get("clf")
    if not isinstance(vectorizer, TfidfVectorizer):
        raise TypeError("explain.py expects a 'tfidf' TfidfVectorizer step.")
    if not isinstance(classifier, LogisticRegression):
        raise TypeError(
            f"SHAP top-token explanation is implemented for LogisticRegression "
            f"only (got {type(classifier).__name__})."
        )
    return vectorizer, classifier


def _shap_values(
    vectorizer: TfidfVectorizer,
    classifier: LogisticRegression,
    X_texts: np.ndarray,
    *,
    background_size: int = 100,
) -> tuple[np.ndarray, list[str]]:
    """Return (shap_values, vocabulary) where shap_values has shape (n_samples, n_features).

    For multi-class LR with K classes, shap returns a list of K arrays; we
    stack into shape (K, n_samples, n_features) so downstream code can
    pick a class index.
    """
    import shap  # local import: optional dependency

    X = vectorizer.transform(X_texts)
    background_size = min(background_size, X.shape[0])
    rng = np.random.default_rng(42)
    bg_idx = rng.choice(X.shape[0], size=background_size, replace=False)
    background = X[bg_idx]

    explainer = shap.LinearExplainer(classifier, background)
    values = explainer.shap_values(X)
    vocab = vectorizer.get_feature_names_out().tolist()

    # Normalize shape to (K, n_samples, n_features) across SHAP versions:
    #   - old (<=0.42)        : list of K arrays, each (N, F)
    #   - new multi-class     : ndarray (N, F, K)
    #   - binary              : ndarray (N, F)
    if isinstance(values, list):
        values = np.stack(values, axis=0)
    elif values.ndim == 3:
        values = np.transpose(values, (2, 0, 1))
    else:
        values = values[np.newaxis, ...]
    return values, vocab


def top_tokens_per_class(
    shap_values: np.ndarray,
    vocab: list[str],
    *,
    top_k: int = 15,
) -> list[list[TokenImpact]]:
    """Return, for each class, the top-k tokens by absolute mean SHAP impact.

    Each token's score keeps its sign so the bar chart shows direction:
    positive bars push *toward* the class, negative bars push *away*.
    """
    out: list[list[TokenImpact]] = []
    n_classes = shap_values.shape[0]
    for cls in range(n_classes):
        mean_shap = shap_values[cls].mean(axis=0)
        # Some sklearn versions emit a 1-D ndarray, others a matrix-like; flatten safely.
        mean_shap = np.asarray(mean_shap).ravel()
        idx = np.argsort(-np.abs(mean_shap))[:top_k]
        rows = [TokenImpact(token=vocab[i], mean_shap=float(mean_shap[i])) for i in idx]
        out.append(rows)
    return out


def _plot_binary(
    tokens: list[TokenImpact],
    *,
    positive_label: str,
    negative_label: str,
    title: str,
    output_path: Path,
) -> None:
    fig, ax = plt.subplots(figsize=(7, 6))
    tokens_sorted = sorted(tokens, key=lambda t: t.mean_shap)
    labels = [t.token for t in tokens_sorted]
    values = [t.mean_shap for t in tokens_sorted]
    colors = ["#1f77b4" if v >= 0 else "#d62728" for v in values]
    ax.barh(labels, values, color=colors)
    ax.axvline(0, color="black", linewidth=0.5)
    ax.set_xlabel(f"← pushes toward '{negative_label}'    |    pushes toward '{positive_label}' →")
    ax.set_title(title)
    fig.tight_layout()
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    logger.info("Saved SHAP plot → %s", output_path)


def _plot_multiclass(
    tokens_per_class: list[list[TokenImpact]],
    label_names: list[str],
    *,
    title: str,
    output_path: Path,
) -> None:
    """Tile one mini bar chart per class so each row of slide 9 has its own vocabulary."""
    n_classes = len(tokens_per_class)
    n_cols = min(3, n_classes)
    n_rows = (n_classes + n_cols - 1) // n_cols
    fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows))
    axes = np.atleast_1d(axes).ravel()
    for idx, (cls_name, tokens) in enumerate(zip(label_names, tokens_per_class)):
        tokens_sorted = sorted(tokens, key=lambda t: t.mean_shap)
        labels = [t.token for t in tokens_sorted]
        values = [t.mean_shap for t in tokens_sorted]
        colors = ["#1f77b4" if v >= 0 else "#d62728" for v in values]
        axes[idx].barh(labels, values, color=colors)
        axes[idx].axvline(0, color="black", linewidth=0.5)
        axes[idx].set_title(cls_name, fontsize=10)
    # Hide any leftover panels in the grid.
    for k in range(len(label_names), len(axes)):
        axes[k].set_visible(False)
    fig.suptitle(title, fontsize=12)
    fig.tight_layout(rect=(0, 0, 1, 0.97))
    fig.savefig(output_path, dpi=150)
    plt.close(fig)
    logger.info("Saved SHAP plot → %s", output_path)


def explain_lr(
    pipe: Pipeline,
    X_test: np.ndarray,
    label_names: list[str],
    *,
    task: str,
    top_k: int = 15,
    output_dir: Path | None = None,
) -> Path:
    """Top entry point: pickled LR pipeline → SHAP PNG under results/."""
    vectorizer, classifier = _ensure_linear(pipe)
    shap_values, vocab = _shap_values(vectorizer, classifier, X_test)
    tokens_per_class = top_tokens_per_class(shap_values, vocab, top_k=top_k)

    out_dir = output_dir or RESULTS_DIR
    tag = "classaction" if task == "class_action" else "casetype"
    output_path = out_dir / f"lr_shap_{tag}.png"

    if len(label_names) == 2:
        # Binary LR has 1 row of shap values, signed toward class index 1.
        _plot_binary(
            tokens_per_class[-1],
            positive_label=label_names[1],
            negative_label=label_names[0],
            title=f"SHAP top tokens — LR · {task}",
            output_path=output_path,
        )
    else:
        _plot_multiclass(
            tokens_per_class,
            label_names,
            title=f"SHAP top tokens — LR · {task}",
            output_path=output_path,
        )
    return output_path


def build_arg_parser():
    import argparse

    parser = argparse.ArgumentParser(description="Generate SHAP top-token plot for an LR pipeline.")
    parser.add_argument("--task", required=True, choices=["class_action", "case_type"])
    parser.add_argument("--pipeline", help="Path to pickled LR pipeline. Defaults to models/lr_{tag}.pkl.")
    parser.add_argument("--top-k", type=int, default=15)
    parser.add_argument("--text-source", default="long_ref",
                        choices=["long_ref", "long_pred", "source_text"])
    return parser


def main() -> None:
    import pickle

    from src.classify.data import load_classification_data
    from src.utils import MODELS_DIR

    args = build_arg_parser().parse_args()
    tag = "classaction" if args.task == "class_action" else "casetype"
    pipeline_path = Path(args.pipeline) if args.pipeline else MODELS_DIR / f"lr_{tag}.pkl"
    with pipeline_path.open("rb") as fh:
        pipe = pickle.load(fh)

    data = load_classification_data(task=args.task, text_source=args.text_source)
    explain_lr(pipe, data.X_test, data.label_names, task=args.task, top_k=args.top_k)


if __name__ == "__main__":
    main()