#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Regenerate Fig 2-5 in a Nature-journal style with colour-universal-design (CUD) palette."""
import json, csv
from collections import defaultdict
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

# ---------- Nature-style global rcParams ----------
mpl.rcParams.update({
    "font.family": "Liberation Sans",      # Arial-metric-compatible
    "font.size": 8,                          # Nature base ~7-8pt
    "axes.titlesize": 9,
    "axes.titleweight": "bold",
    "axes.labelsize": 8,
    "xtick.labelsize": 7.5,
    "ytick.labelsize": 7.5,
    "legend.fontsize": 7,
    "axes.linewidth": 0.6,
    "axes.edgecolor": "#4d4d4d",
    "axes.spines.top": False,
    "axes.spines.right": False,
    "xtick.major.width": 0.6,
    "ytick.major.width": 0.6,
    "xtick.color": "#4d4d4d",
    "ytick.color": "#4d4d4d",
    "xtick.labelcolor": "#1a1a1a",
    "ytick.labelcolor": "#1a1a1a",
    "axes.labelcolor": "#1a1a1a",
    "text.color": "#1a1a1a",
    "figure.dpi": 600,
    "savefig.dpi": 600,
    "savefig.bbox": "tight",
    "savefig.pad_inches": 0.02,
})

# ---------- Okabe-Ito colour-universal-design palette ----------
CUD = {
    "blue":        "#0072B2",
    "vermillion":  "#D55E00",
    "bluishgreen": "#009E73",
    "orange":      "#E69F00",
    "skyblue":     "#56B4E9",
    "reddishpurple":"#CC79A7",
    "yellow":      "#F0E442",
    "black":       "#000000",
    "grey":        "#999999",
}
# Tier -> CUD colour (distinguishable for all colour-vision types)
TIER_COLOR = {"budget": CUD["vermillion"], "budget-mid": CUD["orange"], "mid": CUD["blue"]}
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"}

ca = json.load(open("chapter_alpha.json"))   # m -> [exact_a, chap_a, lo, hi]
bs = json.load(open("bootstrap_ci.json"))     # m -> [ca, ca_lo, ca_hi, al, al_lo, al_hi]
order = sorted(ca.keys(), key=lambda m: bs[m][0], reverse=True)

MM = 1/25.4  # mm to inch
def tier_legend(ax, loc="lower right"):
    handles = [plt.Rectangle((0,0),1,1, fc=TIER_COLOR[t], ec="none") for t in ["mid","budget-mid","budget"]]
    ax.legend(handles, ["mid","budget-mid","budget"], title="vendor tier (price/size)",
              loc=loc, frameon=False, handlelength=1.1, handleheight=1.1,
              borderpad=0.4, labelspacing=0.35, title_fontsize=7)

# ============================================================
# Fig 3 — Chapter-level accuracy with 95% CI (horizontal bars)
# ============================================================
def fig_chapter_accuracy():
    val=[bs[m][0] for m in order]; lo=[bs[m][1] for m in order]; hi=[bs[m][2] for m in order]
    xerr=[[v-l for v,l in zip(val,lo)],[h-v for v,h in zip(val,hi)]]
    y=np.arange(len(order))[::-1]
    fig,ax=plt.subplots(figsize=(150*MM,72*MM))
    ax.barh(y,val,height=0.62,color=CUD["blue"],edgecolor="none",zorder=3)
    ax.errorbar(val,y,xerr=xerr,fmt="none",ecolor="#1a1a1a",elinewidth=0.7,
                capsize=2.2,capthick=0.7,zorder=4)
    for v,h,yy in zip(val,hi,y):
        ax.text(h+1.6,yy,f"{v:.1f}",va="center",ha="left",fontsize=7,color="#1a1a1a")
    ax.set_yticks(y); ax.set_yticklabels(order)
    ax.set_xlim(0,100); ax.set_xlabel("Chapter-level accuracy (%), mean of six languages")
    ax.xaxis.set_major_locator(mpl.ticker.MultipleLocator(20))
    ax.grid(axis="x",color="#e6e6e6",linewidth=0.5,zorder=0)
    ax.set_axisbelow(True)
    fig.savefig("fig1_chapter_accuracy.png"); plt.close(fig)
    print("\u2713 Fig 3 (chapter accuracy, neutral colour)")

