src/explainer.py

"""Explainability for the Proposed Model (BertMetaFusionACSAModel).

Three views:
  1. Cross-Attention weights over metadata tokens   (per aspect, per example)
  2. Integrated Gradients on the input text         (per aspect, per example)
  3. Aspect-level aggregation plots                 (category x aspect)

Saves an HTML report combining (1) and (2), plus PNG figures for (3).
"""
import html as html_lib
import logging
from pathlib import Path
from typing import List, Optional, Dict, Tuple

import numpy as np
import pandas as pd
import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns

from . import config as cfg
from .evaluator import load_meta_acsa
from .meta_encoder import MetaEncoder
from .models import META_NUM_META_TOKENS


logger = logging.getLogger(__name__)

META_TOKEN_NAMES = [f"meta_chunk_{i+1}" for i in range(META_NUM_META_TOKENS)]
SEMANTIC_META_TOKEN_NAMES = ["features", "categories", "numeric"]


# ---------------------------------------------------------------------------
# Low-level rendering helpers
# ---------------------------------------------------------------------------

def _normalize(arr):
    arr = np.asarray(arr, dtype=np.float64)
    if arr.size == 0:
        return np.zeros_like(arr)
    finite = arr[np.isfinite(arr)]
    if finite.size == 0 or finite.max() == finite.min():
        return np.zeros_like(arr)
    return (arr - finite.min()) / (finite.max() - finite.min() + 1e-9)


def _normalize_for_display(scores, mask=None):
    arr = np.asarray(scores, dtype=np.float64)
    if mask is not None:
        arr = arr * np.asarray(mask, dtype=np.float64)
    arr = np.abs(arr)
    if arr.size == 0 or float(arr.max()) <= 0:
        return np.zeros_like(arr)
    denom = np.percentile(arr[arr > 0], 95) if np.any(arr > 0) else arr.max()
    denom = max(float(denom), 1e-9)
    return np.clip(arr / denom, 0.0, 1.0)


def _color_for_score(score: float, label: int) -> str:
    score = float(np.clip(score, 0.0, 1.0))
    if score <= 0:
        return "rgba(255,255,255,0)"
    alpha = 0.18 + 0.72 * score
    if label == 2:
        return f"rgba(255, 80, 80, {alpha:.3f})"
    if label == 1:
        return f"rgba(80, 180, 110, {alpha:.3f})"
    return f"rgba(150, 150, 150, {alpha:.3f})"


def _merge_wordpieces(tokens, scores):
    words, vals = [], []
    for tok, score in zip(tokens, scores):
        tok = str(tok)
        if not tok or (tok.startswith("[") and tok.endswith("]")):
            continue
        if tok.startswith("##") and words:
            words[-1] += tok[2:]
            vals[-1] = max(vals[-1], float(score))
        else:
            words.append(tok)
            vals.append(float(score))
    return words, vals


def _top_terms(tokens, scores, top_k: int = 6):
    words, vals = _merge_wordpieces(tokens, scores)
    pairs = [(w, v) for w, v in zip(words, vals) if _is_informative_term(w) and v > 0]
    pairs.sort(key=lambda x: x[1], reverse=True)
    return _dedupe_terms(pairs, top_k=top_k)


def _explanation_quality(top_terms, meta_source: str, meta_weight: float) -> str:
    n_terms = len(top_terms or [])
    if n_terms >= 3 and meta_source and meta_weight >= 0.45:
        return "strong"
    if n_terms >= 1 and meta_source:
        return "moderate"
    return "weak"


def _make_explanation_sentence(aspect: str, pred_label: str, top_terms, meta_source: str, meta_weight: float) -> str:
    terms = [t for t, _ in (top_terms or [])[:4]]
    text_part = ", ".join(f'"{t}"' for t in terms) if terms else "no strong text token"
    meta_part = f'"{meta_source}" ({meta_weight:.3f})' if meta_source else "no dominant metadata source"
    return (
        f'{aspect} is predicted as {pred_label}. '
        f'The main text evidence is {text_part}, and the strongest metadata source is {meta_part}.'
    )


