#!/usr/bin/env python3
"""
gemma4-e4b Model Card Graph Generator

Generates all SVG visualisations for the gemma4-e4b abliteration comparison
with 23 variants. Reads from analysis results directory + SQLite DB for
HarmBench LLM-judged ASR.

Graph types:
  1. Benchmark Comparison (grouped bar, 6 tasks)
  2. Benchmark Delta (horizontal bar, deltas vs base)
  3. GSM8K Comparison (horizontal bar)
  4. HarmBench ASR Summary (bar chart, LLM-judged)
  5. HarmBench ASR by Category (grouped bar, LLM-judged)
  6. KL Divergence (horizontal bar with ratings)
  7. Aggressiveness (tensors changed)
  8. Cosine Heatmap (23x23 cross-technique)
  9. Layer Edit Magnitude (line plot, all variants)
  10. Edit Distribution (violin plot)
  11. Tensor Type Breakdown (grouped bar)

Usage:
    python3 generate_e4b_card_graphs.py \
        --results-dir comparisons/gemma4-e4b/results \
        --output-dir comparisons/gemma4-e4b/graphs \
        --db-path abliterlitics.db
"""

from __future__ import annotations

import json
import sqlite3
import sys
from collections import defaultdict
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import seaborn as sns

sns.set_theme(style="whitegrid", palette="muted", font_scale=1.1)

# ---------- Variant definitions ----------

VARIANTS = [
    "heretic-std",
    "coder3101",
    "heresy",
    "heretic",
    "mythos",
    "treadon",
    "wwt",
    "nullpo",
    "treadon-disin",
    "treadon-combo",
    "huihui",
    "infinimind",
    "trevorjs",
    "abliterix",
    "apostate",
    "distill",
    "deckard",
    "deckard-expresso",
    "claude-distill",
    "bendernina",
    "physshell",
    "sdft",
    "obliteratus",
]

VARIANT_LABELS = {v: v.replace("-", " ").title() for v in VARIANTS}
VARIANT_LABELS["heretic-std"] = "Heretic Std"
VARIANT_LABELS["deckard-expresso"] = "Deckard Espresso"
VARIANT_LABELS["treadon-disin"] = "Treadon Disin"
VARIANT_LABELS["treadon-combo"] = "Treadon Combo"
VARIANT_LABELS["claude-distill"] = "Claude Distill"
VARIANT_LABELS["coder3101"] = "Coder3101"
VARIANT_LABELS["trevorjs"] = "TrevorJS"

# Colour palette: husl gives perceptually distinct colours
_palette = sns.color_palette("husl", len(VARIANTS))
VARIANT_COLORS = {v: f"#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}" for v, (r, g, b) in zip(VARIANTS, _palette)}

BASE_NAME = "gemma-4-E4B-it"
BASE_COLOR = "#95a5a6"
LABEL = "gemma4-e4b"

# HarmBench categories
HB_CATEGORIES = [
    "chemical_biological",
    "copyright",
    "cybercrime_intrusion",
    "harassment_bullying",
    "harmful",
    "illegal",
    "misinformation_disinformation",
]
HB_CAT_SHORT = {
    "chemical_biological": "Chem/Bio",
    "copyright": "Copyright",
    "cybercrime_intrusion": "Cybercrime",
    "harassment_bullying": "Harassment",
    "harmful": "Harmful",
    "illegal": "Illegal",
    "misinformation_disinformation": "Misinfo",
}

# Benchmark tasks (loglikelihood). Task keys MUST match the lm-eval result JSON
# `results` dict exactly. MMLU migrated from the v1 `mmlu` task to the OLL v2
# `leaderboard_mmlu_pro` task when this comparison moved to the v2 suite
# (AGENTS.md "Benchmark Tasks"). The other entries are v1 forensic supplements
# that are still run alongside the v2 suite.
LOGIT_TASKS = [
    ("MMLU-Pro", "leaderboard_mmlu_pro", "acc,none"),
    ("HellaSwag", "hellaswag", "acc_norm,none"),
    ("ARC-C", "arc_challenge", "acc,none"),
    ("WinoGrande", "winogrande", "acc,none"),
    ("TQA-MC2", "truthfulqa_mc2", "acc,none"),
    ("PiQA", "piqa", "acc_norm,none"),
]


