#!/usr/bin/env python3
# Figures for: Cross-lingual consistency of LLM ICD-11 coding (production run 2026-06-03)
# Data are the per-language summary values from results_prod.csv (N=300, 6 languages, k=1).
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import font_manager

plt.rcParams.update({
    "font.family": "DejaVu Sans", "font.size": 11,
    "axes.spines.top": False, "axes.spines.right": False,
    "axes.titleweight": "bold", "figure.dpi": 140,
})

LANGS = ["en", "es", "fr", "zh", "ar", "ru"]

# ---- production results (per language: en,es,fr,zh,ar,ru) ----
DATA = {
    "gpt-5.4-nano": {
        "tier": "budget",
        "exact":   [0.0,0.0,0.3,0.0,0.3,0.0],
        "stem":    [0.7,0.7,0.7,0.7,0.7,0.7],
        "prefix3": [2.0,2.3,1.7,1.7,2.0,2.0],
        "chapter": [14.3,13.7,17.0,14.0,13.0,14.0],
        "alpha": 0.140,
        "en_chap": {"es":0.530,"fr":0.483,"zh":0.597,"ar":0.553,"ru":0.603},
    },
    "gpt-5.4": {
        "tier": "mid",
        "exact":   [4.3,5.3,3.7,5.3,4.3,4.7],
        "stem":    [17.0,17.3,17.0,17.7,18.3,17.7],
        "prefix3": [39.7,41.0,41.7,39.3,41.3,40.0],
        "chapter": [84.0,85.0,84.0,81.3,84.0,83.0],
        "alpha": 0.394,
        "en_chap": {"es":0.933,"fr":0.940,"zh":0.933,"ar":0.927,"ru":0.920},
    },
    "gemini-2.5-flash-lite": {
        "tier": "budget",
        "exact":   [0.3,0.3,0.7,0.3,0.3,0.3],
        "stem":    [1.0,1.7,1.7,1.3,1.0,2.0],
        "prefix3": [2.0,2.7,3.7,1.7,1.7,3.0],
        "chapter": [26.0,29.7,28.0,23.7,27.3,25.0],
        "alpha": 0.227,
        "en_chap": {"es":0.590,"fr":0.517,"zh":0.637,"ar":0.567,"ru":0.587},
    },
    "gemini-2.5-flash": {
        "tier": "budget-mid",
        "exact":   [0.7,0.3,1.0,0.7,0.0,1.3],
        "stem":    [4.3,7.0,6.3,3.7,4.3,5.7],
        "prefix3": [16.7,21.3,21.3,15.0,17.7,16.7],
        "chapter": [59.3,63.3,65.3,53.7,60.7,60.3],
        "alpha": 0.242,
        "en_chap": {"es":0.713,"fr":0.710,"zh":0.737,"ar":0.753,"ru":0.757},
    },
    "claude-haiku-4-5": {
        "tier": "budget-mid",
        "exact":   [0.0,0.0,0.0,0.0,0.0,0.3],
        "stem":    [2.0,2.0,1.7,1.7,2.0,1.7],
        "prefix3": [5.7,7.0,6.3,6.0,5.7,7.0],
        "chapter": [22.3,25.7,23.3,21.7,25.3,25.3],
        "alpha": 0.334,
        "en_chap": {"es":0.760,"fr":0.763,"zh":0.780,"ar":0.733,"ru":0.770},
    },
    "claude-sonnet-4-6": {
        "tier": "mid",
        "exact":   [5.7,4.0,4.3,4.7,3.3,4.0],
        "stem":    [22.0,21.0,24.3,23.0,21.3,23.0],
        "prefix3": [43.7,46.3,48.7,44.7,46.7,48.0],
        "chapter": [75.0,80.3,80.7,78.3,80.0,78.7],
        "alpha": 0.398,
        "en_chap": {"es":0.860,"fr":0.853,"zh":0.870,"ar":0.857,"ru":0.853},
    },
}

TIER_COLOR = {"budget": "#d98c5f", "budget-mid": "#e0b050", "mid": "#4a8fb5"}

def mean(model, key):  # mean across languages
    return float(np.mean(DATA[model][key]))

# order models by chapter-level accuracy (descending)
order = sorted(DATA.keys(), key=lambda m: mean(m, "chapter"), reverse=True)

# ============ Figure 1: chapter-level accuracy by model ============
fig, ax = plt.subplots(figsize=(8, 4.2))
vals = [mean(m, "chapter") for m in order]
colors = [TIER_COLOR[DATA[m]["tier"]] for m in order]
y = np.arange(len(order))
ax.barh(y, vals, color=colors)
ax.set_yticks(y); ax.set_yticklabels(order)
ax.invert_yaxis()
ax.set_xlabel("Chapter-level accuracy (%)  — mean of 6 languages")
ax.set_title("How often the model places the case in the correct ICD-11 chapter")
for i, v in enumerate(vals):
    ax.text(v + 1, i, f"{v:.0f}%", va="center", fontsize=10)
ax.set_xlim(0, 100)
handles = [plt.Rectangle((0,0),1,1,color=c) for c in TIER_COLOR.values()]
ax.legend(handles, TIER_COLOR.keys(), title="capability tier", loc="lower right", frameon=False)
plt.tight_layout(); plt.savefig("fig1_chapter_accuracy.png", bbox_inches="tight"); plt.close()

# ============ Figure 2: cross-lingual consistency (alpha) ============
fig, ax = plt.subplots(figsize=(8, 4.2))
avals = [DATA[m]["alpha"] for m in order]
ax.barh(y, avals, color=colors)
ax.set_yticks(y); ax.set_yticklabels(order); ax.invert_yaxis()
ax.set_xlabel("Krippendorff's α (nominal) across 6 languages")
ax.set_title("Cross-lingual consistency: same case, same code across languages?")
for i, v in enumerate(avals):
    ax.text(v + 0.005, i, f"{v:.3f}", va="center", fontsize=10)