def _render_token_html(tokens, scores, label, top_k: int = 18):
    norm = _normalize_for_display(scores)
    positive_idx = np.where(norm > 0)[0]
    if len(positive_idx) == 0:
        words, vals = _merge_wordpieces(tokens, norm)
    else:
        ranked = positive_idx[np.argsort(norm[positive_idx])]
        top_idx = set(ranked[-min(top_k, len(ranked)):].tolist())
        display = np.zeros_like(norm)
        for i in top_idx:
            display[i] = max(float(norm[i]), 0.35)
        words, vals = _merge_wordpieces(tokens, display)
    spans = []
    for word, val in zip(words, vals):
        color = _color_for_score(val, label)
        weight = "600" if val >= 0.55 else "400"
        border = " border-bottom:1px solid rgba(0,0,0,.18);" if val >= 0.35 else ""
        spans.append(
            f'<span style="background-color:{color};padding:2px 4px;'
            f'border-radius:4px;margin:1px;font-weight:{weight};{border}">'
            f'{html_lib.escape(word)}</span>'
        )
    return " ".join(spans)


def _meta_summary_from_row(row) -> Dict[str, str]:
    def first_existing(names, default=""):
        for name in names:
            if name not in row:
                continue
            val = row[name]
            if isinstance(val, (list, tuple)):
                return "; ".join(map(str, val[:8]))
            if pd.notna(val):
                return str(val)
        return default

    features = first_existing(["features_text", "features", "description", "title"])
    categories = first_existing(["categories_text", "category", "leaf_category", "main_category"])
    numeric_parts = []
    for col in ["price", "average_rating", "rating_number"]:
        if col in row and pd.notna(row[col]):
            numeric_parts.append(f"{col}={row[col]}")
    return {
        "features": features[:420],
        "categories": categories[:220],
        "numeric": ", ".join(numeric_parts) if numeric_parts else "not available",
    }


def _top_meta_source(meta_attn, token_names=None) -> Tuple[str, float]:
    arr = np.asarray(meta_attn, dtype=np.float64).reshape(-1)
    if arr.size == 0:
        return "", 0.0
    names = token_names or META_TOKEN_NAMES
    idx = int(np.argmax(arr))
    name = names[idx] if idx < len(names) else f"meta_{idx+1}"
    return name, float(arr[idx])

_ASPECT_EVIDENCE_KEYWORDS = {
    "SIZE": {"size", "fit", "fits", "fitting", "small", "large", "big", "tight", "loose", "xl", "medium", "waist", "length"},
    "MATERIAL": {"material", "fabric", "cotton", "polyester", "soft", "scratchy", "thin", "thick", "stretch", "leather", "wool"},
    "QUALITY": {"quality", "stitch", "stitching", "seam", "wash", "washed", "durable", "cheap", "broke", "tear", "torn"},
    "APPEARANCE": {"look", "looks", "color", "colour", "photo", "picture", "beautiful", "cute", "print", "design"},
    "STYLE": {"style", "stylish", "flattering", "casual", "formal", "dress", "shirt", "fashion", "compliments"},
    "VALUE": {"price", "worth", "value", "money", "cheap", "expensive", "discount", "penny", "cost"},
}

_SENTIMENT_EVIDENCE_KEYWORDS = {
    "good", "great", "love", "loved", "perfect", "nice", "excellent", "comfortable", "soft",
    "bad", "poor", "cheap", "terrible", "awful", "small", "large", "tight", "loose", "thin",
    "worth", "disappointed", "return", "returned", "recommend", "flattering", "beautiful",
}

_STOPWORDS = {
    "a", "an", "the", "and", "or", "but", "if", "then", "than", "so", "as", "at", "by", "for", "from",
    "in", "into", "of", "on", "to", "with", "without", "is", "are", "was", "were", "be", "been", "being",
    "it", "its", "this", "that", "these", "those", "i", "me", "my", "we", "our", "you", "your", "he", "she", "they",
    "them", "his", "her", "their", "very", "really", "just", "also", "too", "would", "could", "should", "can",
    "will", "did", "do", "does", "have", "has", "had", "there", "here", "about", "after", "before",
}


def _clean_term(term: str) -> str:
    return str(term).lower().replace("##", "").strip(".,!?;:'\"()[]{}<>/\\|`~@#$%^&*_+=")


def _is_informative_term(term: str) -> bool:
    clean = _clean_term(term)
    if len(clean) < 2 or clean in _STOPWORDS:
        return False
    return any(ch.isalpha() for ch in clean)


