src/visualize.py

"""
Part 3: visualizations comparing an analyzed run against U.S. census data.

Single-model charts, styled to match the standardized Reformation Bias Audit
figures (Times-family type, Wong 2011 palettes, consistent legend chrome).
Cross-model / multi-panel comparison figures are not rendered here.

Figure map (paper -> here, one model at a time):
  Fig 1    skin_tone_colorspace()     L* vs H* cloud, colored by race
  Fig 2.2  race_vs_census_scatter()   All races, colored by race
  Fig 3    race_vs_census_scatter()   one race, colored by proxy + annotations
  Table 2  compute_tvd()              Total Variation Distance vs census
  Fig 4    prompt_race_bar()          per-prompt stacked bars vs census
  Fig 8    homogeneity_gap_bars()     intra-minus-inter cosine gap by race
  Fig 10   homogeneity_violin()       box+strip+bootstrap CI by L* bin
"""
import io
import os
import tempfile
import zlib

import numpy as np
import pandas as pd
import matplotlib

matplotlib.use("Agg")  # headless: render to figures, never open a window
import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb
from matplotlib.lines import Line2D
from matplotlib.patches import Patch, Rectangle
from mpl_toolkits.axes_grid1 import make_axes_locatable

try:
    from adjustText import adjust_text
except ImportError:  # optional; annotated scatter still draws labels
    adjust_text = None

# ── GLOBAL FONT: Times-family serif, matching the paper ───────────────────
plt.rcParams["font.family"] = "serif"
plt.rcParams["font.serif"] = ["Liberation Serif", "Times New Roman", "DejaVu Serif"]
plt.rcParams.update({
    "font.size": 16,
    "axes.titlesize": 17,
    "axes.labelsize": 17,
    "xtick.labelsize": 15,
    "ytick.labelsize": 15,
    "legend.fontsize": 15,
    "legend.title_fontsize": 16,
    "figure.titlesize": 19,
    "axes.linewidth": 0.9,
    "grid.linewidth": 0.5,
    "figure.dpi": 120,
    "savefig.dpi": 200,
})

# ---- canonical races + palettes (Wong 2011 / colorblind-safe) ----------
RACES = ["White", "Black", "Asian", "Hispanic", "Other"]
PRIMARY_RACES = ["White", "Black", "Asian", "Hispanic"]  # paper scatters omit Other
PLOT_ORDER = ["White", "Asian", "Black", "Hispanic", "Other"]  # scatter z-order
LEGEND_RACE_ORDER = ["White", "Black", "Asian", "Hispanic", "Other"]
FIG8_RACE_ORDER = ["White", "Black", "Asian", "Hispanic", "Other"]
PROXY_ORDER = ["occupation", "surname", "geography"]

RACE_CANON = {
    "White/Caucasian": "White", "White": "White",
    "Black/African": "Black", "Black": "Black",
    "Asian": "Asian",
    "Hispanic/Latino": "Hispanic", "Hispanic": "Hispanic",
    "Other": "Other", "Other/Unclear": "Other",
}
RACE_COLORS = {
    "Asian": "#009E73",
    "Black": "#0072B2",
    "Hispanic": "#F0E442",
    "Other": "#CC79A7",
    "White": "#D55E00",
    "Unknown": "#BBBBBB",
}
PROXY_COLORS = {
    "Occupation": "#7570B3",
    "Surname": "#E7298A",
    "Geography": "#66A61E",
    "occupation": "#7570B3",
    "surname": "#E7298A",
    "geography": "#66A61E",
}
L_BINS = [(0, 40), (40, 55), (55, 70), (70, 101)]
FACE_DET_MIN = 0.85
MIN_BIN_SIZE = 5
N_BOOTSTRAP = 1000
CI_LEVEL = 95
SKIN_PALETTE = plt.get_cmap("copper")(np.linspace(0.15, 0.9, 4))


# ====================================================================== style
def apply_paper_style(ax):
    """Clean, publication-ready axis styling."""
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color("#CCCCCC")
    ax.spines["bottom"].set_color("#CCCCCC")
    ax.tick_params(colors="#444444", length=3)
    ax.yaxis.label.set_color("#222222")
    ax.xaxis.label.set_color("#222222")


def style_legend(leg, text_colors=None):
    """White background, thin light-grey border, black text."""
    if leg is None:
        return leg
    frame = leg.get_frame()
    frame.set_facecolor("white")
    frame.set_edgecolor("#CCCCCC")
    frame.set_linewidth(0.8)
    frame.set_alpha(0.95)
    texts = leg.get_texts()
    for i, t in enumerate(texts):
        t.set_color(text_colors[i] if text_colors else "black")
    return leg


def _legend_outside(ax, *args, **kwargs):
    """Park the legend to the right of the axes so it never covers data."""
    kwargs.setdefault("loc", "center left")
    kwargs.setdefault("bbox_to_anchor", (1.04, 0.5))
    kwargs.setdefault("borderaxespad", 0.0)
    kwargs.setdefault("frameon", True)
    leg = ax.legend(*args, **kwargs)
    return style_legend(leg)


def _legend_above(ax, *args, **kwargs):
    """Park the legend above the axes (good when a colorbar sits below)."""
    kwargs.setdefault("loc", "lower center")
    kwargs.setdefault("bbox_to_anchor", (0.5, 1.04))
    kwargs.setdefault("borderaxespad", 0.0)
    kwargs.setdefault("frameon", True)
    kwargs.setdefault("ncol", kwargs.get("ncol", 3))
    leg = ax.legend(*args, **kwargs)
    return style_legend(leg)


def _repel_labels(texts, ax, xs=None, ys=None):
    """Push text labels off each other and off the data points, with leader lines."""
    if not texts:
        return
    if adjust_text is None:
        return
    common = dict(
        ax=ax,
        arrowprops=dict(arrowstyle="-", color="#666666", lw=0.7, alpha=0.85,
                        shrinkA=4, shrinkB=4),
    )
    if xs is not None and ys is not None:
        common["x"] = np.asarray(xs)
        common["y"] = np.asarray(ys)
    try:
        adjust_text(
            texts,
            expand_points=(1.6, 1.6),
            expand_text=(1.3, 1.3),
            force_points=(0.5, 0.6),
            force_text=(0.9, 1.1),
            lim=800,
            avoid_self=True,
            **common,
        )
    except TypeError:
        try:
            adjust_text(
                texts,
                expand=(1.4, 1.5),
                force_text=(0.8, 1.0),
                force_points=(0.5, 0.7),
                **common,
            )
        except TypeError:
            adjust_text(texts, ax=ax, arrowprops=common["arrowprops"])