# ---------- Helpers ----------


def load_json(path: Path, quiet: bool = False) -> dict | None:
    try:
        return json.loads(path.read_text())
    except FileNotFoundError:
        if not quiet:
            print(f"  [WARN] File not found: {path}")
        return None
    except json.JSONDecodeError as e:
        print(f"  [WARN] JSON decode error in {path}: {e}")
        return None


def save_fig(fig: plt.Figure, out_dir: Path, name: str) -> None:
    path = out_dir / name
    fig.savefig(path, format="svg", bbox_inches="tight", dpi=150)
    plt.close(fig)
    print(f"  [OK] {name}")


def load_harmbench_from_db(db_path: Path) -> dict[str, dict]:
    """Load LLM-judged ASR from DB. Returns {model_name: {'overall': asr, 'refused': n, 'categories': {cat: asr}}}"""
    conn = sqlite3.connect(str(db_path))
    c = conn.cursor()

    result: dict[str, dict] = {}

    # Overall ASR
    rows = c.execute("""
        SELECT model_name,
               COUNT(*) as total,
               SUM(CASE WHEN verdict = 'complied' THEN 1 ELSE 0 END) as complied,
               SUM(CASE WHEN verdict = 'refused' THEN 1 ELSE 0 END) as refused
        FROM llm_reviews
        WHERE comparison LIKE '%gemma4-e4b%'
        GROUP BY model_name
    """).fetchall()

    for name, total, complied, refused in rows:
        result[name] = {
            "overall": complied / total * 100 if total else 0,
            "complied": complied,
            "refused": refused,
            "total": total,
            "categories": {},
        }

    # Per-category ASR
    rows = c.execute("""
        SELECT lr.model_name, b.semantic_category,
               COUNT(*) as total,
               SUM(CASE WHEN lr.verdict = 'complied' THEN 1 ELSE 0 END) as complied
        FROM llm_reviews lr
        JOIN behaviors b ON lr.behavior_id = b.behavior_id
        WHERE lr.comparison LIKE '%gemma4-e4b%'
        GROUP BY lr.model_name, b.semantic_category
    """).fetchall()

    for name, cat, total, complied in rows:
        if name in result:
            result[name]["categories"][cat] = complied / total * 100 if total else 0

    conn.close()
    return result


def get_tensor_counts(results_dir: Path) -> dict[str, int]:
    """Get tensor changed count per variant from panel + fingerprints."""
    counts = {}
    panel = load_json(results_dir / "multi_model_panel.json", quiet=True)
    if panel:
        pcc = panel.get("pairwise_changed_counts", {})
        for v in VARIANTS:
            key = f"base->{v}"
            if key in pcc:
                counts[v] = pcc[key]

    # Fill missing from fingerprints
    for v in VARIANTS:
        if v not in counts:
            fp = load_json(results_dir / v / f"fingerprint_{v}.json", quiet=True)
            if fp:
                n = fp.get("scope", {}).get("changed_tensors", 0)
                if n:
                    counts[v] = n

    return counts


# ---------- Graph 1: Benchmark Comparison (grouped bar) ----------