def _dedupe_terms(terms, top_k: int = 6):
    seen, out = set(), []
    for term, score in terms:
        clean = _clean_term(term)
        if not _is_informative_term(clean) or clean in seen:
            continue
        seen.add(clean)
        out.append((clean, float(score)))
        if len(out) >= top_k:
            break
    return out

def _lexical_evidence_scores(tokens, aspect_idx: int):
    aspect = cfg.ASPECTS[aspect_idx]
    aspect_terms = _ASPECT_EVIDENCE_KEYWORDS.get(aspect, set())
    scores = []
    prev = ""
    for tok in tokens:
        clean = str(tok).lower().replace("##", "")
        clean = clean.strip(".,!?;:'\"()[]{}")
        if not clean or clean.startswith("["):
            scores.append(0.0)
            prev = clean
            continue
        score = 0.0
        if clean in aspect_terms:
            score += 1.0
        if clean in _SENTIMENT_EVIDENCE_KEYWORDS:
            score += 0.55
        if prev in {"not", "no", "never", "too", "very", "really"} and score > 0:
            score += 0.25
        scores.append(score)
        prev = clean
    return np.asarray(scores, dtype=np.float32)


def _has_visible_scores(scores) -> bool:
    arr = np.asarray(scores, dtype=np.float64)
    return bool(arr.size and np.isfinite(arr).any() and float(np.nanmax(np.abs(arr))) > 1e-12)
# ---------------------------------------------------------------------------
# Integrated Gradients (per aspect)
# ---------------------------------------------------------------------------

def _ig_for_aspect(model, tokenizer, meta_encoder, row, device,
                   aspect_idx: int, n_steps: int = 30):
    """Gradient x embedding attribution for one aspect head.

    Captum's LayerIntegratedGradients can be brittle with wrapped transformer
    modules in notebooks. This path keeps the same goal, but computes a direct
    first-order attribution from the input embeddings, then falls back to
    lexical evidence if gradients are unavailable.
    """
    model.eval()
    text = str(row["full_text"])
    enc = tokenizer(text, max_length=cfg.MAX_LENGTH, truncation=True,
                    padding="max_length", return_tensors="pt")
    input_ids = enc["input_ids"].to(device)
    attention_mask = enc["attention_mask"].to(device)
    tokens = tokenizer.convert_ids_to_tokens(input_ids[0].detach().cpu().numpy().tolist())
    mask = attention_mask[0].detach().cpu().numpy().astype(bool)

    meta_vec = torch.from_numpy(
        meta_encoder.transform(pd.DataFrame([row]))
    ).float().to(device)

    try:
        model.zero_grad(set_to_none=True)
        embeds = model.bert.embeddings(input_ids=input_ids)
        embeds = embeds.detach().requires_grad_(True)
        bert_out = model.bert(
            inputs_embeds=embeds,
            attention_mask=attention_mask,
            return_dict=True,
        )
        text_vec = bert_out.last_hidden_state[:, 0, :]

        # Mirror GatedAspectSemanticMetaFusionACSAModel.forward from embeddings.
        meta_tokens = model.meta_tokenizer(meta_vec)
        meta_for_concat = meta_vec.clone()
        numeric_start = cfg.META_TFIDF_DIM
        meta_for_concat[:, numeric_start:] = (
            meta_for_concat[:, numeric_start:] * cfg.META_NUMERIC_TOKEN_SCALE
        )
        concat_meta = model.concat_meta_mlp(meta_for_concat)
        concat_base = model.concat_norm(
            model.concat_fusion(torch.cat([text_vec, concat_meta], dim=-1))
        )
        aspect_fused, _, _ = model.aspect_fusion(text_vec, meta_tokens)
        text_expanded = text_vec.unsqueeze(1).expand(-1, model.num_aspects, -1)
        aspect_delta = aspect_fused - text_expanded
        scale = torch.clamp(model.cross_residual_scale, 0.0, 1.0)
        aspect_repr = model.aspect_refine_norm(
            concat_base.unsqueeze(1) + scale * aspect_delta
        )
        logits = model.heads.forward_per_aspect(aspect_repr)
        target_class = int(logits[0, aspect_idx, :].argmax(dim=-1).item())
        target_logit = logits[0, aspect_idx, target_class]
        target_logit.backward()

        grad = embeds.grad.detach()[0]
        attr = (grad * embeds.detach()[0]).sum(dim=-1).detach().cpu().numpy()
        attr = np.abs(attr) * mask
        if not _has_visible_scores(attr):
            attr = _lexical_evidence_scores(tokens, aspect_idx) * mask
        return tokens, attr, mask, target_class
    except Exception as e:
        logger.warning("Gradient attribution failed for aspect %s: %s; using lexical evidence.",
                       cfg.ASPECTS[aspect_idx], e)
        with torch.no_grad():
            out = model(input_ids, attention_mask, meta_vec)
            target_class = int(out["logits"][0, aspect_idx, :].argmax(dim=-1).item())
        attr = _lexical_evidence_scores(tokens, aspect_idx) * mask
        return tokens, attr, mask, target_class