def fig_to_svg_html(fig):
    """Vector (SVG) markup for in-app preview. Dense scatters may be rasterized
    inside the SVG so the file stays light; axes and type stay vector."""
    buf = io.BytesIO()
    fig.savefig(buf, format="svg", bbox_inches="tight", dpi=200)
    svg = buf.getvalue().decode("utf-8")
    if svg.lstrip().startswith("<?xml"):
        svg = svg.split("?>", 1)[-1]
    stripped = svg.lstrip()
    if stripped[:9].upper() == "<!DOCTYPE":
        svg = stripped[stripped.find(">") + 1:]
    if "preserveAspectRatio" not in svg:
        svg = svg.replace("<svg ", '<svg preserveAspectRatio="xMidYMid meet" ', 1)
    return f'<div class="chart-svg">{svg}</div>'


def fig_to_pdf_path(fig, stem):
    """Write a vector PDF and return its path (for per-chart download)."""
    path = os.path.join(tempfile.gettempdir(), f"{stem}.pdf")
    fig.savefig(path, format="pdf", bbox_inches="tight", dpi=200)
    return path


def _empty_fig(msg):
    fig, ax = plt.subplots(figsize=(8, 5))
    ax.text(0.5, 0.5, msg, ha="center", va="center", wrap=True, fontsize=16, color="#444444")
    ax.axis("off")
    return fig


def _canon_race_series(df):
    col = None
    for name in ("race_classification", "Race", "race"):
        if name in df.columns:
            col = name
            break
    if col is None:
        return pd.Series(np.nan, index=df.index)
    return df[col].map(RACE_CANON)


def model_name(df):
    vals = df["model"].dropna().unique()
    return vals[0] if len(vals) else "model"


def _proxy_label(cat):
    return str(cat).strip().title()


def _proxy_color(cat):
    return PROXY_COLORS.get(cat, PROXY_COLORS.get(_proxy_label(cat), "#888888"))


def _display_prompt_name(category, name):
    """Surname census names are stored ALL CAPS; title-case just those."""
    if name is None or (isinstance(name, float) and np.isnan(name)):
        return ""
    name = str(name)
    if str(category).lower() == "surname":
        return name.title()
    return name


def _underscores_to_words(text):
    """Old runs stored prompts as a_person_from_Hawaii — show them as plain English."""
    if text is None or (isinstance(text, float) and np.isnan(text)):
        return ""
    return " ".join(str(text).replace("_", " ").split())


def prompt_sentence(category, prompt_id, census=None, fallback=None):
    """Human-readable prompt, one per (category, prompt_id).

    Surname's three phrasing variants collapse to the same sentence
    (e.g. "A person with the last name Smith").
    """
    name = None
    if census is not None and prompt_id is not None:
        try:
            crow = census.set_index(["category", "prompt_id"]).loc[
                (str(category).lower(), int(prompt_id))
            ]
            if "prompt_name" in crow.index and pd.notna(crow["prompt_name"]):
                name = str(crow["prompt_name"])
        except (KeyError, TypeError, ValueError):
            name = None
    cat = str(category).lower()
    if name:
        if cat == "geography":
            return f"A person from {name}"
        if cat == "surname":
            return f"A person with the last name {name.title()}"
        if cat == "occupation":
            return f"A person that works as {name}"
        return name
    plain = _underscores_to_words(fallback)
    return plain or f"{_proxy_label(category)} {prompt_id}"


def _l_bin_label(lo, hi):
    if lo == 0:
        return f"L* < {hi}"
    if hi >= 101:
        return f"L* \u2265 {lo}"
    return f"L* {lo}\u2013{hi}"


def _model_race_distribution(sub):
    """Fractions of classified images by canonical race for one prompt group."""
    r = _canon_race_series(sub).dropna()
    total = len(r)
    if total == 0:
        return None
    counts = r.value_counts()
    return {race: counts.get(race, 0) / total for race in RACES}


def hue_to_rgb_reference(hue):
    return hsv_to_rgb([hue / 360.0, 1.0, 1.0])


def add_hue_colorbar(ax, x_min, x_max, xlabel="H* (Hue Angle, degrees)"):
    """Reference color bar beneath an L*/H* scatter — no ticks, x-axis label only."""
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("bottom", size="5%", pad=0.35)
    gradient = np.linspace(x_min, x_max, 2000)
    gradient_colors = np.array([hue_to_rgb_reference(h) for h in gradient])
    cax.imshow(gradient_colors.reshape(1, -1, 3), extent=[x_min, x_max, 0, 1], aspect="auto")
    cax.set_xlim(x_min, x_max)
    cax.set_ylim(0, 1)
    cax.set_yticks([])
    cax.set_xticks([])
    cax.set_xlabel(xlabel, fontweight="bold")


def plot_with_uniform_layering(ax, data_dict, x_col, y_col, s=20, rasterized=True, alpha=0.6):
    """Scatter in a fixed race layering order (PLOT_ORDER). Extra groups
    (e.g. Unknown) are drawn underneath so they don't hide classified points."""
    ax.set_facecolor("white")
    handles_labels = []
    extra = [k for k in data_dict if k not in PLOT_ORDER]
    order = extra + list(PLOT_ORDER)
    for idx, race in enumerate(order):
        if race not in data_dict or len(data_dict[race]) == 0:
            continue
        race_data = data_dict[race]
        handle = ax.scatter(
            race_data[x_col], race_data[y_col],
            alpha=alpha, s=s, color=RACE_COLORS.get(race, "#BBBBBB"),
            edgecolors="white", linewidths=0.4, zorder=2 + idx,
            rasterized=rasterized,
        )
        handles_labels.append((handle, race, len(race_data)))
    return handles_labels


# ====================================================================== TVD
TVD_SUMMARY_COLS = ["group", "mean", "ci_low", "ci_high", "std", "n"]


