"""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

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)]


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

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


def _color_for_score(score: float, label: int) -> str:
    score = float(np.clip(score, 0.0, 1.0))
    if label == 2:    # Negative
        return f"rgba(255, 80, 80, {score:.3f})"
    if label == 1:    # Positive
        return f"rgba(80, 200, 120, {score:.3f})"
    return f"rgba(180, 180, 180, {score:.3f})"  # Not_Mentioned


def _render_token_html(tokens, scores, label):
    norm = _normalize(scores)
    spans = []
    for tok, s in zip(tokens, norm):
        clean = tok.lstrip("Ġ").replace("##", "")
        if not clean or (clean.startswith("[") and clean.endswith("]")):
            continue
        color = _color_for_score(s, label)
        spans.append(
            f'<span style="background-color:{color};padding:1px 3px;'
            f'border-radius:3px;margin:1px;">{html_lib.escape(clean)}</span>'
        )
    return " ".join(spans)


# ---------------------------------------------------------------------------
# Integrated Gradients (per aspect)
# ---------------------------------------------------------------------------

def _ig_for_aspect(model, tokenizer, meta_encoder, row, device,
                   aspect_idx: int, n_steps: int = 30):
    """Integrated Gradients on the BERT input embedding for one aspect head."""
    from captum.attr import LayerIntegratedGradients

    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)

    # encode the row's meta features
    meta_vec = torch.from_numpy(
        meta_encoder.transform(pd.DataFrame([row]))
    ).float().to(device)

    def forward_for_aspect(ids, mask):
        out = model(ids, mask, meta_vec)
        return out["logits"][:, aspect_idx, :]

    embed_layer = model.bert.embeddings
    lig = LayerIntegratedGradients(forward_for_aspect, embed_layer)

    with torch.no_grad():
        logits = forward_for_aspect(input_ids, attention_mask)
        target_class = int(logits.argmax(dim=-1).item())

    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
    baseline = torch.full_like(input_ids, pad_id)

    attributions = lig.attribute(
        inputs=input_ids, baselines=baseline,
        additional_forward_args=(attention_mask,),
        target=target_class, n_steps=n_steps, internal_batch_size=4,
    )
    attributions = attributions.sum(dim=-1).squeeze(0).cpu().numpy()
    tokens = tokenizer.convert_ids_to_tokens(input_ids[0].cpu().numpy().tolist())
    mask = attention_mask[0].cpu().numpy().astype(bool)
    return tokens, attributions, 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)
    return tokens, cls_attn, mask, pred_class


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

def get_meta_attention(model, tokenizer, meta_encoder, row, device):
    """Returns a vector (num_meta_tokens,) of attention over meta chunks for
    this row, plus the predicted labels for each aspect.

    NB: in the current model the cross-attention happens BEFORE the per-aspect
    heads, so meta-attn is shared across aspects for a given example. We still
    return per-aspect preds for the report.
    """
    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)
    meta_attn = out["meta_attn_weights"][0].cpu().numpy()      # (T,)
    preds = out["logits"][0].argmax(dim=-1).cpu().numpy()       # (num_aspects,)
    return meta_attn, preds


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

def _render_meta_bar_html(meta_attn: np.ndarray) -> str:
    """Render a tiny inline bar chart for the meta-attention vector."""
    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)):
        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'{META_TOKEN_NAMES[i]}<br/>{raw:.3f}</div></div>'
        )
    parts.append('</div>')
    return "".join(parts)


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, sans-serif; max-width: 1100px; margin: 24px auto;"
        " padding: 0 16px; color: #222; }",
        ".example { border: 1px solid #ddd; padding: 16px; margin-bottom: 24px; border-radius: 8px; }",
        ".meta { font-size: 13px; color: #666; margin-bottom: 8px; }",
        ".aspect-row { padding: 8px 0; border-bottom: 1px dashed #eee; line-height: 1.7; }",
        ".aspect-label { display: inline-block; min-width: 140px; font-weight: 600; }",
        ".pred-pos { color: #2a9d4f; font-weight: 600; }",
        ".pred-neg { color: #c0392b; font-weight: 600; }",
        ".pred-na  { color: #888; }",
        "h2 { margin-top: 24px; }",
        ".legend { background: #f8f8f8; padding: 8px 12px; border-radius: 6px; font-size: 13px; }",
        ".section-h { font-weight:600;margin-top:14px;color:#444;font-size:14px; }",
        "</style></head><body>",
        "<h1>Aspect-Level Sentiment with Metadata Fusion — Explanations</h1>",
        '<div class="legend">'
        'Color intensity = attribution magnitude. '
        '<span style="background:rgba(80,200,120,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(180,180,180,0.7);padding:2px 6px;border-radius:3px">grey</span> = Not_Mentioned. '
        'Bottom bar chart shows the model\'s attention over metadata chunks for this review.'
        '</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", "?")} | '
                     f'Category: {ex.get("category", "?")}</div>')
        parts.append(
            f'<div style="background:#fafafa;padding:8px;border-radius:4px;'
            f'margin-bottom:12px">{html_lib.escape(ex["text"])}</div>'
        )

        # Per-aspect rows with token highlights
        parts.append('<div class="section-h">Per-aspect prediction &amp; token attribution</div>')
        for a in 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> '
                         f'<span class="{css}">{name}</span>')
            if pred != 0 and "tokens" in a and "scores" in a:
                tok_html = _render_token_html(a["tokens"], a["scores"], pred)
                parts.append(f'<div style="margin-top:6px">{tok_html}</div>')
            parts.append('</div>')

        # Metadata cross-attention bar
        parts.append('<div class="section-h">Metadata cross-attention (shared across aspects)</div>')
        parts.append(_render_meta_bar_html(np.asarray(ex["meta_attn"])))
        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)


# ---------------------------------------------------------------------------
# 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, _ = 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_attn": meta_attn.tolist(),
            "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
                ex["aspects"].append({
                    "aspect": aspect, "pred_label": int(pred),
                    "tokens": tokens, "scores": scores,
                })
            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)
    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