def _attention_for_aspect(model, tokenizer, meta_encoder, row, device,
                          aspect_idx: int):
    """Fallback when captum is not installed: use the model's own
    cross-attention output as a coarse word-level proxy is impossible
    (cross-attn is over meta tokens, not BERT tokens), so we return BERT
    self-attention from [CLS] -> tokens as a rough text attribution.
    """
    model.eval()
    text = str(row["full_text"])
    enc = tokenizer(text, max_length=cfg.MAX_LENGTH, truncation=True,
                    padding="max_length", return_tensors="pt")
    input_ids = enc["input_ids"].to(device)
    attention_mask = enc["attention_mask"].to(device)
    meta_vec = torch.from_numpy(
        meta_encoder.transform(pd.DataFrame([row]))
    ).float().to(device)

    with torch.no_grad():
        out = model(input_ids, attention_mask, meta_vec, output_attentions=True)

    pred_class = int(out["logits"][0, aspect_idx, :].argmax().item())
    tokens = tokenizer.convert_ids_to_tokens(input_ids[0].cpu().numpy().tolist())
    mask = attention_mask[0].cpu().numpy().astype(bool)

    bert_attn = out.get("bert_attentions")
    if bert_attn:
        last = bert_attn[-1][0]            # (heads, seq, seq)
        cls_attn = last[:, 0, :].mean(0).cpu().numpy()
    else:
        cls_attn = np.zeros_like(mask, dtype=np.float32)
    cls_attn = np.asarray(cls_attn, dtype=np.float32) * mask
    if not _has_visible_scores(cls_attn):
        cls_attn = _lexical_evidence_scores(tokens, aspect_idx) * mask
    return tokens, cls_attn, mask, pred_class


# ---------------------------------------------------------------------------
# Cross-Attention weights over meta tokens (per aspect, per example)
# ---------------------------------------------------------------------------

def _meta_token_names_from_output(out, attn):
    names = out.get("meta_token_names")
    if names is not None:
        return list(names)
    arr = np.asarray(attn)
    width = int(arr.shape[-1]) if arr.ndim else int(arr.size)
    if width == 3:
        return SEMANTIC_META_TOKEN_NAMES
    return [f"meta_chunk_{i+1}" for i in range(width)]


def get_meta_attention(model, tokenizer, meta_encoder, row, device):
    """Return display meta attention, predictions, names, and optional per-aspect attention."""
    text = str(row["full_text"])
    enc = tokenizer(text, max_length=cfg.MAX_LENGTH, truncation=True,
                    padding="max_length", return_tensors="pt")
    input_ids = enc["input_ids"].to(device)
    attention_mask = enc["attention_mask"].to(device)
    meta_vec = torch.from_numpy(
        meta_encoder.transform(pd.DataFrame([row]))
    ).float().to(device)
    with torch.no_grad():
        out = model(input_ids, attention_mask, meta_vec)

    raw_attn = out["meta_attn_weights"][0].cpu().numpy()
    token_names = _meta_token_names_from_output(out, raw_attn)
    if raw_attn.ndim == 1:
        display_attn = raw_attn
        aspect_attn = None
    else:
        aspect_attn = raw_attn
        if "global_meta_attn_weights" in out and out["global_meta_attn_weights"] is not None:
            display_attn = out["global_meta_attn_weights"][0].cpu().numpy()
        else:
            display_attn = raw_attn.mean(axis=0)

    preds = out["logits"][0].argmax(dim=-1).cpu().numpy()
    return display_attn, preds, token_names, aspect_attn