def _bootstrap_mean_ci(values, n_boot=N_BOOTSTRAP, ci=CI_LEVEL, seed=0):
    """Percentile bootstrap CI for the mean of a 1-d sample. Returns (lo, hi)."""
    vals = np.asarray(values, dtype=float)
    vals = vals[np.isfinite(vals)]
    if len(vals) == 0:
        return float("nan"), float("nan")
    rng = np.random.default_rng(seed)
    n = len(vals)
    draws = rng.choice(vals, size=(n_boot, n), replace=True)
    boot_means = draws.mean(axis=1)
    lo_pct = (100 - ci) / 2
    lo, hi = np.percentile(boot_means, [lo_pct, 100 - lo_pct])
    return float(lo), float(hi)


def _tvd_summary_row(group, values):
    values = np.asarray(values, dtype=float)
    seed = zlib.adler32(f"tvd|{group}".encode()) & 0xFFFFFFFF
    lo, hi = _bootstrap_mean_ci(values, seed=seed)
    return {
        "group": group,
        "mean": float(np.mean(values)),
        "ci_low": lo,
        "ci_high": hi,
        "std": float(np.std(values, ddof=1)) if len(values) > 1 else 0.0,
        "n": int(len(values)),
    }


def compute_tvd(df, census):
    """Per-prompt TVD (0.5 * sum|model - census| over races), then summarized.

    Returns (per_prompt_df, summary_df). TVD is computed for each prompt that
    has at least one classified image, matching the paper's per-prompt
    definition. Summary is mean + 95% bootstrap CI (over prompts) per proxy
    category, plus an Overall row.
    """
    census_idx = census.set_index(["category", "prompt_id"])
    rows = []
    for (cat, pid), sub in df.groupby(["category", "prompt_id"]):
        cat = str(cat).lower()
        try:
            pid = int(pid)
        except (TypeError, ValueError):
            continue
        dist = _model_race_distribution(sub)
        if dist is None or (cat, pid) not in census_idx.index:
            continue
        crow = census_idx.loc[(cat, pid)]
        tvd = 0.5 * sum(abs(dist[r] - float(crow[r])) for r in RACES)
        rows.append({"category": cat, "prompt_id": pid, "tvd": tvd})

    per_prompt = pd.DataFrame(rows)
    if per_prompt.empty:
        return per_prompt, pd.DataFrame(columns=TVD_SUMMARY_COLS)

    summary = []
    seen = set()
    for cat in PROXY_ORDER:
        g = per_prompt[per_prompt["category"] == cat]
        if g.empty:
            continue
        summary.append(_tvd_summary_row(cat, g["tvd"].to_numpy()))
        seen.add(cat)
    for cat, g in per_prompt.groupby("category"):
        if cat in seen:
            continue
        summary.append(_tvd_summary_row(str(cat), g["tvd"].to_numpy()))
    summary.append(_tvd_summary_row("Overall", per_prompt["tvd"].to_numpy()))
    summary_df = pd.DataFrame(summary, columns=TVD_SUMMARY_COLS)
    for col in ("mean", "ci_low", "ci_high", "std"):
        summary_df[col] = summary_df[col].round(3)
    summary_df["n"] = summary_df["n"].astype(int)
    return per_prompt, summary_df


def _agg_points(df, census):
    """One row per (prompt, race): census % vs model output %."""
    census_idx = census.set_index(["category", "prompt_id"])
    recs = []
    for (cat, pid), sub in df.groupby(["category", "prompt_id"]):
        cat = str(cat).lower()
        try:
            pid = int(pid)
        except (TypeError, ValueError):
            continue
        dist = _model_race_distribution(sub)
        if dist is None or (cat, pid) not in census_idx.index:
            continue
        crow = census_idx.loc[(cat, pid)]
        prompt_name = crow["prompt_name"] if "prompt_name" in crow.index else f"{cat} {pid}"
        for race in RACES:
            recs.append({
                "category": cat, "prompt_id": pid, "race": race,
                "prompt_name": prompt_name,
                "census_pct": float(crow[race]) * 100.0,
                "model_pct": dist[race] * 100.0,
            })
    return pd.DataFrame(recs)


# ====================================================================== Fig 2.2 / Fig 3
def race_vs_census_scatter(df, census, race_filter=None):
    """Census % (x) vs model output % (y), one point per (prompt, race).

    race_filter=None/'All' -> all primary races, colored by race (Fig 2.2).
    race_filter='Black' (etc.) -> that race only, colored by proxy, with the
    furthest/closest points labeled (Fig 3).
    """
    pts = _agg_points(df, census)
    if pts.empty:
        return _empty_fig("No race-classification data to plot.\nRun Part 2 with 'Race classification' enabled.")

    if race_filter and race_filter != "All":
        return _annotated_race_scatter(pts, race_filter)

    fig, ax = plt.subplots(figsize=(10.5, 7.5))
    ax.set_facecolor("white")
    model_data = pts[pts["race"].isin(PRIMARY_RACES)]
    for race in PRIMARY_RACES:
        race_data = model_data[model_data["race"] == race]
        ax.scatter(
            race_data["census_pct"], race_data["model_pct"],
            alpha=0.6, s=55, c=RACE_COLORS[race], marker="o",
            label=race, edgecolors="white", linewidth=0.5, zorder=2,
            rasterized=True,
        )
    ax.plot([0, 100], [0, 100], "k--", linewidth=2, alpha=0.7, zorder=1)
    ax.grid(True, alpha=0.3, linewidth=0.5)
    ax.set_xlim(0, 100)
    ax.set_ylim(0, 100)
    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel("Census Demographic %", fontweight="bold")
    ax.set_ylabel("Model Output %", fontweight="bold")
    _legend_outside(ax)
    return fig


def _spaced_annotation_rows(df, n, min_dist=8.0):
    """Take up to n rows, skipping any whose (x, y) sits on top of one already kept."""
    picked = []
    pts = []
    for _, row in df.iterrows():
        xy = np.array([row["census_pct"], row["model_pct"]], dtype=float)
        if pts and np.min(np.linalg.norm(np.vstack(pts) - xy, axis=1)) < min_dist:
            continue
        picked.append(row)
        pts.append(xy)
        if len(picked) >= n:
            break
    return picked