# ============================================================
# Fig 4 — Chapter-level Krippendorff alpha with 95% CI + 0.667 ref
# ============================================================
def fig_chapter_alpha():
    val=[ca[m][1] for m in order]; lo=[ca[m][2] for m in order]; hi=[ca[m][3] for m in order]
    xerr=[[v-l for v,l in zip(val,lo)],[h-v for v,h in zip(val,hi)]]
    # M3: colour by alpha band (verifiable quantity), not vendor tier
    def band_color(a):
        if a>=0.800: return CUD["blue"]
        if a>=0.667: return CUD["skyblue"]
        return CUD["grey"]
    colors=[band_color(v) for v in val]
    y=np.arange(len(order))[::-1]
    fig,ax=plt.subplots(figsize=(150*MM,82*MM))
    # two conventional reference lines (drawn behind)
    ax.axvline(0.667,color="#9a9a9a",lw=0.7,ls=(0,(4,2)),zorder=2)
    ax.axvline(0.800,color="#4d4d4d",lw=0.9,ls=(0,(5,2)),zorder=2)
    ax.text(0.667,-0.46,"0.667\n(tentative)",fontsize=5.8,color="#7a7a7a",ha="center",va="top")
    ax.text(0.800,-0.46,"0.800\n(reliable)",fontsize=5.8,color="#4d4d4d",ha="center",va="top")
    ax.barh(y,val,height=0.62,color=colors,edgecolor="none",zorder=3)
    ax.errorbar(val,y,xerr=xerr,fmt="none",ecolor="#1a1a1a",elinewidth=0.7,
                capsize=2.2,capthick=0.7,zorder=4)
    for v,h,yy in zip(val,hi,y):
        ax.text(h+0.018,yy,f"{v:.3f}",va="center",ha="left",fontsize=7,color="#1a1a1a")
    ax.set_yticks(y); ax.set_yticklabels(order)
    ax.set_xlim(0,1.0); ax.set_xlabel("Chapter-level Krippendorff's \u03b1 across six languages")
    ax.xaxis.set_major_locator(mpl.ticker.MultipleLocator(0.2))
    ax.grid(axis="x",color="#e6e6e6",linewidth=0.5,zorder=0)
    ax.set_axisbelow(True)
    handles=[plt.Rectangle((0,0),1,1, fc=CUD["blue"], ec="none"),
             plt.Rectangle((0,0),1,1, fc=CUD["skyblue"], ec="none"),
             plt.Rectangle((0,0),1,1, fc=CUD["grey"], ec="none")]
    ax.legend(handles, ["\u03b1 \u2265 0.800 (reliable)","0.667\u20130.800 (tentative)","\u03b1 < 0.667"],
              loc="upper center", bbox_to_anchor=(0.5,-0.22), ncol=3, frameon=False,
              handlelength=1.1, handleheight=1.1, borderpad=0.3, columnspacing=1.4,
              labelspacing=0.3, fontsize=6.5)
    fig.savefig("fig2_alpha.png"); plt.close(fig)
    print("\u2713 Fig 4 (two thresholds, alpha-band colour)")