def gen_benchmark_comparison(results_dir: Path, out_dir: Path) -> None:
    """Grouped bar chart of Phase 1 benchmarks across all 24 models."""
    all_models = [BASE_NAME] + VARIANTS
    scores: dict[str, list[float]] = {v: [] for v in all_models}
    task_names = [t[0] for t in LOGIT_TASKS]

    for _tname, task_key, metric in LOGIT_TASKS:
        for v in all_models:
            path = results_dir / "lm_eval" / f"lm_eval_{v}.json"
            d = load_json(path, quiet=True)
            val = 0.0
            if d and "results" in d:
                try:
                    val = d["results"][task_key][metric] * 100
                except (KeyError, TypeError):
                    pass
            scores[v].append(val)

    n_tasks = len(task_names)
    n_vars = len(all_models)
    x = np.arange(n_tasks)
    width = 0.85 / n_vars
    offsets = np.arange(n_vars) - (n_vars - 1) / 2

    fig, ax = plt.subplots(figsize=(24, 9))
    for i, v in enumerate(all_models):
        label = "Base" if v == BASE_NAME else VARIANT_LABELS.get(v, v)
        color = BASE_COLOR if v == BASE_NAME else VARIANT_COLORS.get(v, "#888")
        ax.bar(x + offsets[i] * width, scores[v], width,
               label=label, color=color, alpha=0.85, edgecolor="white", linewidth=0.3)

    ax.set_xticks(x)
    ax.set_xticklabels(task_names, fontsize=11)
    ax.set_ylabel("Score (%)")
    ax.set_ylim(0, 65)
    ax.legend(loc="upper right", fontsize=6, ncol=4)
    ax.set_title(f"{LABEL} Benchmark Comparison (Loglikelihood Tasks)",
                 fontsize=14, fontweight="bold")
    save_fig(fig, out_dir, f"{LABEL}_benchmark_comparison.svg")


# ---------- Graph 2: Benchmark Delta (horizontal bar) ----------


def gen_benchmark_delta(results_dir: Path, out_dir: Path) -> None:
    """Horizontal bar chart of deltas vs base for each variant."""
    base_data = load_json(results_dir / "lm_eval" / f"lm_eval_{BASE_NAME}.json")
    if not base_data:
        print("  [SKIP] benchmark_delta: no base data")
        return

    base_scores = {}
    for tname, task_key, metric in LOGIT_TASKS:
        try:
            base_scores[tname] = base_data["results"][task_key][metric] * 100
        except (KeyError, TypeError):
            base_scores[tname] = 0

    delta_data: dict[str, dict[str, float]] = {}
    for v in VARIANTS:
        d = load_json(results_dir / "lm_eval" / f"lm_eval_{v}.json", quiet=True)
        if not d or "results" not in d:
            continue
        delta_data[v] = {}
        for tname, task_key, metric in LOGIT_TASKS:
            try:
                variant_score = d["results"][task_key][metric] * 100
                delta_data[v][tname] = variant_score - base_scores[tname]
            except (KeyError, TypeError):
                delta_data[v][tname] = 0

    if not delta_data:
        print("  [SKIP] benchmark_delta: no variant data")
        return

    task_names = [t[0] for t in LOGIT_TASKS]
    n_tasks = len(task_names)
    n_vars = len(delta_data)
    variant_order = [v for v in VARIANTS if v in delta_data]

    fig, ax = plt.subplots(figsize=(14, 16))
    y = np.arange(n_tasks)
    height = 0.85 / n_vars
    offsets = np.arange(n_vars) - (n_vars - 1) / 2

    for i, v in enumerate(variant_order):
        vals = [delta_data[v].get(t, 0) for t in task_names]
        label = VARIANT_LABELS.get(v, v)
        color = VARIANT_COLORS.get(v, "#888")
        ax.barh(y + offsets[i] * height, vals, height,
                label=label, color=color, alpha=0.85, edgecolor="white", linewidth=0.3)

    ax.set_yticks(y)
    ax.set_yticklabels(task_names, fontsize=11)
    ax.set_xlabel("Delta vs Base (percentage points)")
    ax.axvline(x=0, color="black", linewidth=0.8)
    ax.legend(loc="lower right", fontsize=6, ncol=3)
    ax.set_title(f"{LABEL} Benchmark Delta vs Base",
                 fontsize=14, fontweight="bold")
    save_fig(fig, out_dir, f"{LABEL}_benchmark_delta.svg")


# ---------- Graph 3: GSM8K Comparison ----------