def _annotated_race_scatter(pts, race, annotate_top_n=10):
    """Fig 3: one race, colored by proxy, furthest/closest points labeled."""
    model_data = pts[pts["race"] == race].copy()
    if model_data.empty:
        return _empty_fig(f"No {race} classification data to plot.")
    model_data["residual"] = (model_data["model_pct"] - model_data["census_pct"]).abs()

    fig, ax = plt.subplots(figsize=(13, 8.5))
    ax.set_facecolor("white")
    for proxy in PROXY_ORDER:
        proxy_data = model_data[model_data["category"].astype(str).str.lower() == proxy]
        if proxy_data.empty:
            continue
        ax.scatter(
            proxy_data["census_pct"], proxy_data["model_pct"],
            alpha=0.6, s=150, c=_proxy_color(proxy), marker="s",
            label=_proxy_label(proxy), edgecolors="white", linewidth=1, zorder=2,
            rasterized=True,
        )

    furthest = _spaced_annotation_rows(
        model_data.sort_values("residual", ascending=False), annotate_top_n, min_dist=12.0,
    )
    closest = []
    ax.set_xlim(-12, 112)
    ax.set_ylim(-12, 112)

    texts = []
    for row, face in [(r, "white") for r in furthest] + [(r, "#ccffcc") for r in closest]:
        x, y = float(row["census_pct"]), float(row["model_pct"])
        lx, ly = x, y
        if y < 8:
            ly = 14
        elif y > 92:
            ly = 86
        texts.append(ax.annotate(
            _display_prompt_name(row["category"], row["prompt_name"]),
            xy=(x, y), xytext=(lx, ly),
            xycoords="data", textcoords="data",
            fontsize=11, color=_proxy_color(row["category"]), ha="center", va="center",
            bbox=dict(boxstyle="round,pad=0.3", fc=face,
                      ec=_proxy_color(row["category"]), alpha=0.92, linewidth=1.1),
            zorder=6,
        ))

    ax.plot([0, 100], [0, 100], "k--", linewidth=2, alpha=0.7, zorder=1)
    _repel_labels(texts, ax)

    ax.set_xlabel("Census Demographic %", fontweight="bold")
    ax.set_ylabel("Model Output %", fontweight="bold")
    _legend_outside(ax)
    ax.grid(True, alpha=0.3)
    ax.set_aspect("equal", adjustable="box")
    return fig


# ====================================================================== Fig 4
def prompt_tvd(df, census, category, prompt_id):
    """TVD of one prompt vs its census row. None if there is nothing to compare."""
    cat = str(category).lower()
    sub = df[
        (df["category"].astype(str).str.lower() == cat)
        & (df["prompt_id"] == int(prompt_id))
    ]
    dist = _model_race_distribution(sub)
    if dist is None:
        return None
    try:
        crow = census.set_index(["category", "prompt_id"]).loc[(cat, int(prompt_id))]
    except KeyError:
        return None
    return 0.5 * sum(abs(dist[r] - float(crow[r])) for r in RACES)


def prompt_race_bar(df, census, category, prompt_id, count_labels=False):
    """Stacked race breakdown for one prompt: census row under the model's row.

    All phrasing variants that share a prompt_id (surname has three) are pooled.
    count_labels: also print n/N on the model bar (playground, small n).
    """
    cat = str(category).lower()
    sub = df[
        (df["category"].astype(str).str.lower() == cat)
        & (df["prompt_id"] == int(prompt_id))
    ]
    dist = _model_race_distribution(sub)
    if dist is None:
        return _empty_fig("No classified images for this prompt.")

    crow = census.set_index(["category", "prompt_id"]).loc[(cat, int(prompt_id))]
    race_order = ["White", "Black", "Asian", "Hispanic", "Other"]
    races = _canon_race_series(sub).dropna()
    n_total = len(races)
    counts = races.value_counts()
    model_vals = [dist[r] * 100 for r in race_order]
    census_vals = [float(crow[r]) * 100 for r in race_order]

    # Census at the bottom (y=0), model above — same stacking as the paper Fig 4.
    row_labels = ["Census", str(model_name(df))]
    data = np.array([census_vals, model_vals])

    fig, ax = plt.subplots(figsize=(11, 3.6))
    y_pos = np.arange(len(row_labels))
    left = np.zeros(len(row_labels))
    for i, race in enumerate(race_order):
        values = data[:, i]
        bars = ax.barh(
            y_pos, values, left=left, color=RACE_COLORS[race],
            edgecolor="white", linewidth=1.2, label=race, height=0.75,
        )
        n_race = int(counts.get(race, 0))
        for row_i, (bar, val, start) in enumerate(zip(bars, values, left)):
            if val <= 6:
                continue
            weight = "bold" if val >= 10 else "normal"
            if count_labels and row_i == 1 and n_total:
                label = f"{n_race}/{n_total}"
            else:
                label = f"{val:.0f}%"
            ax.text(
                start + val / 2, bar.get_y() + bar.get_height() / 2,
                label, ha="center", va="center", fontsize=13,
                fontweight=weight, color="#2d2d2d", clip_on=True,
            )
        left += values

    ax.set_yticks(y_pos)
    ax.set_yticklabels(row_labels)
    ax.set_xlim(0, 100)
    ax.set_xlabel("Percentage (%)", labelpad=8)
    fallback = None
    if "prompt_text" in sub.columns and len(sub):
        texts = sub["prompt_text"].dropna().astype(str)
        if len(texts):
            fallback = max(texts, key=lambda t: len(_underscores_to_words(t)))
    title = prompt_sentence(category, prompt_id, census, fallback=fallback)
    ax.set_title(f'"{title}"', fontweight="bold", pad=12)
    ax.tick_params(axis="y", length=0, pad=8)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_visible(False)
    ax.spines["bottom"].set_color("#666666")
    ax.set_axisbelow(True)
    ax.grid(axis="x", alpha=0.2, linestyle="--", linewidth=0.5, color="#999999")

    handles = [
        Rectangle((0, 0), 1, 1, facecolor=RACE_COLORS[r], edgecolor="white", linewidth=1.2)
        for r in race_order
    ]
    # Below the bars so it cannot collide with the prompt title.
    leg = fig.legend(
        handles, race_order, loc="upper center", ncol=5, bbox_to_anchor=(0.5, -0.12),
        columnspacing=1.5, handletextpad=0.5, handlelength=1.5,
    )
    style_legend(leg)
    fig.subplots_adjust(bottom=0.28)
    return fig


