import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import numpy as np
 
# 資料
categories = [
    "Direct Injection",
    "Indirect Injection",
    "Role-play Jailbreak",
    "Encoding Bypass",
    "System Prompt Extraction",
    "Tool Abuse",
    "Multi-turn Manipulation",
    "Safety Bypass",
]
asr_values = [45, 20, 70, 15, 60, 10, 35, 55]
 
# 依嚴重性門檻著色
colors = []
for v in asr_values:
    if v >= 50:
        colors.append("#dc3545")   # 紅 - 嚴重
    elif v >= 30:
        colors.append("#fd7e14")   # 橘 - 警示
    else:
        colors.append("#28a745")   # 綠 - 可接受
 
fig, ax = plt.subplots(figsize=(10, 6))
bars = ax.barh(categories, asr_values, color=colors, height=0.6)
 
# 加上數值標籤
for bar, val in zip(bars, asr_values):
    ax.text(bar.get_width() + 1.5, bar.get_y() + bar.get_height() / 2,
            f"{val}%", va="center", fontsize=11, fontweight="bold")
 
# 門檻線
ax.axvline(x=50, color="#dc3545", linestyle="--", alpha=0.5, label="Critical threshold")
ax.axvline(x=30, color="#fd7e14", linestyle="--", alpha=0.5, label="Warning threshold")
 
ax.set_xlabel("Attack Success Rate (%)", fontsize=12)
ax.set_title("Attack Success Rate by Category", fontsize=14, fontweight="bold")
ax.set_xlim(0, 100)
ax.xaxis.set_major_formatter(mticker.PercentFormatter())
ax.legend(loc="lower right")
ax.invert_yaxis()
plt.tight_layout()
plt.savefig("asr_by_category.png", dpi=150, bbox_inches="tight")
plt.show()

import matplotlib.pyplot as plt
import numpy as np
 
categories = ["Injection", "Jailbreak", "Extraction", "Tool Abuse", "Safety"]
model_a = [45, 70, 60, 10, 55]
model_b = [25, 40, 30, 5, 20]
 
x = np.arange(len(categories))
width = 0.35
 
fig, ax = plt.subplots(figsize=(10, 6))
bars1 = ax.bar(x - width / 2, model_a, width, label="Before Remediation",
               color="#dc3545", alpha=0.8)
bars2 = ax.bar(x + width / 2, model_b, width, label="After Remediation",
               color="#28a745", alpha=0.8)
 
ax.set_ylabel("Attack Success Rate (%)")
ax.set_title("Remediation Effectiveness by Attack Category")
ax.set_xticks(x)
ax.set_xticklabels(categories)
ax.legend()
ax.set_ylim(0, 100)
 
# 加上數值標籤
for bars in [bars1, bars2]:
    for bar in bars:
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width() / 2, height + 1,
                f"{height}%", ha="center", va="bottom", fontsize=9)
 
plt.tight_layout()
plt.savefig("remediation_comparison.png", dpi=150, bbox_inches="tight")
plt.show()

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
 
# 列：攻擊類型；欄：目標元件
attack_types = [
    "Direct Injection", "Indirect Injection", "Role-play Jailbreak",
    "Encoding Bypass", "Prompt Extraction", "Tool Manipulation",
]
components = ["Chat API", "RAG Pipeline", "Agent Tools", "Safety Filter"]
 
# 數值：ASR 百分比（-1 = 未測試）
data = np.array([
    [45, 20, -1, 30],
    [15, 55, 40, 10],
    [70, -1, 25, 50],
    [15, 5,  -1, 20],
    [60, 30, 10, -1],
    [-1, -1, 35, 5],
])
 
# 為未測試格遮罩
mask = data == -1
display_data = np.where(mask, 0, data)
 
fig, ax = plt.subplots(figsize=(10, 7))
sns.heatmap(
    display_data, annot=True, fmt="d", mask=mask,
    xticklabels=components, yticklabels=attack_types,
    cmap="RdYlGn_r", vmin=0, vmax=100,
    linewidths=1, linecolor="white",
    cbar_kws={"label": "Attack Success Rate (%)"},
    ax=ax,
)
 
# 標示未測試格
for i in range(data.shape[0]):
    for j in range(data.shape[1]):
        if mask[i, j]:
            ax.text(j + 0.5, i + 0.5, "N/T", ha="center", va="center",
                    fontsize=10, color="gray", fontstyle="italic")
 
ax.set_title("Attack Taxonomy Coverage Heatmap", fontsize=14, fontweight="bold")
ax.set_ylabel("Attack Type")
ax.set_xlabel("Target Component")
plt.tight_layout()
plt.savefig("attack_heatmap.png", dpi=150, bbox_inches="tight")
plt.show()

import matplotlib.pyplot as plt
 
labels = ["Critical", "High", "Medium", "Low"]
sizes = [2, 4, 6, 3]
colors = ["#dc3545", "#fd7e14", "#ffc107", "#28a745"]
explode = (0.05, 0, 0, 0)
 
fig, ax = plt.subplots(figsize=(8, 8))
wedges, texts, autotexts = ax.pie(
    sizes, explode=explode, labels=labels, colors=colors,
    autopct=lambda pct: f"{pct:.0f}%\n({int(round(pct / 100 * sum(sizes)))})",
    shadow=False, startangle=90, pctdistance=0.75,
    textprops={"fontsize": 12},
)
 
# 畫中心圓營造甜甜圈效果
centre_circle = plt.Circle((0, 0), 0.50, fc="white")
ax.add_artist(centre_circle)
ax.text(0, 0, f"{sum(sizes)}\nFindings", ha="center", va="center",
        fontsize=16, fontweight="bold")
 
ax.set_title("Finding Severity Distribution", fontsize=14, fontweight="bold")
plt.tight_layout()
plt.savefig("severity_distribution.png", dpi=150, bbox_inches="tight")
plt.show()