def gen_gsm8k_comparison(results_dir: Path, out_dir: Path) -> None:
    """Horizontal bar chart of GSM8K scores."""
    all_models = [BASE_NAME] + VARIANTS
    data: list[tuple[str, float]] = []

    for v in all_models:
        path = results_dir / "lm_eval" / f"lm_eval_{v}.json"
        d = load_json(path, quiet=True)
        if not d or "results" not in d:
            continue
        try:
            score = d["results"]["gsm8k"]["exact_match,flexible-extract"] * 100
            data.append((v, score))
        except (KeyError, TypeError):
            continue

    if len(data) < 2:
        print("  [SKIP] gsm8k_comparison: not enough data")
        return

    data.sort(key=lambda x: x[1], reverse=True)

    labels = ["Base" if v == BASE_NAME else VARIANT_LABELS.get(v, v) for v, _ in data]
    values = [s for _, s, in data]
    colors = [BASE_COLOR if v == BASE_NAME else VARIANT_COLORS.get(v, "#888") for v, _ in data]

    fig, ax = plt.subplots(figsize=(12, 10))
    bars = ax.barh(range(len(labels)), values, color=colors, alpha=0.85, edgecolor="white")

    for i, (bar, val) in enumerate(zip(bars, values)):
        ax.text(bar.get_width() + 0.3, bar.get_y() + bar.get_height() / 2,
                f"{val:.1f}%", va="center", fontsize=8)

    ax.set_yticks(range(len(labels)))
    ax.set_yticklabels(labels, fontsize=8)
    ax.set_xlabel("GSM8K Score (%)")
    ax.invert_yaxis()
    ax.set_title(f"{LABEL} GSM8K with Thinking Enabled",
                 fontsize=14, fontweight="bold")
    save_fig(fig, out_dir, f"{LABEL}_gsm8k_comparison.svg")


# ---------- Graph 4: HarmBench ASR Summary (from DB) ----------


def gen_harmbench_summary(hb_data: dict, out_dir: Path) -> None:
    """Bar chart of overall HarmBench ASR for all 24 models."""
    all_models = [BASE_NAME] + VARIANTS
    asr_list: list[tuple[str, float, int]] = []

    for v in all_models:
        if v in hb_data:
            d = hb_data[v]
            asr_list.append((v, d["overall"], d["refused"]))

    if len(asr_list) < 2:
        print("  [SKIP] harmbench_summary: not enough data")
        return

    # Sort by ASR descending
    asr_list.sort(key=lambda x: x[1], reverse=True)

    labels = ["Base" if v == BASE_NAME else VARIANT_LABELS.get(v, v) for v, _, _ in asr_list]
    values = [a for _, a, _ in asr_list]
    colors = [BASE_COLOR if v == BASE_NAME else VARIANT_COLORS.get(v, "#888") for v, _, _ in asr_list]

    fig, ax = plt.subplots(figsize=(16, 8))
    bars = ax.bar(labels, values, color=colors, alpha=0.9, edgecolor="white")

    for bar, val, ref in zip(bars, values, [r for _, _, r in asr_list]):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 1,
                f"{val:.1f}%\n({ref} ref)", ha="center", va="bottom",
                fontsize=7, fontweight="bold")

    ax.set_ylabel("Attack Success Rate (%)")
    ax.set_ylim(0, 115)
    ax.axhline(y=100, color="#e74c3c", linestyle="--", alpha=0.3)
    plt.xticks(rotation=45, ha="right", fontsize=8)
    ax.set_title(f"{LABEL} HarmBench Overall ASR (LLM-Judged)",
                 fontsize=14, fontweight="bold")
    fig.tight_layout()
    save_fig(fig, out_dir, f"{LABEL}_harmbench_summary.svg")


# ---------- Graph 5: HarmBench ASR by Category (from DB) ----------