# ====================================================================== Fig 1
def _ensure_h_star(df):
    """Fill missing H* from a*/b*. Does not drop or quality-filter any rows."""
    sub = df.copy()
    if "h_star" not in sub.columns:
        sub["h_star"] = np.nan
    a_col = "a" if "a" in sub.columns else ("A" if "A" in sub.columns else None)
    b_col = "b" if "b" in sub.columns else ("B" if "B" in sub.columns else None)
    if a_col and b_col:
        missing = sub["h_star"].isna()
        if missing.any():
            h = np.degrees(np.arctan2(sub.loc[missing, b_col], sub.loc[missing, a_col])) % 360
            sub.loc[missing, "h_star"] = h
    return sub


def normalize_viz_df(df):
    """Map older/HF column names onto the schema the charts expect."""
    if df is None:
        return df
    out = df.copy()
    aliases = {
        "race_classification": ("Race", "race"),
        "category": ("proxy_cat", "proxy"),
        "prompt_id": ("promptID", "promptId"),
        "prompt_text": ("prompt",),
        "a": ("A",),
        "b": ("B",),
        "h_star": ("H", "h", "H_star"),
        "L": ("l", "L_star"),
    }
    for dest, srcs in aliases.items():
        if dest in out.columns:
            continue
        for src in srcs:
            if src in out.columns:
                out[dest] = out[src]
                break
    if "category" in out.columns:
        out["category"] = out["category"].astype(str).str.lower()
    if "prompt_id" in out.columns:
        out["prompt_id"] = pd.to_numeric(out["prompt_id"], errors="coerce")
    return _ensure_h_star(out)


def skin_tone_colorspace(df, marker_size=20, rasterized=True):
    """L* vs H* skin-tone cloud, colored by race, with a hue reference bar.

    Every image with a CIELAB reading is plotted — no face-detection /
    embedding-quality filter (those apply only to homogeneity charts).
    Unclassified race is kept as Unknown.
    """
    if df is None or "L" not in df.columns:
        return _empty_fig("No skin-tone data to plot.\nRun Part 2 with 'Skin tone' enabled.")
    sub = _ensure_h_star(df)
    # Include all images that have lightness + hue. Do not filter on
    # face_det_score, n_faces_detected, or embed_status.
    sub = sub[sub["L"].notna() & sub["h_star"].notna()].copy()
    if sub.empty:
        return _empty_fig("No skin-tone data to plot.\nRun Part 2 with 'Skin tone' enabled.")
    if "race_classification" in sub.columns:
        sub["race"] = _canon_race_series(sub).fillna("Unknown")
    else:
        sub["race"] = "Unknown"

    h = sub["h_star"].to_numpy()
    h_2nd, h_98th = np.percentile(h, [2, 98])
    h_range = max(h_98th - h_2nd, 5.0)
    x_margin = h_range * 0.08
    x_min = np.floor(max(0, h_2nd - x_margin) / 5) * 5
    x_max = np.ceil(min(360.0, h_98th + x_margin) / 5) * 5
    if x_max <= x_min:
        x_max = x_min + 25

    fig, ax = plt.subplots(figsize=(9, 7.5))
    data_by_race = {race: sub[sub["race"] == race] for race in RACE_COLORS}
    handles_labels = plot_with_uniform_layering(
        ax, data_by_race, x_col="h_star", y_col="L",
        s=marker_size, rasterized=rasterized,
        alpha=0.85 if marker_size >= 50 else 0.6,
    )

    ax.set_ylabel("L* (Lightness)", fontweight="bold")
    hl_by_race = {race: (hnd, n) for hnd, race, n in handles_labels}
    legend_order = LEGEND_RACE_ORDER + [r for r in hl_by_race if r not in LEGEND_RACE_ORDER]
    ordered = [(race, hl_by_race[race]) for race in legend_order if race in hl_by_race]
    handles = [hnd for _, (hnd, n) in ordered]
    labels = [f"{race} (n={n})" for race, (hnd, n) in ordered]
    if handles:
        _legend_above(ax, handles, labels, ncol=min(len(labels), 3))

    ax.grid(True, alpha=0.25, linewidth=0.5, zorder=1)
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(0, 100)
    for spine in ax.spines.values():
        spine.set_edgecolor("black")
        spine.set_linewidth(1)
    add_hue_colorbar(ax, x_min, x_max)
    ax.set_xlabel("")
    return fig


def skin_tone_strip(df):
    """Playground-scale skin-tone chart: each face as a point on the L* axis."""
    from .skin_tone import lab_to_hex

    if df is None or "L" not in df.columns:
        return _empty_fig("No skin-tone data to plot.")
    sub = df[df["L"].notna()].copy()
    if sub.empty:
        return _empty_fig("No face detected for skin-tone measurement.")

    n_missing = len(df) - len(sub)
    lo = float(sub["L"].min())
    hi = float(sub["L"].max())

    fig, ax = plt.subplots(figsize=(9.5, 2.8))
    ax.set_xlim(0, 100)
    ax.set_ylim(-0.7, 0.7)
    ax.axhline(0, color="#B0B0B0", lw=1.3, zorder=1)
    ax.plot([lo, hi], [0, 0], color="#555555", lw=2.2, solid_capstyle="round", zorder=2)
    for _, row in sub.iterrows():
        try:
            hex_color = lab_to_hex(row["L"], row["a"], row["b"])
        except Exception:
            hex_color = "#CCCCCC"
        ax.scatter(
            float(row["L"]), 0, s=280, c=hex_color,
            edgecolors="#222222", linewidths=0.8, zorder=3,
        )
    ax.set_yticks([])
    ax.set_xlabel("L*  (0 = black, 100 = white)", fontweight="bold")
    ax.set_xticks([0, 25, 50, 75, 100])
    apply_paper_style(ax)
    ax.spines["left"].set_visible(False)
    ax.tick_params(axis="y", length=0)
    spread = hi - lo
    ax.text(
        50, -0.55,
        f"Lightness span  {lo:.0f}–{hi:.0f}   (spread {spread:.0f})"
        + (f"  ·  {n_missing} with no face" if n_missing else ""),
        ha="center", va="top", fontsize=13, color="#444444",
    )
    return fig


