scripts/plot/_plot_rq2_query_translation.py

"""
Generate RQ2 query-translation result plots.

Produces:
  artifacts/results/rq2_query_translation_overview.png   -- QAR/QCR bar chart
  artifacts/results/rq2_query_translation_perspec.png    -- per-spec heatmap
  artifacts/results/rq2_query_translation_breakdown.png  -- correct / wrong-verdict / compile-fail stacked bar
"""
import csv
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np

ROOT = Path(__file__).resolve().parents[1]
OUT  = ROOT / "artifacts" / "results"
OUT.mkdir(parents=True, exist_ok=True)

# ── data ──────────────────────────────────────────────────────────────────────
SYSTEMS = {
    "Claude":   ROOT / "artifacts/baselines/claude/trans_query_eval.csv",
    "Grok":     ROOT / "artifacts/baselines/grok/trans_query_eval.csv",
    "GPT-4.1":  ROOT / "artifacts/baselines/gpt_4_1/trans_query_eval.csv",
    "Ours":     ROOT / "artifacts/ours/trans_query.csv",
}

COLORS = {
    "Claude":  "#4C72B0",
    "Grok":    "#55A868",
    "GPT-4.1": "#C44E52",
    "Ours":    "#DD8452",
}

SPEC_LABELS = [
    "S01\nCoffee", "S02\nTraffic", "S03\nCounter", "S04\nProd-Con",
    "S05\nBank",   "S06\nTrain1",  "S07\nCoffee2", "S08\nGearbox",
    "S09\nMutex",  "S10\nMaster",  "S11\nPump",    "S12\nPacemaker",
    "S13\nPatient","S14\nGDPR1",   "S15\nGDPR2",   "S16\nTrain2",
    "S17\nRBC",    "S18\nAircraft","S19\nCable",   "S20\nFire",
]


def load(path):
    return list(csv.DictReader(open(path, encoding="utf-8")))


def stats(rows):
    total = len(rows)
    correct    = sum(1 for r in rows if r.get("correct", "") == "Y")
    compile_ok = sum(1 for r in rows if r.get("verdict", "") in ("T", "F"))
    wrong_v    = compile_ok - correct      # compiled but wrong verdict
    comp_fail  = total - compile_ok        # failed to compile
    return {
        "total": total, "correct": correct, "compile_ok": compile_ok,
        "wrong_verdict": wrong_v, "compile_fail": comp_fail,
        "qar": 100 * correct / total,
        "qcr": 100 * compile_ok / total,
    }


def per_spec(rows):
    sp = {i: {"correct": 0, "wrong_v": 0, "comp_fail": 0} for i in range(1, 21)}
    for r in rows:
        sid = int(r["spec_id"])
        if r.get("correct", "") == "Y":
            sp[sid]["correct"] += 1
        elif r.get("verdict", "") in ("T", "F"):
            sp[sid]["wrong_v"] += 1
        else:
            sp[sid]["comp_fail"] += 1
    return sp


data = {name: {"rows": load(p)} for name, p in SYSTEMS.items()}
for name, d in data.items():
    d.update(stats(d["rows"]))
    d["per_spec"] = per_spec(d["rows"])

names  = list(data.keys())
qar    = [data[n]["qar"]  for n in names]
qcr    = [data[n]["qcr"]  for n in names]


# ── Figure 1: Overview bar chart (QAR + QCR) ──────────────────────────────────
fig, ax = plt.subplots(figsize=(8, 5))
x = np.arange(len(names))
w = 0.35

bars1 = ax.bar(x - w/2, qar, w, label="QAR (correct verdicts)", color=[COLORS[n] for n in names],
               zorder=3)
bars2 = ax.bar(x + w/2, qcr, w, label="QCR (compilation rate)", color=[COLORS[n] for n in names],
               alpha=0.45, zorder=3, hatch="///")

# value labels
for bar in bars1:
    ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.6,
            f"{bar.get_height():.0f}%", ha="center", va="bottom", fontsize=10, fontweight="bold")
for bar in bars2:
    ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.6,
            f"{bar.get_height():.0f}%", ha="center", va="bottom", fontsize=10, color="#444")

ax.set_xticks(x)
ax.set_xticklabels(names, fontsize=11)
ax.set_ylabel("Percentage (%)", fontsize=11)
ax.set_title("RQ2 – Query Translation: QAR and QCR by System\n(100 queries, 20 UPPAAL models)",
             fontsize=12, fontweight="bold")
ax.set_ylim(0, 112)
ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0)
ax.set_axisbelow(True)

# custom legend
solid = mpatches.Patch(facecolor="#888", label="QAR – Query Answer Rate")
hatch = mpatches.Patch(facecolor="#888", alpha=0.45, hatch="///", label="QCR – Query Compilation Rate")
ax.legend(handles=[solid, hatch], fontsize=10, loc="lower right")

# highlight "Ours"
idx = names.index("Ours")
ax.get_xticklabels()[idx].set_color(COLORS["Ours"])
ax.get_xticklabels()[idx].set_fontweight("bold")