ax.set_xlim(0, 0.5)
plt.tight_layout(); plt.savefig("fig2_alpha.png", bbox_inches="tight"); plt.close()

# ============ Figure 3: graded accuracy funnel (mean across languages) ============
fig, ax = plt.subplots(figsize=(9, 4.6))
levels = ["exact", "stem", "prefix3", "chapter"]
level_labels = ["Exact code", "Same stem", "Same 3-char", "Same chapter"]
x = np.arange(len(levels))
w = 0.13
for j, m in enumerate(order):
    means = [mean(m, lv) for lv in levels]
    ax.plot(x, means, marker="o", label=m, color=TIER_COLOR[DATA[m]["tier"]],
            linewidth=2 if DATA[m]["tier"]=="mid" else 1.3,
            alpha=0.95 if DATA[m]["tier"]=="mid" else 0.8)
ax.set_xticks(x); ax.set_xticklabels(level_labels)
ax.set_ylabel("Accuracy (%) — mean of 6 languages")
ax.set_title("Accuracy by matching stringency: exact match is near-zero, chapter-level separates models")
ax.legend(frameon=False, fontsize=9, ncol=2)
ax.set_ylim(0, 100)
plt.tight_layout(); plt.savefig("fig3_graded_accuracy.png", bbox_inches="tight"); plt.close()

# ============ Figure 4: heatmap model x language (chapter-level English concordance) ============
nonen = ["es", "fr", "zh", "ar", "ru"]
M = np.array([[DATA[m]["en_chap"][l] for l in nonen] for m in order])
fig, ax = plt.subplots(figsize=(7.6, 4.6))
im = ax.imshow(M, cmap="YlGnBu", vmin=0.4, vmax=1.0, aspect="auto")
ax.set_xticks(range(len(nonen))); ax.set_xticklabels([l.upper() for l in nonen])
ax.set_yticks(range(len(order))); ax.set_yticklabels(order)
ax.set_title("Agreement with the English code (chapter level)")
for i in range(M.shape[0]):
    for k in range(M.shape[1]):
        ax.text(k, i, f"{M[i,k]:.2f}", ha="center", va="center",
                color="white" if M[i,k] > 0.78 else "black", fontsize=9)
cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
cbar.set_label("fraction matching English chapter")
plt.tight_layout(); plt.savefig("fig4_english_concordance_heatmap.png", bbox_inches="tight"); plt.close()

# ============ Combined 2x2 panel ============
fig, axes = plt.subplots(2, 2, figsize=(15, 9))
# (a)
ax = axes[0,0]; ax.barh(y, vals, color=colors); ax.set_yticks(y); ax.set_yticklabels(order); ax.invert_yaxis()
ax.set_xlabel("Chapter-level accuracy (%)"); ax.set_title("(a) Correct ICD-11 chapter (mean of 6 languages)"); ax.set_xlim(0,100)
for i,v in enumerate(vals): ax.text(v+1,i,f"{v:.0f}%",va="center",fontsize=9)
# (b)
ax = axes[0,1]; ax.barh(y, avals, color=colors); ax.set_yticks(y); ax.set_yticklabels(order); ax.invert_yaxis()
ax.set_xlabel("Krippendorff's α"); ax.set_title("(b) Cross-lingual consistency (α)"); ax.set_xlim(0,0.5)
for i,v in enumerate(avals): ax.text(v+0.005,i,f"{v:.3f}",va="center",fontsize=9)
# (c)
ax = axes[1,0]
for m in order:
    ax.plot(x, [mean(m,lv) for lv in levels], marker="o", label=m,
            color=TIER_COLOR[DATA[m]["tier"]], linewidth=2 if DATA[m]["tier"]=="mid" else 1.3)
ax.set_xticks(x); ax.set_xticklabels(level_labels); ax.set_ylabel("Accuracy (%)")
ax.set_title("(c) Accuracy by matching stringency"); ax.legend(frameon=False, fontsize=8, ncol=2); ax.set_ylim(0,100)
# (d)
ax = axes[1,1]; im = ax.imshow(M, cmap="YlGnBu", vmin=0.4, vmax=1.0, aspect="auto")
ax.set_xticks(range(len(nonen))); ax.set_xticklabels([l.upper() for l in nonen])
ax.set_yticks(range(len(order))); ax.set_yticklabels(order); ax.set_title("(d) Agreement with English chapter")
for i in range(M.shape[0]):
    for k in range(M.shape[1]):
        ax.text(k,i,f"{M[i,k]:.2f}",ha="center",va="center",color="white" if M[i,k]>0.78 else "black",fontsize=8)
plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
plt.suptitle("Cross-lingual LLM ICD-11 coding — pilot/production results (N=300, 6 languages, 2026-06-03)",
             fontsize=13, fontweight="bold")
plt.tight_layout(rect=[0,0,1,0.97]); plt.savefig("fig_combined_panel.png", bbox_inches="tight"); plt.close()

print("saved: fig1_chapter_accuracy.png, fig2_alpha.png, fig3_graded_accuracy.png,")
print("       fig4_english_concordance_heatmap.png, fig_combined_panel.png")
print("\nmodel order (by chapter accuracy):")
for m in order:
    print(f"  {m:24s} tier={DATA[m]['tier']:11s} chapter={mean(m,'chapter'):.1f}%  alpha={DATA[m]['alpha']:.3f}")