# ====================================================================== embeddings
def _stack_embeddings(df, emb_col, dim, require_face=False):
    """Return (matrix (n,dim), sub_df) for rows with a valid embedding.

    require_face matches the paper's homogeneity subset: detection score
    >= 0.85 and a single face (n_faces_detected == 1 when that column exists).
    """
    sub = df[df[emb_col].notna()].copy()
    if require_face:
        if "embed_status" in sub.columns:
            sub = sub[sub["embed_status"].isin(["ok"])]
        if "face_det_score" in sub.columns:
            sub = sub[sub["face_det_score"].fillna(0) >= FACE_DET_MIN]
        if "n_faces_detected" in sub.columns:
            sub = sub[sub["n_faces_detected"] == 1]
    if sub.empty:
        return np.empty((0, dim)), sub
    vecs = []
    keep = []
    for idx, v in sub[emb_col].items():
        arr = np.asarray(v, dtype=np.float32)
        if arr.shape == (dim,) and not np.all(np.isnan(arr)):
            vecs.append(arr)
            keep.append(idx)
    if not vecs:
        return np.empty((0, dim)), sub.loc[[]]
    mat = np.vstack(vecs)
    mat = mat / (np.linalg.norm(mat, axis=1, keepdims=True) + 1e-9)
    return mat, sub.loc[keep]


EMB_DIMS = {"face_emb": 512, "clip_emb": 1024}


def _pairwise_sims(mat):
    """Upper-triangle pairwise cosine similarities (rows assumed unit-norm)."""
    n = len(mat)
    return (mat @ mat.T)[np.triu_indices(n, k=1)]


def _bootstrap_ci_mean(embeddings, n_boot=N_BOOTSTRAP, ci=CI_LEVEL, rng=None):
    """Bootstrap CI for the mean pairwise cosine similarity within a set."""
    if rng is None:
        rng = np.random.default_rng(0)
    n = len(embeddings)
    boot_means = np.empty(n_boot)
    for b in range(n_boot):
        idx = rng.choice(n, n, replace=True)
        boot_means[b] = _pairwise_sims(embeddings[idx]).mean()
    lo_pct = (100 - ci) / 2
    hi_pct = 100 - lo_pct
    return np.percentile(boot_means, [lo_pct, hi_pct]), float(boot_means.mean())


def homogeneity_violin(df, emb_col="face_emb"):
    """Fig 10: box + strip + bootstrap CI of pairwise cosine sim by L* bin.

    Equal-n within the model: each plotted bin is randomly subsampled to the
    size of the smallest eligible bin (n >= MIN_BIN_SIZE), so bin comparisons
    aren't confounded by uneven sample sizes.
    """
    dim = EMB_DIMS.get(emb_col, 512)
    mat, sub = _stack_embeddings(df, emb_col, dim, require_face=True)
    if len(sub) == 0 or "L" not in sub.columns or sub["L"].notna().sum() == 0:
        return _empty_fig(
            f"No usable {emb_col} + skin-tone data.\n"
            f"(CLIP embeddings aren't present yet.)" if emb_col == "clip_emb"
            else f"No usable {emb_col} + skin-tone data to plot."
        )

    L = sub["L"].to_numpy()
    bin_indices = []
    for (lo, hi), color in zip(L_BINS, SKIN_PALETTE):
        idx = np.where((L >= lo) & (L < hi))[0]
        bin_indices.append((idx, _l_bin_label(lo, hi), color))

    eligible = [len(idx) for idx, _, _ in bin_indices if len(idx) >= MIN_BIN_SIZE]
    if not eligible:
        return _empty_fig(
            f"Not enough images per skin-tone bin for a homogeneity plot.\n"
            f"Need >={MIN_BIN_SIZE} images in a bin."
        )
    equal_n = min(eligible)
    if equal_n < 2:
        return _empty_fig(
            f"Not enough images per skin-tone bin for a homogeneity plot.\n"
            f"Need >={MIN_BIN_SIZE} images in a bin."
        )

    rng = np.random.default_rng(42)
    strip_rng = np.random.default_rng(1)
    plot_data, plot_means, valid_labels, valid_colors = [], [], [], []
    ci_los, ci_his = [], []

    for idx, lab, color in bin_indices:
        if len(idx) < equal_n:
            continue
        chosen = rng.choice(idx, equal_n, replace=False)
        block = mat[chosen]
        sims = _pairwise_sims(block)
        boot_seed = zlib.adler32(f"{emb_col}|{lab}|{equal_n}".encode()) & 0xFFFFFFFF
        boot_rng = np.random.default_rng(boot_seed)
        (ci_lo, ci_hi), _ = _bootstrap_ci_mean(block, rng=boot_rng)

        plot_data.append(sims)
        plot_means.append(float(sims.mean()))
        valid_labels.append(lab)
        valid_colors.append(color)
        ci_los.append(ci_lo)
        ci_his.append(ci_hi)

    if not plot_data:
        return _empty_fig(
            f"Not enough images per skin-tone bin for a homogeneity plot.\n"
            f"Need >={MIN_BIN_SIZE} images in a bin."
        )

    all_sims = np.concatenate(plot_data)
    y_lo = min(float(all_sims.min()), float(min(ci_los)))
    y_hi = max(float(all_sims.max()), float(max(ci_his)))
    y_pad = (y_hi - y_lo) * 0.10 or 0.02
    y_lim = (y_lo - y_pad, y_hi + y_pad)

    fig, ax = plt.subplots(figsize=(8.5, 7.5))
    fig.patch.set_facecolor("white")
    ax.set_facecolor("white")
    positions = list(range(len(plot_data)))

    for i, sims_arr in enumerate(plot_data):
        jitter = strip_rng.uniform(-0.10, 0.10, size=len(sims_arr))
        ax.scatter(
            np.full(len(sims_arr), positions[i]) + jitter,
            sims_arr,
            color=valid_colors[i], alpha=0.12, s=8,
            linewidths=0, zorder=2, rasterized=True,
        )

    bp = ax.boxplot(
        plot_data, positions=positions, widths=0.45,
        patch_artist=True, showfliers=False, zorder=4,
        medianprops=dict(color="#111111", linewidth=1.8),
        boxprops=dict(edgecolor="#111111", linewidth=1.0),
        whiskerprops=dict(color="#111111", linewidth=1.0),
        capprops=dict(color="#111111", linewidth=1.0),
    )
    for patch, color in zip(bp["boxes"], valid_colors):
        patch.set_facecolor(color)
        patch.set_alpha(0.55)

    ci_offset = [p + 0.32 for p in positions]
    err_lower = [max(0.0, m - lo) for m, lo in zip(plot_means, ci_los)]
    err_upper = [max(0.0, hi - m) for m, hi in zip(plot_means, ci_his)]
    ax.errorbar(
        ci_offset, plot_means,
        yerr=[err_lower, err_upper],
        fmt="D", color="#111111",
        ecolor="#111111", elinewidth=1.8,
        capsize=5, capthick=1.8,
        markersize=8, markerfacecolor="white", markeredgewidth=1.6,
        zorder=7,
    )

    ax.set_ylim(y_lim)
    ax.set_xlim(-0.6, len(positions) - 1 + 0.6)
    ax.set_xticks(positions)
    ax.set_xticklabels(valid_labels, rotation=15, ha="right", fontweight="bold")
    ax.set_xlabel("\u2190 Darker                    Lighter \u2192", fontweight="bold")
    ax.set_ylabel("Cosine Similarity", fontweight="bold")
    for tick_label in ax.get_yticklabels():
        tick_label.set_fontweight("bold")
    legend_elems = [
        Line2D([0], [0], marker="D", color="none", markerfacecolor="white",
               markeredgecolor="#111111", markersize=9,
               label=f"Mean \u00b1 {CI_LEVEL}% bootstrap CI"),
    ]
    _legend_above(ax, handles=legend_elems, ncol=1)
    apply_paper_style(ax)
    ax.grid(axis="y", color="#DDDDDD", linewidth=0.7)
    ax.set_axisbelow(True)
    return fig


