#!/usr/bin/env python3
# Recompute metrics + figures directly from results_prod.csv (self-contained alpha).
import csv
from collections import defaultdict
import numpy as np
import matplotlib.pyplot as plt

CSV = "results_prod.csv"
LANGS = ["en", "es", "fr", "zh", "ar", "ru"]
TIER_COLOR = {"budget": "#d98c5f", "budget-mid": "#e0b050", "mid": "#4a8fb5"}

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

rows = list(csv.DictReader(open(CSV, encoding="utf-8")))
models, tier = [], {}
for r in rows:
    if r["model"] not in models:
        models.append(r["model"]); tier[r["model"]] = r.get("tier", "mid")
def rep0(r): return str(r.get("rep", "0")) == "0"

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

def graded_mean(model, lv):
    vals = []
    for lang in LANGS:
        ls = [r for r in rows if r["model"] == model and r["lang"] == lang and rep0(r)]
        if ls: vals.append(sum(1 for r in ls if r["grade"] in lv) / len(ls) * 100)
    return float(np.mean(vals)) if vals else 0.0

def krippendorff_alpha_nominal(model):
    # units = entities; ratings = predicted codes across languages (non-missing only)
    by_ent = defaultdict(dict)
    for r in rows:
        if r["model"] != model or not rep0(r): continue
        code = norm(r["pred_code"])
        if code:
            by_ent[r["entity_url"]][r["lang"]] = code
    # coincidence matrix
    o = defaultdict(float)
    for u, d in by_ent.items():
        vals = list(d.values())
        m = len(vals)
        if m < 2: continue
        for i in range(m):
            for j in range(m):
                if i != j:
                    o[(vals[i], vals[j])] += 1.0 / (m - 1)
    # marginals
    n_c = defaultdict(float)
    for (c, k), v in o.items():
        n_c[c] += v
    n = sum(n_c.values())
    if n < 2: return float("nan")
    Do = sum(v for (c, k), v in o.items() if c != k) / n
    De = sum(n_c[c] * n_c[k] for c in n_c for k in n_c if c != k) / (n * (n - 1))
    return 1 - Do / De if De > 0 else float("nan")

def en_chapter_concordance(model):
    by_ent = defaultdict(dict)
    for r in rows:
        if r["model"] != model or not rep0(r): continue
        by_ent[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: krippendorff_alpha_nominal(m) for m in models}
conc = {m: en_chapter_concordance(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
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)
h = [plt.Rectangle((0, 0), 1, 1, color=c) for c in TIER_COLOR.values()]
ax.legend(h, 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
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, 0.5)
plt.tight_layout(); plt.savefig("fig2_alpha.png", bbox_inches="tight"); plt.close()

# Fig 3
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
nonen = ["es", "fr", "zh", "ar", "ru"]
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)
cb = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04); cb.set_label("fraction matching English chapter")
plt.tight_layout(); plt.savefig("fig4_english_concordance_heatmap.png", bbox_inches="tight"); plt.close()

# numbers for cross-check
print("=== CSV直読みの確定値（原稿検算用）===")
print(f"{'model':24s} {'tier':11s} {'exact%':>7s} {'stem%':>6s} {'3char%':>7s} {'chap%':>6s} {'alpha':>7s}")
for m in order:
    print(f"{m:24s} {tier[m]:11s} {graded_mean(m,LV['exact']):7.1f} {graded_mean(m,LV['stem']):6.1f} "
          f"{graded_mean(m,LV['prefix3']):7.1f} {chap[m]:6.1f} {alph[m]:7.3f}")
print("\n英語との章一致（非英語）:")
for m in order:
    print(" ", f"{m:24s}", {l: round(conc[m][l], 3) for l in nonen})