def gen_harmbench_asr_category(hb_data: dict, out_dir: Path) -> None:
    """Grouped bar: ASR by category for all 24 models."""
    all_models = [BASE_NAME] + VARIANTS

    cats_present = [c for c in HB_CATEGORIES
                    if any(c in hb_data.get(v, {}).get("categories", {}) for v in all_models)]
    if not cats_present:
        print("  [SKIP] harmbench_asr_category: no category data")
        return

    cat_labels = [HB_CAT_SHORT[c] for c in cats_present]
    n_cats = len(cats_present)
    n_vars = len(all_models)
    x = np.arange(n_cats)
    width = 0.85 / n_vars
    offsets = np.arange(n_vars) - (n_vars - 1) / 2

    fig, ax = plt.subplots(figsize=(24, 9))
    for i, v in enumerate(all_models):
        cats = hb_data.get(v, {}).get("categories", {})
        vals = [cats.get(c, 0) for c in cats_present]
        label = "Base" if v == BASE_NAME else VARIANT_LABELS.get(v, v)
        color = BASE_COLOR if v == BASE_NAME else VARIANT_COLORS.get(v, "#888")
        ax.bar(x + offsets[i] * width, vals, width,
               label=label, color=color, alpha=0.85, edgecolor="white", linewidth=0.3)

    ax.set_xticks(x)
    ax.set_xticklabels(cat_labels, fontsize=10)
    ax.set_ylabel("Attack Success Rate (%)")
    ax.set_ylim(0, 115)
    ax.axhline(y=100, color="#e74c3c", linestyle="--", alpha=0.3)
    ax.legend(fontsize=5, ncol=6, loc="upper left")
    ax.set_title(f"{LABEL} HarmBench ASR by Category (LLM-Judged)",
                 fontsize=14, fontweight="bold")
    save_fig(fig, out_dir, f"{LABEL}_harmbench_asr.svg")


# ---------- Graph 6: KL Divergence ----------


def gen_kl_divergence(results_dir: Path, out_dir: Path) -> None:
    """Horizontal bar chart of KL divergence with colour-coded ratings."""
    kl_data: list[tuple[str, float, str]] = []

    for v in VARIANTS:
        path = results_dir / "kl" / f"kl_{v}.json"
        d = load_json(path, quiet=True)
        if d and "kl_divergence_batchmean" in d:
            kl = d["kl_divergence_batchmean"]
            if kl < 0.01:
                rating = "excellent"
            elif kl < 0.1:
                rating = "very good"
            elif kl < 0.4:
                rating = "moderate"
            elif kl < 1.0:
                rating = "significant"
            else:
                rating = "heavy"
            kl_data.append((v, kl, rating))

    if not kl_data:
        print("  [SKIP] kl_divergence: no data")
        return

    kl_data.sort(key=lambda x: x[1])

    labels = [f"{VARIANT_LABELS.get(v, v)}\n({r})" for v, _, r in kl_data]
    values = [kl for _, kl, _ in kl_data]

    rating_colors = {
        "excellent": "#2ecc71",
        "very good": "#27ae60",
        "moderate": "#3498db",
        "significant": "#f39c12",
        "heavy": "#e74c3c",
    }
    colors = [rating_colors.get(r, "#888") for _, _, r in kl_data]

    fig, ax = plt.subplots(figsize=(10, 12))
    bars = ax.barh(range(len(labels)), values, color=colors, alpha=0.9, edgecolor="white")

    for bar, val in zip(bars, values):
        ax.text(bar.get_width() + max(values) * 0.02, bar.get_y() + bar.get_height() / 2,
                f"{val:.4f}", va="center", fontsize=8, fontweight="bold")

    ax.set_yticks(range(len(labels)))
    ax.set_yticklabels(labels, fontsize=8)
    ax.set_xlabel("KL Divergence (batchmean)")
    ax.invert_yaxis()

    if max(values) > 0.01:
        ax.axvspan(0, 0.01, alpha=0.05, color="green")
        ax.axvspan(0.01, 0.1, alpha=0.03, color="lime")
        if max(values) > 0.1:
            ax.axvspan(0.1, 0.4, alpha=0.03, color="blue")
            if max(values) > 0.4:
                ax.axvspan(0.4, 1.0, alpha=0.03, color="orange")
                if max(values) > 1.0:
                    ax.axvspan(1.0, max(values) * 1.1, alpha=0.03, color="red")

    ax.set_title(f"{LABEL} KL Divergence from Base",
                 fontsize=14, fontweight="bold")
    fig.tight_layout()
    save_fig(fig, out_dir, f"{LABEL}_kl_divergence.svg")