# ============================================================
# Fig 2 — Accuracy by matching stringency (nested levels)
# ============================================================
def fig_stringency():
    rows=list(csv.DictReader(open("results_prod.csv",encoding="utf-8")))
    def rep0(r): return str(r.get("rep","0"))=="0"
    levels=["exact","stem","prefix3","chapter"]
    rank={"exact":4,"stem":3,"prefix3":2,"chapter":1,"miss":0,"empty":0}
    # cumulative accuracy at each stringency level, per model (mean over all lang rows)
    acc=defaultdict(lambda: defaultdict(lambda:[0,0]))
    for r in rows:
        if not rep0(r): continue
        m=r["model"]; g=r["grade"]
        for lv in levels:
            acc[m][lv][1]+=1
            if rank.get(g,0)>=rank[lv]: acc[m][lv][0]+=1
    xlabels=["Exact","Stem","3-character","Chapter"]
    x=np.arange(len(levels))
    fig,ax=plt.subplots(figsize=(150*MM,82*MM))
    # performance-ordered single-hue palette (dark=strongest) + distinct marker/linestyle for B/W
    seq=["#08306B","#0072B2","#3690C0","#74A9CF","#A6BDDB","#BDBDBD"]
    line_color={m:seq[i] for i,m in enumerate(order)}
    styles={"gpt-5.4":("o","-"),"claude-sonnet-4-6":("s","-"),
            "gemini-2.5-flash":("^","--"),"claude-haiku-4-5":("D","--"),
            "gemini-2.5-flash-lite":("v",":"),"gpt-5.4-nano":("P",":")}
    for m in order:
        ys=[acc[m][lv][0]/acc[m][lv][1]*100 for lv in levels]
        mk,ls=styles[m]
        ax.plot(x,ys,ls=ls,lw=1.0,marker=mk,ms=4,color=line_color[m],
                markerfacecolor=line_color[m],markeredgecolor="white",
                markeredgewidth=0.4,label=m,zorder=3,clip_on=False)
    ax.set_xticks(x); ax.set_xticklabels(xlabels)
    ax.set_xlabel("Matching criterion (increasingly lenient \u2192)")
    ax.set_ylabel("Accuracy at each nested matching level (%)")
    ax.set_ylim(0,100); ax.set_xlim(-0.15,3.15)
    ax.yaxis.set_major_locator(mpl.ticker.MultipleLocator(20))
    ax.grid(axis="y",color="#e6e6e6",linewidth=0.5,zorder=0); ax.set_axisbelow(True)
    ax.legend(loc="upper left",frameon=False,handlelength=1.8,labelspacing=0.3,ncol=1)
    fig.savefig("fig_graded_accuracy.png"); plt.close(fig)
    print("✓ Fig 2 (stringency)")

# ============================================================
# Fig 5 — English-concordance heatmap (chapter level), CUD sequential
# ============================================================
def fig_heatmap():
    rows=list(csv.DictReader(open("results_prod.csv",encoding="utf-8")))
    def rep0(r): return str(r.get("rep","0"))=="0"
    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
    LANGS_NONEN=["es","fr","zh","ar","ru"]
    by=defaultdict(lambda: defaultdict(dict))
    for r in rows:
        if rep0(r): by[r["model"]][r["entity_url"]][r["lang"]]=norm(r["pred_code"])
    M=np.zeros((len(order),len(LANGS_NONEN)))
    for i,m in enumerate(order):
        for j,lg in enumerate(LANGS_NONEN):
            a=t=0
            for u,d in by[m].items():
                en=d.get("en",""); ot=d.get(lg,"")
                if en and ot:
                    t+=1
                    if en[:1]==ot[:1]: a+=1
            M[i,j]=a/t if t else np.nan
    # CUD-friendly sequential colormap: light yellow -> CUD blue (perceptually safe)
    cud_seq=LinearSegmentedColormap.from_list("cud_seq",
            ["#FFF7BC","#56B4E9","#0072B2","#08306B"])
    fig,ax=plt.subplots(figsize=(140*MM,82*MM))
    im=ax.imshow(M,cmap=cud_seq,vmin=0.4,vmax=1.0,aspect="auto")
    ax.set_xticks(range(len(LANGS_NONEN))); ax.set_xticklabels([l.upper() for l in LANGS_NONEN])
    ax.set_yticks(range(len(order))); ax.set_yticklabels(order)
    ax.set_xlabel("Input language (non-English)")
    # annotate; choose text colour for contrast (>=0.72 -> white)
    for i in range(M.shape[0]):
        for j in range(M.shape[1]):
            v=M[i,j]
            ax.text(j,i,f"{v:.2f}",ha="center",va="center",fontsize=6.8,
                    color="white" if v>=0.72 else "#1a1a1a")
    cb=fig.colorbar(im,ax=ax,fraction=0.046,pad=0.02)
    cb.set_label("Fraction matching the English chapter",fontsize=7)
    cb.ax.tick_params(labelsize=6.5,width=0.6)
    cb.outline.set_linewidth(0.6)
    ax.set_xticks(np.arange(-.5,len(LANGS_NONEN),1),minor=True)
    ax.set_yticks(np.arange(-.5,len(order),1),minor=True)
    ax.grid(which="minor",color="white",linewidth=1.0)
    ax.tick_params(which="minor",length=0)
    fig.savefig("fig4_heatmap.png"); plt.close(fig)
    print("✓ Fig 5 (heatmap)")

fig_stringency()
fig_chapter_accuracy()
fig_chapter_alpha()
fig_heatmap()
print("\nAll figures regenerated at 600 DPI (Nature style + Okabe-Ito CUD palette).")