# ---------------------------------------------------------------------------
# HTML report builder (combines text IG + meta attention)
# ---------------------------------------------------------------------------

def _render_meta_bar_html(meta_attn: np.ndarray, token_names: Optional[List[str]] = None) -> str:
    """Render a tiny inline bar chart for one metadata-attention vector."""
    meta_attn = np.asarray(meta_attn, dtype=np.float64).reshape(-1)
    token_names = token_names or META_TOKEN_NAMES
    norm = _normalize(meta_attn)
    parts = ['<div style="display:flex;gap:6px;align-items:end;height:60px;'
             'margin-top:8px;">']
    for i, (n, raw) in enumerate(zip(norm, meta_attn)):
        name = token_names[i] if i < len(token_names) else f"meta_{i+1}"
        h = int(8 + 50 * float(n))
        parts.append(
            f'<div style="text-align:center;width:80px;">'
            f'<div style="background:#5b8def;height:{h}px;'
            f'border-radius:3px 3px 0 0;"></div>'
            f'<div style="font-size:11px;margin-top:2px;color:#666;">'
            f'{html_lib.escape(name)}<br/>{float(raw):.3f}</div></div>'
        )
    parts.append('</div>')
    return "".join(parts)


def _render_meta_matrix_html(meta_attn_by_aspect, token_names: Optional[List[str]] = None) -> str:
    """Render aspect-specific metadata attention as compact rows."""
    if meta_attn_by_aspect is None:
        return ""
    arr = np.asarray(meta_attn_by_aspect, dtype=np.float64)
    if arr.ndim != 2:
        return ""
    token_names = token_names or META_TOKEN_NAMES
    parts = ['<div style="margin-top:8px;display:grid;gap:6px;">']
    for i, aspect in enumerate(cfg.ASPECTS):
        if i >= arr.shape[0]:
            break
        weights = arr[i]
        norm = _normalize(weights)
        parts.append('<div style="display:flex;align-items:center;gap:8px;">')
        parts.append(f'<div style="width:95px;font-size:12px;font-weight:600;">{aspect}</div>')
        for j, (n, raw) in enumerate(zip(norm, weights)):
            name = token_names[j] if j < len(token_names) else f"meta_{j+1}"
            w = int(35 + 75 * float(n))
            parts.append(
                f'<div title="{html_lib.escape(name)}: {float(raw):.3f}" '
                f'style="background:#dbe7ff;border-left:4px solid #5b8def;'
                f'width:{w}px;padding:2px 4px;border-radius:4px;font-size:11px;">'
                f'{html_lib.escape(name)} {float(raw):.2f}</div>'
            )
        parts.append('</div>')
    parts.append('</div>')
    return "".join(parts)

def _render_meta_summary_html(meta_summary: Dict[str, str]) -> str:
    rows = []
    for name in ["features", "categories", "numeric"]:
        value = html_lib.escape(str(meta_summary.get(name, "")))
        rows.append(
            f'<div><b>{html_lib.escape(name)}:</b> '
            f'<span style="color:#444">{value or "not available"}</span></div>'
        )
    return '<div class="meta-box">' + "".join(rows) + '</div>'