# ---------- Graph 7: Aggressiveness ----------


def gen_aggressiveness(results_dir: Path, out_dir: Path) -> None:
    """Bar chart: tensors changed per variant."""
    counts = get_tensor_counts(results_dir)

    if not counts:
        print("  [SKIP] aggressiveness: no data")
        return

    # Sort by count ascending
    sorted_items = sorted(counts.items(), key=lambda x: x[1])

    labels = [VARIANT_LABELS.get(v, v) for v, _ in sorted_items]
    values = [c for _, c in sorted_items]
    colors = [VARIANT_COLORS.get(v, "#888") for v, _ in sorted_items]

    fig, ax = plt.subplots(figsize=(10, 12))
    bars = ax.barh(range(len(labels)), values, color=colors, alpha=0.9, edgecolor="white")

    for bar, val in zip(bars, values):
        ax.text(bar.get_width() + max(values) * 0.02, bar.get_y() + bar.get_height() / 2,
                str(val), va="center", fontsize=9, fontweight="bold")

    ax.set_yticks(range(len(labels)))
    ax.set_yticklabels(labels, fontsize=8)
    ax.set_xlabel("Tensors Changed")
    ax.invert_yaxis()
    ax.set_title(f"{LABEL} Abliteration Aggressiveness",
                 fontsize=14, fontweight="bold")
    fig.tight_layout()
    save_fig(fig, out_dir, f"{LABEL}_aggressiveness.svg")


# ---------- Graph 8: Cosine Heatmap ----------


def gen_cosine_heatmap(results_dir: Path, out_dir: Path) -> None:
    """23x23 heatmap of cross-technique cosine similarities."""
    cosine_map: dict[tuple[str, str], float] = {}
    for corr_file in sorted(results_dir.glob("correlation_*_vs_*.json")):
        d = load_json(corr_file, quiet=True)
        if not d:
            continue
        stem = corr_file.stem.replace("correlation_", "")
        parts = stem.split("_vs_")
        if len(parts) != 2:
            continue
        a, b = parts[0], parts[1]
        pc = d.get("pairwise_cosines", {})
        for _key, pdata in pc.items():
            cosine_map[(a, b)] = pdata.get("mean", 0)
            cosine_map[(b, a)] = pdata.get("mean", 0)
            break

    if not cosine_map:
        print("  [SKIP] cosine_heatmap: no correlation data")
        return

    n = len(VARIANTS)
    matrix = np.zeros((n, n))
    for i in range(n):
        matrix[i, i] = 1.0
        for j in range(n):
            if i != j:
                matrix[i, j] = cosine_map.get((VARIANTS[i], VARIANTS[j]), 0.0)

    labels = [VARIANT_LABELS.get(v, v) for v in VARIANTS]

    fig, ax = plt.subplots(figsize=(18, 16))
    sns.heatmap(matrix, ax=ax, annot=True, fmt=".2f", cmap="RdYlGn",
                vmin=0, vmax=1.0, center=0.5,
                xticklabels=labels, yticklabels=labels,
                linewidths=0.5, linecolor="white",
                annot_kws={"size": 5},
                cbar_kws={"label": "Mean Cosine Similarity"})
    ax.set_title(f"{LABEL} Cross-Technique Edit Vector Cosine Similarity",
                 fontsize=14, fontweight="bold")
    plt.xticks(rotation=45, ha="right", fontsize=7)
    plt.yticks(fontsize=7)
    fig.tight_layout()
    save_fig(fig, out_dir, f"{LABEL}_cosine_heatmap.svg")