def _equal_n_race_gaps(mat, races, emb_col="face_emb"):
    """Intra−inter cosine gap per race, with equal-n sampling within the model.

    Each eligible race (n >= MIN_BIN_SIZE) is subsampled to the size of the
    smallest eligible race so gap comparisons aren't confounded by uneven N.
    Returns (gaps_dict, equal_n) where gaps_dict maps race -> gap float.
    """
    race_indices = {}
    for race in RACES:
        idx = np.where(races == race)[0]
        if len(idx) >= MIN_BIN_SIZE:
            race_indices[race] = idx

    if len(race_indices) < 2:
        return {}, 0

    equal_n = min(len(idx) for idx in race_indices.values())
    if equal_n < 2:
        return {}, 0

    sampled = {}
    for race, idx in race_indices.items():
        seed = zlib.adler32(f"{emb_col}|gap|{race}|{equal_n}".encode()) & 0xFFFFFFFF
        race_rng = np.random.default_rng(seed)
        sampled[race] = race_rng.choice(idx, equal_n, replace=False)

    gaps = {}
    for race, idx in sampled.items():
        others = np.concatenate([sampled[r] for r in sampled if r != race])
        if len(others) < 1:
            continue
        iu = np.triu_indices(equal_n, k=1)
        block = mat[idx]
        intra = float((block @ block.T)[iu].mean())
        inter = float((mat[idx] @ mat[others].T).mean())
        gaps[race] = intra - inter
    return gaps, equal_n


def homogeneity_gap_bars(df, emb_col="face_emb"):
    """Fig 8 (single model): intra-minus-inter mean cosine similarity, per race.

    Positive gap = that race's images are more self-similar than they are
    similar to other races (a signal of less-diverse / more stereotyped output).

    Uses equal-n sampling: every plotted race is randomly subsampled to the
    size of the smallest eligible race (n >= MIN_BIN_SIZE).
    """
    dim = EMB_DIMS.get(emb_col, 512)
    mat, sub = _stack_embeddings(df, emb_col, dim, require_face=True)
    if len(sub) == 0:
        return _empty_fig(
            f"No usable {emb_col} data.\n"
            "(CLIP embeddings aren't present yet.)" if emb_col == "clip_emb"
            else f"No usable {emb_col} data to plot."
        )

    races = _canon_race_series(sub).to_numpy()
    gaps_dict, equal_n = _equal_n_race_gaps(mat, races, emb_col=emb_col)
    if not gaps_dict:
        return _empty_fig(
            f"Not enough images per race for a homogeneity gap.\n"
            f"Need >={MIN_BIN_SIZE} images in at least two race groups."
        )

    plot_races = [r for r in FIG8_RACE_ORDER if r in gaps_dict]
    gap_data = [gaps_dict[r] for r in plot_races]
    x = np.arange(len(plot_races))
    bar_width = 0.55

    fig, ax = plt.subplots(figsize=(9, 7))
    ax.set_facecolor("white")
    for i, race in enumerate(plot_races):
        ax.bar(
            x[i], gap_data[i], width=bar_width, color=RACE_COLORS[race],
            label=race, edgecolor="white", linewidth=0.5, zorder=3,
        )

    y_min_data, y_max_data = min(gap_data), max(gap_data)
    y_range = (y_max_data - y_min_data) or 0.05
    ax.set_ylim(y_min_data - y_range * 0.1, y_max_data + y_range * 0.12)
    ax.set_xticks(x)
    ax.set_xticklabels(plot_races, fontweight="bold")
    ax.set_ylabel("Homogeneity Gap (intra \u2212 inter)")
    ax.axhline(0, color="#AAAAAA", linewidth=0.8, linestyle="--")
    handles = [Patch(facecolor=RACE_COLORS[r], label=r) for r in plot_races]
    _legend_above(ax, handles=handles, title="Race", ncol=len(plot_races))
    apply_paper_style(ax)
    ax.grid(axis="y", color="#EEEEEE", linewidth=0.8)
    ax.set_axisbelow(True)
    return fig


# ====================================================== cross-model comparison
# Kept for callers that pool several models; the Gradio UI does not use these.
def _models_in(df):
    return sorted(df["model"].dropna().unique())


