#!/usr/bin/env python3
"""
make_figures_from_csv.py
results_prod.csv を直接読み込んで Fig1-4 を生成する（数値の転記を一切しない確定版）。
ユーザPCの who フォルダ（results_prod.csv がある場所）で実行:
    pip install matplotlib numpy krippendorff
    python make_figures_from_csv.py
出力: fig1_chapter_accuracy.png / fig2_alpha.png / fig3_graded_accuracy.png /
      fig4_english_concordance_heatmap.png  （現行と同一レイアウト）
"""
import csv, sys
from collections import defaultdict
import numpy as np
import matplotlib.pyplot as plt

CSV = "results_prod.csv"
LANGS = ["en", "es", "fr", "zh", "ar", "ru"]
# 表示順・ティア色（CSVの tier 列があれば自動。なければこの既定を使用）
TIER_COLOR = {"budget": "#d98c5f", "budget-mid": "#e0b050", "mid": "#4a8fb5"}
DEFAULT_TIER = {
    "gpt-5.4-nano": "budget", "gpt-5.4": "mid",
    "gemini-2.5-flash-lite": "budget", "gemini-2.5-flash": "budget-mid",
    "claude-haiku-4-5": "budget-mid", "claude-sonnet-4-6": "mid",
}

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

def norm(c):
    if not c: return ""
    c = str(c).upper().strip().strip('"').strip("'").rstrip(".").strip()
    return c.split()[0] if c.split() else c

# ---- load ----
try:
    rows = list(csv.DictReader(open(CSV, encoding="utf-8")))
except FileNotFoundError:
    sys.exit(f"{CSV} が見つかりません。who フォルダで実行してください。")

models = []
tier = {}
for r in rows:
    m = r["model"]
    if m not in models:
        models.append(m)
        tier[m] = r.get("tier") or DEFAULT_TIER.get(m, "mid")

# rep==0 のみ（k=1）
def rep0(r): return str(r.get("rep", "0")) == "0"

# ---- metrics from CSV ----
def graded_mean(model, level_set):
    vals = []
    for lang in LANGS:
        ls = [r for r in rows if r["model"] == model and r["lang"] == lang and rep0(r)]
        if not ls: continue
        hit = sum(1 for r in ls if r["grade"] in level_set) / len(ls) * 100
        vals.append(hit)
    return float(np.mean(vals)) if vals else 0.0

LV = {"exact": {"exact"},
      "stem": {"exact", "stem"},
      "prefix3": {"exact", "stem", "prefix3"},
      "chapter": {"exact", "stem", "prefix3", "chapter"}}

def alpha_of(model):
    try:
        import krippendorff
    except ImportError:
        return float("nan")
    by_ent, ids, nid = {}, {}, 0
    for r in rows:
        if r["model"] != model or not rep0(r): continue
        code = norm(r["pred_code"]) or None
        if code is not None and code not in ids:
            ids[code] = nid; nid += 1
        by_ent.setdefault(r["entity_url"], {})[r["lang"]] = code
    if not by_ent: return float("nan")
    data = [[ids[by_ent[u].get(l)] if by_ent[u].get(l) in ids and by_ent[u].get(l) is not None
             else np.nan for u in by_ent] for l in LANGS]
    try:
        return krippendorff.alpha(reliability_data=np.array(data, dtype=float),
                                  level_of_measurement="nominal")
    except Exception:
        return float("nan")

def en_chapter_concordance(model):
    by_ent = {}
    for r in rows:
        if r["model"] != model or not rep0(r): continue
        by_ent.setdefault(r["entity_url"], {})[r["lang"]] = norm(r["pred_code"])
    out = {}
    for lang in LANGS:
        if lang == "en": continue
        m = t = 0
        for u, d in by_ent.items():
            en, ot = d.get("en", ""), d.get(lang, "")
            if en and ot:
                t += 1
                if ot[:1] == en[:1]: m += 1
        out[lang] = (m / t) if t else float("nan")
    return out

chap = {m: graded_mean(m, LV["chapter"]) for m in models}
alph = {m: alpha_of(m) for m in models}
order = sorted(models, key=lambda m: chap[m], reverse=True)
colors = [TIER_COLOR.get(tier[m], "#888") for m in order]
y = np.arange(len(order))

# ---- Fig 1: chapter accuracy ----
fig, ax = plt.subplots(figsize=(8, 4.2))
vals = [chap[m] for m in 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()

# ---- Fig 2: alpha ----
fig, ax = plt.subplots(figsize=(8, 4.2))
avals = [alph[m] 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, max(0.5, max(avals) * 1.15 if avals else 0.5))
plt.tight_layout(); plt.savefig("fig2_alpha.png", bbox_inches="tight"); plt.close()

# ---- Fig 3: graded funnel ----
fig, ax = plt.subplots(figsize=(9, 4.6))
levels = ["exact", "stem", "prefix3", "chapter"]
labels = ["Exact code", "Same stem", "Same 3-char", "Same chapter"]
x = np.arange(len(levels))
for m in order:
    ax.plot(x, [graded_mean(m, LV[lv]) for lv in levels], marker="o", label=m,
            color=TIER_COLOR.get(tier[m], "#888"),
            linewidth=2 if tier[m] == "mid" else 1.3, alpha=0.95 if tier[m] == "mid" else 0.8)
ax.set_xticks(x); ax.set_xticklabels(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()

# ---- Fig 4: heatmap ----
nonen = ["es", "fr", "zh", "ar", "ru"]
conc = {m: en_chapter_concordance(m) for m in order}
M = np.array([[conc[m][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()

# ---- console: print exact numbers for the manuscript ----
print("=== CSVから直接計算した数値（原稿の検算用）===")
print(f"{'model':24s} {'tier':11s} {'chapter%':>9s} {'alpha':>7s}")
for m in order:
    print(f"{m:24s} {tier[m]:11s} {chap[m]:9.1f} {alph[m]:7.3f}")
print("\n英語との章一致（非英語言語）:")
for m in order:
    print(" ", m, {l: round(conc[m][l], 3) for l in nonen})
print("\n保存: fig1_chapter_accuracy.png, fig2_alpha.png, fig3_graded_accuracy.png, fig4_english_concordance_heatmap.png")