# ---------- Graph 9: Layer Edit Magnitude ----------


def gen_layer_comparison(results_dir: Path, out_dir: Path) -> None:
    """Line plot: mean edit norm by layer for all variants."""
    fig, ax = plt.subplots(figsize=(18, 9))

    has_data = False
    all_layers = None
    for v in VARIANTS:
        d = load_json(results_dir / v / f"layer_analysis_{v}.json", quiet=True)
        if not d or "layer_progression" not in d:
            continue
        lp = d["layer_progression"]
        layers = sorted(lp.keys(), key=lambda k: int(k))
        if all_layers is None:
            all_layers = layers
        edit_norms = [lp[l].get("mean_edit_norm", 0) for l in layers]

        label = VARIANT_LABELS.get(v, v)
        color = VARIANT_COLORS.get(v, "#888")
        ax.plot(range(len(layers)), edit_norms,
                label=label, color=color, alpha=0.7, linewidth=1.2)
        has_data = True

    if not has_data:
        print("  [SKIP] layer_comparison: no data")
        plt.close(fig)
        return

    ax.set_ylabel("Mean Edit Norm", fontsize=12)
    ax.set_xlabel("Layer", fontsize=12)
    ax.legend(fontsize=6, ncol=4, loc="upper left")
    ax.set_title(f"{LABEL} Layer-wise Edit Magnitude",
                 fontsize=14, fontweight="bold")

    if all_layers:
        tick_step = max(1, len(all_layers) // 16)
        ax.set_xticks(range(0, len(all_layers), tick_step))
        ax.set_xticklabels([f"L{all_layers[i]}" for i in range(0, len(all_layers), tick_step)])

    save_fig(fig, out_dir, f"{LABEL}_layer_comparison.svg")


# ---------- Graph 10: Edit Distribution ----------


def gen_edit_distribution(results_dir: Path, out_dir: Path) -> None:
    """Violin plot of per-tensor edit norms per variant."""
    all_norms: dict[str, list[float]] = {}

    for v in VARIANTS:
        svd = load_json(results_dir / v / f"svd_{v}.json", quiet=True)
        if not svd or "tensor_results" not in svd:
            continue
        norms = []
        for r in svd["tensor_results"]:
            for k, val in r.items():
                if isinstance(val, dict) and "frobenius_norm" in val:
                    n = val["frobenius_norm"]
                    if n > 0.01:
                        norms.append(n)
        if norms:
            all_norms[v] = norms

    if not all_norms:
        print("  [SKIP] edit_distribution: no per-tensor data")
        return

    variant_order = sorted(all_norms.keys(), key=lambda v: np.median(all_norms[v]))
    plot_data = [all_norms[v] for v in variant_order]
    plot_labels = [VARIANT_LABELS.get(v, v) for v in variant_order]
    plot_colors = [VARIANT_COLORS.get(v, "#888") for v in variant_order]

    fig, ax = plt.subplots(figsize=(16, 9))
    parts = ax.violinplot(plot_data, showmeans=True, showmedians=True)
    for i, pc in enumerate(parts["bodies"]):
        pc.set_facecolor(plot_colors[i])
        pc.set_alpha(0.7)

    ax.set_xticks(range(1, len(plot_labels) + 1))
    ax.set_xticklabels(plot_labels, fontsize=8, rotation=45, ha="right")
    ax.set_ylabel("Edit Norm (Frobenius)", fontsize=12)
    ax.set_title(f"{LABEL} Distribution of Per-Tensor Edit Magnitudes",
                 fontsize=14, fontweight="bold")
    fig.tight_layout()
    save_fig(fig, out_dir, f"{LABEL}_edit_distribution.svg")


# ---------- Graph 11: Tensor Type Breakdown ----------


def gen_tensor_type_breakdown(results_dir: Path, out_dir: Path) -> None:
    """Grouped bar: which tensor types each variant modifies."""
    all_types: set[str] = set()
    type_counts: dict[str, dict[str, int]] = {}

    for v in VARIANTS:
        fp = load_json(results_dir / v / f"fingerprint_{v}.json", quiet=True)
        if fp and "targeting" in fp:
            tt = fp["targeting"].get("tensor_types", {})
            type_counts[v] = tt
            all_types.update(tt.keys())

    if not type_counts or not all_types:
        print("  [SKIP] tensor_type_breakdown: no data")
        return

    def shorten(t: str) -> str:
        return (t.replace(".weight", "")
                 .replace("_proj", "")
                 .replace("self_attn.", "attn_")
                 .replace("per_layer_", "pl_"))

    type_order = sorted(all_types,
                        key=lambda t: sum(type_counts.get(v, {}).get(t, 0) for v in VARIANTS),
                        reverse=True)

    short_names = [shorten(t) for t in type_order]
    n = len(type_order)
    x = np.arange(n)
    variant_order = [v for v in VARIANTS if v in type_counts]
    n_vars = len(variant_order)
    width = 0.85 / max(n_vars, 1)

    fig, ax = plt.subplots(figsize=(max(16, n * 2), 9))

    for i, v in enumerate(variant_order):
        vals = [type_counts.get(v, {}).get(t, 0) for t in type_order]
        offset = (i - (n_vars - 1) / 2) * width
        ax.bar(x + offset, vals, width,
               label=VARIANT_LABELS[v], color=VARIANT_COLORS[v], alpha=0.85)

    ax.set_xticks(x)
    ax.set_xticklabels(short_names, fontsize=9, rotation=25, ha="right")
    ax.set_ylabel("Tensors Modified", fontsize=12)
    ax.legend(fontsize=5, ncol=5)
    ax.set_title(f"{LABEL} Tensor Type Targeting by Variant",
                 fontsize=14, fontweight="bold")
    fig.tight_layout()
    save_fig(fig, out_dir, f"{LABEL}_tensor_type_breakdown.svg")


# ---------- Main ----------


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(description=f"Generate {LABEL} model card SVGs")
    parser.add_argument("--results-dir", type=Path,
                        default=Path("comparisons/gemma4-e4b/results"))
    parser.add_argument("--output-dir", type=Path,
                        default=Path("comparisons/gemma4-e4b/graphs"))
    parser.add_argument("--db-path", type=Path,
                        default=Path("abliterlitics.db"))
    args = parser.parse_args()

    results_dir = args.results_dir
    out_dir = args.output_dir
    out_dir.mkdir(parents=True, exist_ok=True)

    print("=" * 60)
    print(f"{LABEL} Card Graph Generator")
    print(f"Results: {results_dir}")
    print(f"Output:  {out_dir}")
    print(f"DB:      {args.db_path}")
    print("=" * 60)

    # Load HarmBench from DB
    print("\nLoading HarmBench LLM-judged data from DB...")
    hb_data = load_harmbench_from_db(args.db_path)
    print(f"  Loaded {len(hb_data)} models")

    # Generate all graphs
    print("\nGenerating graphs...")
    gen_benchmark_comparison(results_dir, out_dir)
    gen_benchmark_delta(results_dir, out_dir)
    gen_gsm8k_comparison(results_dir, out_dir)
    gen_harmbench_summary(hb_data, out_dir)
    gen_harmbench_asr_category(hb_data, out_dir)
    gen_kl_divergence(results_dir, out_dir)
    gen_aggressiveness(results_dir, out_dir)
    gen_cosine_heatmap(results_dir, out_dir)
    gen_layer_comparison(results_dir, out_dir)
    gen_edit_distribution(results_dir, out_dir)
    gen_tensor_type_breakdown(results_dir, out_dir)

    all_svgs = sorted(out_dir.glob(f"{LABEL}_*.svg"))
    print(f"\n{'=' * 60}")
    print(f"Done. {len(all_svgs)} report SVGs saved to {out_dir}/")
    for svg in all_svgs:
        print(f"  {svg.name}")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()