def race_vs_census_grid(df, census):
    """Census-vs-model scatter, one panel per model (not shown in the UI)."""
    models = _models_in(df)
    if not models:
        return _empty_fig("No race-classification data to compare.")
    census_idx = census.set_index(["category", "prompt_id"])
    fig, axes = plt.subplots(1, len(models), figsize=(7.0 * len(models), 7.5), squeeze=False)
    for ax, model in zip(axes[0], models):
        sub = df[df["model"] == model]
        pts = {r: ([], []) for r in PRIMARY_RACES}
        for (cat, pid), g in sub.groupby(["category", "prompt_id"]):
            dist = _model_race_distribution(g)
            if dist is None or (cat, pid) not in census_idx.index:
                continue
            crow = census_idx.loc[(cat, pid)]
            for r in PRIMARY_RACES:
                pts[r][0].append(float(crow[r]) * 100)
                pts[r][1].append(dist[r] * 100)
        ax.plot([0, 100], [0, 100], "k--", linewidth=2, alpha=0.7, zorder=1)
        for r in PRIMARY_RACES:
            ax.scatter(
                pts[r][0], pts[r][1], s=55, alpha=0.6, color=RACE_COLORS[r],
                marker="o", edgecolors="white", linewidth=0.5, zorder=2,
            )
        ax.set_xlim(0, 100)
        ax.set_ylim(0, 100)
        ax.set_aspect("equal", adjustable="box")
        ax.set_title(str(model), fontweight="bold", pad=12)
        ax.set_xlabel("Census Demographic %", fontweight="bold")
        ax.grid(True, alpha=0.3, linewidth=0.5)
    axes[0][0].set_ylabel("Model Output %", fontweight="bold")
    handles = [
        Line2D([0], [0], marker="o", color="w", markerfacecolor=RACE_COLORS[r],
               markeredgecolor="white", markersize=10, label=r)
        for r in PRIMARY_RACES
    ]
    leg = fig.legend(handles=handles, loc="upper center", ncol=4, bbox_to_anchor=(0.5, 1.02))
    style_legend(leg)
    fig.tight_layout()
    return fig


def tvd_comparison_bars(df, census):
    """Mean TVD per model, by proxy (not shown in the UI)."""
    models = _models_in(df)
    if not models:
        return _empty_fig("No race-classification data to compare.")
    groups = ["geography", "occupation", "surname", "Overall"]
    per_model = {}
    for m in models:
        _, summary = compute_tvd(df[df["model"] == m], census)
        per_model[m] = {row["group"]: row["mean"] for _, row in summary.iterrows()}

    x = np.arange(len(groups))
    w = 0.8 / len(models)
    fig, ax = plt.subplots(figsize=(9, 6))
    ax.set_facecolor("white")
    for i, m in enumerate(models):
        vals = [per_model[m].get(g, np.nan) for g in groups]
        ax.bar(x + i * w, vals, w, label=m, edgecolor="white", linewidth=0.5)
    ax.set_xticks(x + w * (len(models) - 1) / 2)
    ax.set_xticklabels([g.capitalize() for g in groups], fontweight="bold")
    ax.set_ylabel("Mean TVD vs census (lower = closer)")
    leg = ax.legend(loc="upper right")
    style_legend(leg)
    apply_paper_style(ax)
    ax.grid(axis="y", color="#EEEEEE", linewidth=0.8)
    ax.set_axisbelow(True)
    fig.tight_layout()
    return fig


def skin_tone_grid(df):
    """L* vs H* skin-tone, one panel per model (not shown in the UI)."""
    models = _models_in(df)
    sub_all = df[df["L"].notna() & df["h_star"].notna()]
    if sub_all.empty or not models:
        return _empty_fig("No skin-tone data to compare.")
    h = sub_all["h_star"].to_numpy()
    h_2nd, h_98th = np.percentile(h, [2, 98])
    h_range = h_98th - h_2nd
    x_margin = h_range * 0.05
    x_min = np.floor(max(0, h_2nd - x_margin) / 5) * 5
    x_max = 70
    fig, axes = plt.subplots(1, len(models), figsize=(7.0 * len(models), 7.5), squeeze=False)
    for ax, model in zip(axes[0], models):
        sub = sub_all[sub_all["model"] == model].copy()
        sub["race"] = _canon_race_series(sub)
        data_by_race = {race: sub[sub["race"] == race] for race in RACE_COLORS}
        plot_with_uniform_layering(ax, data_by_race, x_col="h_star", y_col="L")
        ax.set_xlim(x_min, x_max)
        ax.set_ylim(0, 100)
        ax.set_title(str(model), fontweight="bold", pad=10)
        ax.set_ylabel("L* (Lightness)", fontweight="bold")
        ax.grid(True, alpha=0.25, linewidth=0.5, zorder=1)
        for spine in ax.spines.values():
            spine.set_edgecolor("black")
            spine.set_linewidth(1)
        add_hue_colorbar(ax, x_min, x_max)
    fig.tight_layout()
    return fig


def homogeneity_gap_grouped(df, emb_col="face_emb"):
    """Homogeneity gap per race, grouped bars by model (not shown in the UI)."""
    models = _models_in(df)
    dim = EMB_DIMS.get(emb_col, 512)
    gaps = {}
    equal_ns = {}
    for m in models:
        mat, sub = _stack_embeddings(df[df["model"] == m], emb_col, dim, require_face=True)
        if len(sub) == 0:
            continue
        races = _canon_race_series(sub).to_numpy()
        d, equal_n = _equal_n_race_gaps(mat, races, emb_col=f"{emb_col}|{m}")
        if d:
            gaps[m] = d
            equal_ns[m] = equal_n
    if not gaps:
        kind_note = "(CLIP embeddings aren't present yet.)" if emb_col == "clip_emb" else ""
        return _empty_fig(f"No usable {emb_col} data to compare across models.\n{kind_note}")

    races_present = [r for r in FIG8_RACE_ORDER if any(r in gaps[m] for m in gaps)]
    models_present = [m for m in models if m in gaps]
    x = np.arange(len(models_present))
    n_races = len(races_present)
    bar_width = 0.15
    fig, ax = plt.subplots(figsize=(11, 7.5))
    ax.set_facecolor("white")
    for i, race in enumerate(races_present):
        gap_data = [gaps[m].get(race, 0) for m in models_present]
        offset = (i - n_races / 2 + 0.5) * (bar_width + 0.01)
        ax.bar(
            x + offset, gap_data, width=bar_width, color=RACE_COLORS[race],
            label=race, edgecolor="white", linewidth=0.5, zorder=3,
        )
    ax.set_xticks(x)
    ax.set_xticklabels([str(m) for m in models_present], fontweight="bold")
    ax.set_ylabel("Homogeneity Gap (intra \u2212 inter)")
    ax.axhline(0, color="#AAAAAA", linewidth=0.8, linestyle="--")
    handles = [Patch(facecolor=RACE_COLORS[r], label=r) for r in races_present]
    leg = ax.legend(
        handles=handles, title="Race", loc="upper center",
        ncol=len(races_present), bbox_to_anchor=(0.5, 1.0),
    )
    style_legend(leg)
    apply_paper_style(ax)
    ax.grid(axis="y", color="#EEEEEE", linewidth=0.8)
    ax.set_axisbelow(True)
    fig.tight_layout()
    return fig