def build_explanation_html(examples: List[dict], output_path: Path):
    parts = [
        "<!doctype html><html><head><meta charset='utf-8'>",
        "<title>ACSA + Meta Fusion Explanations</title>",
        "<style>",
        "body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; max-width: 1180px; margin: 24px auto; padding: 0 16px; color: #222; }",
        ".example { border: 1px solid #d8d8d8; padding: 16px; margin-bottom: 24px; border-radius: 8px; background:#fff; }",
        ".meta { font-size: 13px; color: #666; margin-bottom: 8px; }",
        ".meta-box { background:#f6f8fb; padding:10px 12px; border-radius:6px; font-size:13px; line-height:1.55; margin:10px 0 12px; }",
        ".review-box { background:#fafafa; padding:10px 12px; border-radius:6px; margin-bottom:12px; line-height:1.55; }",
        ".aspect-row { padding: 10px 0; border-bottom: 1px dashed #e6e6e6; line-height: 1.7; }",
        ".aspect-label { display: inline-block; min-width: 130px; font-weight: 700; }",
        ".pred-pos { color: #21834a; font-weight: 700; }",
        ".pred-neg { color: #bd2f24; font-weight: 700; }",
        ".pred-na  { color: #777; font-weight: 600; }",
        ".evidence { margin-top:6px; font-size:13px; color:#444; }",
        ".chip { display:inline-block; background:#eef2ff; border:1px solid #d6def8; padding:1px 6px; border-radius:999px; margin:1px 3px 1px 0; font-size:12px; }",
        "h1 { margin-bottom: 8px; } h2 { margin-top: 8px; }",
        ".legend { background: #f8f8f8; padding: 10px 12px; border-radius: 6px; font-size: 13px; line-height:1.5; }",
        ".section-h { font-weight:700;margin-top:16px;color:#333;font-size:14px; }",
        "</style></head><body>",
        "<h1>Aspect-Level Sentiment with Metadata Fusion: Explanations</h1>",
        '<div class="legend">'
        'Highlights show the strongest text evidence for each aspect. '
        '<span style="background:rgba(80,180,110,0.7);padding:2px 6px;border-radius:3px">green</span> = Positive, '
        '<span style="background:rgba(255,80,80,0.7);padding:2px 6px;border-radius:3px">red</span> = Negative, '
        '<span style="background:rgba(150,150,150,0.7);padding:2px 6px;border-radius:3px">grey</span> = Not_Mentioned. '
        'Metadata rows show which source the model attends to for each aspect.'
        '</div>',
    ]
    for i, ex in enumerate(examples):
        parts.append(f'<div class="example"><h2>Example {i+1}</h2>')
        parts.append(f'<div class="meta">Rating: {ex.get("rating", "?")} | Category: {html_lib.escape(str(ex.get("category", "?")))}</div>')
        parts.append(f'<div class="review-box">{html_lib.escape(ex["text"])}</div>')
        parts.append(_render_meta_summary_html(ex.get("meta_summary", {})))

        parts.append('<div class="section-h">Per-aspect prediction, text evidence, and metadata source</div>')
        token_names = ex.get("meta_token_names")
        meta_by_aspect = ex.get("meta_attn_by_aspect")
        for idx, a in enumerate(ex["aspects"]):
            pred = a["pred_label"]
            name = cfg.LABEL_NAMES[pred]
            css = {0: "pred-na", 1: "pred-pos", 2: "pred-neg"}[pred]
            parts.append('<div class="aspect-row">')
            parts.append(f'<span class="aspect-label">{a["aspect"]}:</span> <span class="{css}">{name}</span>')

            if "tokens" in a and "scores" in a:
                terms = a.get("top_terms") or []
                chips = "".join(f'<span class="chip">{html_lib.escape(t)}</span>' for t, _ in terms[:6])
                if chips:
                    parts.append(f'<div class="evidence"><b>Top text evidence:</b> {chips}</div>')
                tok_html = _render_token_html(a["tokens"], a["scores"], pred)
                parts.append(f'<div style="margin-top:6px">{tok_html}</div>')

            if meta_by_aspect is not None and idx < len(meta_by_aspect):
                source, weight = _top_meta_source(meta_by_aspect[idx], token_names)
            else:
                source, weight = _top_meta_source(ex.get("meta_attn", []), token_names)
            quality = a.get("evidence_quality", "weak")
            explanation = a.get("explanation_sentence", "")
            if source:
                parts.append(
                    f'<div class="evidence"><b>Top metadata source:</b> '
                    f'<span class="chip">{html_lib.escape(source)} {weight:.3f}</span> '
                    f'<span class="chip">quality: {html_lib.escape(quality)}</span></div>'
                )
            if explanation:
                parts.append(
                    f'<div class="evidence"><b>Explanation:</b> '
                    f'{html_lib.escape(explanation)}</div>'
                )
            parts.append('</div>')

        parts.append('<div class="section-h">Metadata cross-attention (global / averaged)</div>')
        parts.append(_render_meta_bar_html(np.asarray(ex["meta_attn"]), token_names))
        if meta_by_aspect is not None:
            parts.append('<div class="section-h">Aspect-specific metadata attention</div>')
            parts.append(_render_meta_matrix_html(np.asarray(meta_by_aspect), token_names))
        parts.append('</div>')

    parts.append('</body></html>')
    output_path = Path(output_path)
    output_path.write_text("\n".join(parts), encoding="utf-8")
    logger.info("Wrote explanation HTML to %s", output_path)