plt.tight_layout()
p1 = OUT / "rq2_query_translation_overview.png"
fig.savefig(p1, dpi=150, bbox_inches="tight")
plt.close(fig)
print(f"Saved {p1}")


# ── Figure 2: Per-spec QAR heatmap ────────────────────────────────────────────
mat = np.zeros((len(names), 20))
for i, name in enumerate(names):
    sp = data[name]["per_spec"]
    for sid in range(1, 21):
        mat[i, sid-1] = sp[sid]["correct"]

fig, ax = plt.subplots(figsize=(14, 3.6))
im = ax.imshow(mat, cmap="RdYlGn", vmin=0, vmax=5, aspect="auto")

# cell annotations
for i in range(len(names)):
    for j in range(20):
        v = int(mat[i, j])
        c = "white" if v <= 1 else "black"
        ax.text(j, i, f"{v}/5", ha="center", va="center", fontsize=8.5,
                color=c, fontweight="bold")

ax.set_xticks(range(20))
ax.set_xticklabels(SPEC_LABELS, fontsize=7.5, rotation=0)
ax.set_yticks(range(len(names)))
ax.set_yticklabels(names, fontsize=10)
ax.set_title("RQ2 – Correct Queries per Specification (out of 5)\nGreen = all correct, Red = all wrong",
             fontsize=11, fontweight="bold")

cbar = fig.colorbar(im, ax=ax, orientation="vertical", fraction=0.02, pad=0.01)
cbar.set_label("Correct (0–5)", fontsize=9)

plt.tight_layout()
p2 = OUT / "rq2_query_translation_perspec.png"
fig.savefig(p2, dpi=150, bbox_inches="tight")
plt.close(fig)
print(f"Saved {p2}")


# ── Figure 3: Stacked bar — correct / wrong-verdict / compile-fail ─────────────
correct_v   = [data[n]["correct"]       for n in names]
wrong_v     = [data[n]["wrong_verdict"] for n in names]
comp_fail_v = [data[n]["compile_fail"]  for n in names]

fig, ax = plt.subplots(figsize=(8, 4.5))
x = np.arange(len(names))
w = 0.5

b1 = ax.bar(x, correct_v,   w, label="Correct verdict",  color="#4CAF50", zorder=3)
b2 = ax.bar(x, wrong_v,     w, bottom=correct_v,          label="Wrong verdict",   color="#FFC107", zorder=3)
b3 = ax.bar(x, comp_fail_v, w,
            bottom=[c+w for c, w in zip(correct_v, wrong_v)],
            label="Compile failure", color="#F44336", zorder=3)

# QAR label on top
for i, (c, w_v, cf) in enumerate(zip(correct_v, wrong_v, comp_fail_v)):
    ax.text(i, c + w_v + cf + 1, f"{c}%", ha="center", va="bottom",
            fontsize=10, fontweight="bold")

ax.set_xticks(x)
ax.set_xticklabels(names, fontsize=11)
ax.set_ylabel("Number of queries (out of 100)", fontsize=11)
ax.set_title("RQ2 – Query Outcome Breakdown by System", fontsize=12, fontweight="bold")
ax.set_ylim(0, 112)
ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0)
ax.set_axisbelow(True)
ax.legend(fontsize=10, loc="lower right")

plt.tight_layout()
p3 = OUT / "rq2_query_translation_breakdown.png"
fig.savefig(p3, dpi=150, bbox_inches="tight")
plt.close(fig)
print(f"Saved {p3}")


# ── Figure 4: Radar / spider chart ─────────────────────────────────────────────
# Group specs into 4 domains
GROUPS = {
    "Embedded\n(S01-S10)": list(range(1, 11)),
    "Legal/GDPR\n(S13-S15)": [13, 14, 15],
    "Transport\n(S06,S16,S17,S18)": [6, 16, 17, 18],
    "Industrial\n(S11,S12,S19,S20)": [11, 12, 19, 20],
}

fig, ax = plt.subplots(figsize=(8, 4.5))
bar_w = 0.18
group_names = list(GROUPS.keys())
gx = np.arange(len(group_names))

for i, name in enumerate(names):
    sp = data[name]["per_spec"]
    scores = []
    for specs_in_group in GROUPS.values():
        total_q = len(specs_in_group) * 5
        correct_q = sum(sp[s]["correct"] for s in specs_in_group)
        scores.append(100 * correct_q / total_q)
    offset = (i - (len(names)-1)/2) * bar_w
    bars = ax.bar(gx + offset, scores, bar_w, label=name, color=COLORS[name],
                  zorder=3, alpha=0.88)

ax.set_xticks(gx)
ax.set_xticklabels(group_names, fontsize=9)
ax.set_ylabel("QAR (%)", fontsize=11)
ax.set_title("RQ2 – QAR by Domain Group", fontsize=12, fontweight="bold")
ax.set_ylim(0, 112)
ax.yaxis.grid(True, linestyle="--", alpha=0.5, zorder=0)
ax.set_axisbelow(True)
ax.legend(fontsize=10)

plt.tight_layout()
p4 = OUT / "rq2_query_translation_by_domain.png"
fig.savefig(p4, dpi=150, bbox_inches="tight")
plt.close(fig)
print(f"Saved {p4}")