def write_explanation_summary(examples: List[dict], output_path: Path):
    rows = []
    for ex_idx, ex in enumerate(examples, start=1):
        token_names = ex.get("meta_token_names")
        meta_by_aspect = ex.get("meta_attn_by_aspect")
        for aspect_idx, aspect_data in enumerate(ex.get("aspects", [])):
            if meta_by_aspect is not None and aspect_idx < len(meta_by_aspect):
                source, weight = _top_meta_source(meta_by_aspect[aspect_idx], token_names)
            else:
                source, weight = _top_meta_source(ex.get("meta_attn", []), token_names)
            pred_name = cfg.LABEL_NAMES[int(aspect_data.get("pred_label", 0))]
            rows.append({
                "example": ex_idx,
                "rating": ex.get("rating"),
                "category": ex.get("category"),
                "aspect": aspect_data.get("aspect"),
                "pred_label": pred_name,
                "evidence_quality": aspect_data.get("evidence_quality", "weak"),
                "explanation_sentence": aspect_data.get("explanation_sentence", ""),
                "top_text_terms": "; ".join(t for t, _ in aspect_data.get("top_terms", [])[:8]),
                "top_meta_source": source,
                "top_meta_weight": weight,
                "features": ex.get("meta_summary", {}).get("features", ""),
                "categories": ex.get("meta_summary", {}).get("categories", ""),
                "numeric": ex.get("meta_summary", {}).get("numeric", ""),
                "text": ex.get("text", ""),
            })
    pd.DataFrame(rows).to_csv(output_path, index=False)
    logger.info("Wrote explanation summary to %s", output_path)

# ---------------------------------------------------------------------------
# High-level: pick examples, run all attributions, save HTML
# ---------------------------------------------------------------------------

def explain_examples(test_df, n_examples: int = 8, method: str = "ig",
                     output_path: Optional[Path] = None,
                     checkpoint_dir: Optional[Path] = None,
                     meta_encoder: Optional[MetaEncoder] = None):
    """Pick a mix of ratings from test_df and produce an explanation HTML."""
    if output_path is None:
        output_path = cfg.REPORT_DIR / f"explanation_{method}.html"

    model, tokenizer, meta_encoder, device = load_meta_acsa(checkpoint_dir, meta_encoder)

    # Sample across ratings to get diversity
    per_rating = max(1, n_examples // 5)
    chosen = []
    for r in [1, 2, 3, 4, 5]:
        sub = test_df[test_df["rating"] == r]
        if len(sub) > 0:
            chosen.append(sub.sample(n=min(per_rating, len(sub)),
                                     random_state=cfg.RANDOM_SEED))
    if not chosen:
        chosen = [test_df.sample(n=min(n_examples, len(test_df)),
                                 random_state=cfg.RANDOM_SEED)]
    selected = pd.concat(chosen).head(n_examples).reset_index(drop=True)

    try:
        import captum  # noqa: F401
        captum_ok = True
    except ImportError:
        logger.warning("captum not installed; falling back to BERT [CLS]-attention as proxy.")
        captum_ok = False
        method = "attention"

    examples_data = []
    for _, row in selected.iterrows():
        meta_attn, _, meta_token_names, meta_attn_by_aspect = get_meta_attention(model, tokenizer, meta_encoder, row, device)
        ex = {
            "text": str(row["full_text"]),
            "rating": int(row["rating"]),
            "category": str(row.get("leaf_category", "")),
            "meta_summary": _meta_summary_from_row(row),
            "meta_attn": meta_attn.tolist(),
            "meta_token_names": meta_token_names,
            "meta_attn_by_aspect": (meta_attn_by_aspect.tolist() if meta_attn_by_aspect is not None else None),
            "aspects": [],
        }
        for i, aspect in enumerate(cfg.ASPECTS):
            try:
                if method == "ig" and captum_ok:
                    tokens, attr, mask, pred = _ig_for_aspect(
                        model, tokenizer, meta_encoder, row, device, i,
                    )
                else:
                    tokens, attr, mask, pred = _attention_for_aspect(
                        model, tokenizer, meta_encoder, row, device, i,
                    )
                scores = np.abs(attr) * mask
                display_scores = _normalize_for_display(scores)
                top_terms = _top_terms(tokens, display_scores)
                if meta_attn_by_aspect is not None and i < len(meta_attn_by_aspect):
                    source, weight = _top_meta_source(meta_attn_by_aspect[i], meta_token_names)
                else:
                    source, weight = _top_meta_source(meta_attn, meta_token_names)
                pred_name = cfg.LABEL_NAMES[int(pred)]
                quality = _explanation_quality(top_terms, source, weight)
                ex["aspects"].append({
                    "aspect": aspect, "pred_label": int(pred),
                    "tokens": tokens, "scores": scores,
                    "top_terms": top_terms,
                    "evidence_quality": quality,
                    "explanation_sentence": _make_explanation_sentence(
                        aspect, pred_name, top_terms, source, weight,
                    ),
                })
            except Exception as e:
                logger.warning("Attribution failed for aspect %s: %s", aspect, e)
                # Still record the prediction even without scores
                ex["aspects"].append({"aspect": aspect, "pred_label": 0})
        examples_data.append(ex)

    build_explanation_html(examples_data, output_path)
    summary_path = Path(output_path).with_name(Path(output_path).stem + "_summary.csv")
    write_explanation_summary(examples_data, summary_path)
    return examples_data


# ---------------------------------------------------------------------------
# Aggregation plots (application-layer)
# ---------------------------------------------------------------------------

def plot_aspect_distribution(df_with_preds: pd.DataFrame,
                             output_path: Optional[Path] = None):
    """Bar chart of Positive/Negative share per aspect (over mentioned rows)."""
    if output_path is None:
        output_path = cfg.REPORT_DIR / "aspect_distribution.png"
    rows = []
    for a in cfg.ASPECTS:
        col = f"pred_{a}" if f"pred_{a}" in df_with_preds.columns else f"aspect_{a}"
        if col not in df_with_preds.columns:
            continue
        vc = df_with_preds[col].value_counts()
        n_pos = int(vc.get(1, 0)); n_neg = int(vc.get(2, 0))
        n_mentioned = n_pos + n_neg
        if n_mentioned == 0:
            continue
        rows.append({"aspect": a, "positive": n_pos / n_mentioned,
                     "negative": n_neg / n_mentioned, "n_mentioned": n_mentioned})
    if not rows:
        logger.warning("No data to plot for aspect distribution.")
        return None

    dfp = pd.DataFrame(rows)
    fig, ax = plt.subplots(figsize=(10, 5))
    x = np.arange(len(dfp)); w = 0.35
    ax.bar(x - w/2, dfp["positive"], w, label="Positive share", color="#2a9d4f")
    ax.bar(x + w/2, dfp["negative"], w, label="Negative share", color="#c0392b")
    ax.set_xticks(x); ax.set_xticklabels(dfp["aspect"], rotation=20)
    ax.set_ylabel("Share among mentioned reviews")
    ax.set_title("Aspect-level sentiment distribution")
    ax.legend()
    for i, r in dfp.iterrows():
        ax.text(i, max(r["positive"], r["negative"]) + 0.02,
                f"n={r['n_mentioned']}", ha="center", fontsize=9)
    plt.tight_layout(); plt.savefig(output_path, dpi=120); plt.close()
    logger.info("Saved aspect distribution plot to %s", output_path)
    return dfp


def plot_category_aspect_heatmap(agg_df: pd.DataFrame, metric: str = "negative_share",
                                 output_path: Optional[Path] = None):
    if agg_df is None or agg_df.empty:
        return None
    if output_path is None:
        output_path = cfg.REPORT_DIR / f"category_aspect_{metric}.png"
    pivot = agg_df.pivot(index="category", columns="aspect", values=metric)
    fig, ax = plt.subplots(figsize=(10, max(4, len(pivot) * 0.4)))
    cmap = "RdYlGn_r" if "negative" in metric else "RdYlGn"
    sns.heatmap(pivot, annot=True, fmt=".2f", cmap=cmap, ax=ax,
                cbar_kws={"label": metric})
    ax.set_title(f"{metric.replace('_', ' ').title()} by category 闂?aspect")
    plt.tight_layout(); plt.savefig(output_path, dpi=120); plt.close()
    logger.info("Saved heatmap to %s", output_path)
    return pivot