python
from pathlib import Path
import py_compile

code = r'''#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Batch visualization for client_scaling_20261031_224126.

What this script does
---------------------
1. Scan the 6 CSV files under client_scaling_20261031_224126/.
2. Generate one publication-quality PDF for every CSV:
      Training Loss / AUROC / AUPRC vs. Communication Round
3. Extract run-level statistics and write:
      figures/client_scaling_summary.csv
4. Aggregate results by number of clients and generate an effect-style summary:
      figures/effect.pdf
      figures/effect.png

Required metric columns
-----------------------
Round, Train_Loss, AUROC, AUPRC

Number of clients is inferred in this order:
1. CSV column such as: num_clients / Num_Clients / N
2. MANUAL_N_CLIENTS mapping below
3. filename pattern such as: n3, n6, n12, clients_24, clients24

Selected clients is inferred similarly. If unavailable, it defaults to 1.
"""

import csv
import math
import re
from collections import defaultdict
from pathlib import Path
from statistics import mean, stdev

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker


# ---------------------------------------------------------------------
# Paths / experiment settings
# ---------------------------------------------------------------------
HERE = Path(__file__).resolve().parent
DATA_DIR = HERE / "client_scaling_20261031_224126"

OUT_DIR = DATA_DIR / "figures"
INDIVIDUAL_DIR = OUT_DIR / "individual"

SUMMARY_CSV = OUT_DIR / "client_scaling_summary.csv"
EFFECT_PDF = OUT_DIR / "effect.pdf"
EFFECT_PNG = OUT_DIR / "effect.png"

EXPECTED_CSV_COUNT = 6

# Match the behavior of your previous effect.py:
# "final"     -> use the final communication round
# "tail_mean" -> average the last TAIL_K communication rounds
SUMMARY_MODE = "final"
TAIL_K = 5

# Preferred baseline. If N=3 is not present, the smallest N is used.
BASELINE_N = 3

# If filenames do not contain n3/n6/... and the CSV has no metadata column,
# fill these dictionaries, for example:
# MANUAL_N_CLIENTS = {
#     "pflad_scaling_01_results.csv": 3,
#     "pflad_scaling_02_results.csv": 6,
# }
MANUAL_N_CLIENTS = {}

# Optional override. Otherwise selected clients defaults to 1.
MANUAL_SELECTED_CLIENTS = {}


# ---------------------------------------------------------------------
# Plot style
# ---------------------------------------------------------------------
plt.rcParams.update({
    "font.size": 11,
    "font.family": "serif",
    "font.serif": ["Times New Roman", "DejaVu Serif"],
    "axes.labelsize": 12,
    "axes.titlesize": 12,
    "xtick.labelsize": 10,
    "ytick.labelsize": 10,
    "legend.fontsize": 9.5,
    "lines.linewidth": 1.8,
    "axes.linewidth": 0.8,
    "figure.dpi": 300,
    "savefig.dpi": 300,
    "savefig.bbox": "tight",
    "savefig.pad_inches": 0.05,
    "pdf.fonttype": 42,
    "ps.fonttype": 42,
})

COLOR_LOSS = "#377eb8"
COLOR_AUROC = "#e41a1c"
COLOR_AUPRC = "#4daf4a"

COLOR_GAP_POS = "#2166ac"
COLOR_GAP_NEG = "#b2182b"
COLOR_BASELINE = "#4d4d4d"
COLOR_GAP_BOX = "#f4b6b6"


# ---------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------
def _find_key(fieldnames, candidates):
    """Case-insensitive column lookup."""
    if not fieldnames:
        return None
    lookup = {name.strip().lower(): name for name in fieldnames}
    for candidate in candidates:
        key = lookup.get(candidate.lower())
        if key is not None:
            return key
    return None


def _as_int(value):
    if value is None:
        return None
    text = str(value).strip()
    if not text:
        return None
    return int(float(text))


def infer_num_clients(csv_path, first_row, fieldnames):
    """Infer N from CSV metadata, manual mapping, or filename."""
    key = _find_key(
        fieldnames,
        ["num_clients", "number_of_clients", "clients", "n_clients", "N"],
    )
    if key is not None:
        value = _as_int(first_row.get(key))
        if value is not None:
            return value

    if csv_path.name in MANUAL_N_CLIENTS:
        return int(MANUAL_N_CLIENTS[csv_path.name])

    stem = csv_path.stem.lower()
    patterns = [
        r"(?:^|[_-])n[_-]?(\d+)(?:[_-]|$)",
        r"(?:^|[_-])clients?[_-]?(\d+)(?:[_-]|$)",
        r"(?:^|[_-])num[_-]?clients?[_-]?(\d+)(?:[_-]|$)",
    ]
    for pattern in patterns:
        match = re.search(pattern, stem)
        if match:
            return int(match.group(1))

    raise ValueError(
        f"Cannot infer number of clients from {csv_path.name}. "
        "Add num_clients to the CSV, include n3/n6/... in the filename, "
        "or fill MANUAL_N_CLIENTS."
    )


def infer_num_selected(csv_path, first_row, fieldnames):
    """Infer selected clients; default to 1 when not encoded."""
    key = _find_key(
        fieldnames,
        [
            "num_selected_clients",
            "selected_clients",
            "num_selected",
            "clients_per_round",
            "k",
        ],
    )
    if key is not None:
        value = _as_int(first_row.get(key))
        if value is not None:
            return value

    if csv_path.name in MANUAL_SELECTED_CLIENTS:
        return int(MANUAL_SELECTED_CLIENTS[csv_path.name])

    stem = csv_path.stem.lower()
    patterns = [
        r"(?:^|[_-])selected[_-]?(\d+)(?:[_-]|$)",
        r"(?:^|[_-])sel[_-]?(\d+)(?:[_-]|$)",
        r"(?:^|[_-])k[_-]?(\d+)(?:[_-]|$)",
        r"(?:^|[_-])m[_-]?(\d+)(?:[_-]|$)",
    ]
    for pattern in patterns:
        match = re.search(pattern, stem)
        if match:
            return int(match.group(1))

    return 1


def read_run(csv_path):
    """Read one run and normalize the required columns."""
    with csv_path.open(newline="", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        fieldnames = reader.fieldnames or []

        round_key = _find_key(fieldnames, ["Round", "communication_round", "round"])
        loss_key = _find_key(fieldnames, ["Train_Loss", "train_loss", "loss"])
        auroc_key = _find_key(fieldnames, ["AUROC", "auroc"])
        auprc_key = _find_key(fieldnames, ["AUPRC", "auprc"])

        required = {
            "Round": round_key,
            "Train_Loss": loss_key,
            "AUROC": auroc_key,
            "AUPRC": auprc_key,
        }
        missing = [name for name, key in required.items() if key is None]
        if missing:
            raise ValueError(
                f"{csv_path.name}: missing required column(s): {', '.join(missing)}. "
                f"Available columns: {fieldnames}"
            )

        raw_rows = list(reader)

    if not raw_rows:
        raise ValueError(f"{csv_path.name}: CSV is empty.")

    first_row = raw_rows[0]
    num_clients = infer_num_clients(csv_path, first_row, fieldnames)
    num_selected = infer_num_selected(csv_path, first_row, fieldnames)

    rows = []
    for row in raw_rows:
        try:
            rows.append({
                "round": int(float(row[round_key])),
                "loss": float(row[loss_key]),
                "auroc": float(row[auroc_key]),
                "auprc": float(row[auprc_key]),
            })
        except (TypeError, ValueError):
            # Skip malformed / incomplete rows instead of crashing on one row.
            continue

    if not rows:
        raise ValueError(f"{csv_path.name}: no valid metric rows found.")

    rows.sort(key=lambda r: r["round"])

    return {
        "path": csv_path,
        "name": csv_path.name,
        "stem": csv_path.stem,
        "num_clients": num_clients,
        "num_selected": num_selected,
        "participation_ratio": num_selected / num_clients,
        "rows": rows,
    }


def summarize_values(values):
    """Return final value, tail mean, and best value."""
    tail = values[-min(TAIL_K, len(values)):]
    return {
        "final": values[-1],
        "tail_mean": mean(tail),
        "best_max": max(values),
    }


def summarize_run(run):
    rows = run["rows"]
    losses = [r["loss"] for r in rows]
    aurocs = [r["auroc"] for r in rows]
    auprcs = [r["auprc"] for r in rows]

    loss_stats = summarize_values(losses)
    auroc_stats = summarize_values(aurocs)
    auprc_stats = summarize_values(auprcs)

    return {
        "file": run["name"],
        "num_clients": run["num_clients"],
        "num_selected_clients": run["num_selected"],
        "participation_ratio": run["participation_ratio"],
        "num_rounds": len(rows),
        "last_round": rows[-1]["round"],
        "final_train_loss": loss_stats["final"],
        "final_auroc": auroc_stats["final"],
        "final_auprc": auprc_stats["final"],
        f"mean_last_{TAIL_K}_train_loss": loss_stats["tail_mean"],
        f"mean_last_{TAIL_K}_auroc": auroc_stats["tail_mean"],
        f"mean_last_{TAIL_K}_auprc": auprc_stats["tail_mean"],
        "best_auroc": auroc_stats["best_max"],
        "best_auprc": auprc_stats["best_max"],
    }


def choose_metric(summary_row, metric):
    if SUMMARY_MODE == "final":
        return summary_row[f"final_{metric}"]
    if SUMMARY_MODE == "tail_mean":
        return summary_row[f"mean_last_{TAIL_K}_{metric}"]
    raise ValueError("SUMMARY_MODE must be 'final' or 'tail_mean'.")


# ---------------------------------------------------------------------
# Per-CSV figure
# ---------------------------------------------------------------------
def plot_individual_run(run):
    rows = run["rows"]
    rounds = [r["round"] for r in rows]
    loss = [r["loss"] for r in rows]
    auroc = [r["auroc"] for r in rows]
    auprc = [r["auprc"] for r in rows]

    fig, axes = plt.subplots(1, 3, figsize=(10.5, 3.2))

    panels = [
        (axes[0], loss, "Training Loss", COLOR_LOSS, "o", "(a)"),
        (axes[1], auroc, "AUROC", COLOR_AUROC, "s", "(b)"),
        (axes[2], auprc, "AUPRC", COLOR_AUPRC, "^", "(c)"),
    ]

    for ax, values, ylabel, color, marker, label in panels:
        ax.plot(
            rounds,
            values,
            color=color,
            marker=marker,
            markersize=4.5,
            markerfacecolor="white",
            markeredgewidth=1.2,
            markevery=1,
            alpha=0.95,
            zorder=3,
        )
        ax.set_xlabel("Communication Round")
        ax.set_ylabel(ylabel)
        ax.set_title(label, fontweight="bold", loc="left", fontsize=11)
        ax.grid(True, linestyle="--", linewidth=0.5, alpha=0.4, zorder=0)
        ax.xaxis.set_major_locator(mticker.MaxNLocator(nbins=6, integer=True))
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.margins(x=0.02)

    fig.suptitle(
        "PFLAD Client-Scaling Experiment "
        f"(N = {run['num_clients']} clients, "
        f"{run['num_selected']} selected/round, "
        f"{run['participation_ratio'] * 100:.1f}% participation)",
        fontsize=12,
        fontweight="bold",
        y=1.03,
    )

    fig.tight_layout()
    out_pdf = INDIVIDUAL_DIR / f"{run['stem']}.pdf"
    fig.savefig(out_pdf)
    plt.close(fig)
    return out_pdf


# ---------------------------------------------------------------------
# Run-level summary CSV
# ---------------------------------------------------------------------
def write_summary_csv(summary_rows):
    fieldnames = [
        "file",
        "num_clients",
        "num_selected_clients",
        "participation_ratio",
        "num_rounds",
        "last_round",
        "final_train_loss",
        "final_auroc",
        "final_auprc",
        f"mean_last_{TAIL_K}_train_loss",
        f"mean_last_{TAIL_K}_auroc",
        f"mean_last_{TAIL_K}_auprc",
        "best_auroc",
        "best_auprc",
    ]

    with SUMMARY_CSV.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(summary_rows)


# ---------------------------------------------------------------------
# Aggregate by N
# ---------------------------------------------------------------------
def aggregate_by_num_clients(summary_rows):
    groups = defaultdict(list)
    for row in summary_rows:
        groups[row["num_clients"]].append(row)

    aggregated = []
    for n in sorted(groups):
        group = groups[n]
        selected_values = [r["num_selected_clients"] for r in group]
        ratio_values = [r["participation_ratio"] for r in group]
        auroc_values = [choose_metric(r, "auroc") for r in group]
        auprc_values = [choose_metric(r, "auprc") for r in group]

        def safe_std(values):
            return stdev(values) if len(values) > 1 else 0.0

        aggregated.append({
            "num_clients": n,
            "num_selected": int(round(mean(selected_values))),
            "participation_ratio": mean(ratio_values),
            "n_runs": len(group),
            "auroc_mean": mean(auroc_values),
            "auroc_std": safe_std(auroc_values),
            "auprc_mean": mean(auprc_values),
            "auprc_std": safe_std(auprc_values),
        })

    return aggregated


# ---------------------------------------------------------------------
# effect.png-style summary figure
# ---------------------------------------------------------------------
def plot_effect(aggregated):
    if not aggregated:
        raise ValueError("No aggregated rows to plot.")

    baseline = next(
        (row for row in aggregated if row["num_clients"] == BASELINE_N),
        aggregated[0],
    )

    x = [row["num_clients"] for row in aggregated]
    xtick_labels = [
        f"{row['num_clients']}\n({row['participation_ratio'] * 100:.0f}% sel.)"
        for row in aggregated
    ]

    fig, axes = plt.subplots(1, 2, figsize=(9.5, 3.8))

    panels = [
        (
            axes[0],
            "auroc_mean",
            "auroc_std",
            "AUROC",
            COLOR_AUROC,
            "s",
            "(a)",
        ),
        (
            axes[1],
            "auprc_mean",
            "auprc_std",
            "AUPRC",
            COLOR_AUPRC,
            "^",
            "(b)",
        ),
    ]

    for ax, mean_key, std_key, ylabel, color, marker, panel_label in panels:
        values = [row[mean_key] for row in aggregated]
        errors = [row[std_key] for row in aggregated]
        baseline_value = baseline[mean_key]

        has_error = any(err > 0 for err in errors)
        if has_error:
            ax.errorbar(
                x,
                values,
                yerr=errors,
                color=color,
                marker=marker,
                markersize=6,
                markerfacecolor="white",
                markeredgewidth=1.4,
                capsize=3,
                alpha=0.95,
                zorder=3,
                label=ylabel,
            )
        else:
            ax.plot(
                x,
                values,
                color=color,
                marker=marker,
                markersize=6,
                markerfacecolor="white",
                markeredgewidth=1.4,
                alpha=0.95,
                zorder=3,
                label=ylabel,
            )

        ax.axhline(
            baseline_value,
            color=COLOR_BASELINE,
            linestyle="--",
            linewidth=1.0,
            alpha=0.65,
            zorder=1,
            label=f"N={baseline['num_clients']} baseline",
        )

        # Gap annotations relative to the baseline.
        value_span = max(values) - min(values)
        pad = max(value_span * 0.08, 0.002)

        for row, value in zip(aggregated, values):
            n = row["num_clients"]
            if n == baseline["num_clients"]:
                # Show the baseline value itself.
                ax.annotate(
                    f"{value:.4f}",
                    xy=(n, value),
                    xytext=(0, 12),
                    textcoords="offset points",
                    ha="center",
                    fontsize=8.5,
                    color=color,
                    fontweight="bold",
                    zorder=5,
                )
                continue

            gap = value - baseline_value
            gap_color = COLOR_GAP_POS if gap >= 0 else COLOR_GAP_NEG
            sign = "+" if gap >= 0 else ""

            offset_y = 13 if gap >= 0 else -20
            ax.annotate(
                f"Gap: {sign}{gap:.4f}",
                xy=(n, value),
                xytext=(0, offset_y),
                textcoords="offset points",
                ha="center",
                fontsize=8.2,
                color="black",
                bbox=dict(
                    boxstyle="round,pad=0.25",
                    facecolor=COLOR_GAP_BOX,
                    edgecolor="#888888",
                    linewidth=0.7,
                    alpha=0.9,
                ),
                arrowprops=dict(
                    arrowstyle="->",
                    color="#666666",
                    linewidth=0.7,
                    shrinkA=3,
                    shrinkB=4,
                ),
                zorder=5,
            )

        ax.set_xlabel("Number of Clients (N)")
        ax.set_ylabel(ylabel)
        ax.set_title(panel_label, fontweight="bold", loc="left", fontsize=11)
        ax.grid(True, linestyle="--", linewidth=0.5, alpha=0.4, zorder=0)
        ax.xaxis.set_major_locator(mticker.FixedLocator(x))
        ax.set_xticklabels(xtick_labels)
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.margins(x=0.12, y=0.25)
        ax.legend(loc="best", frameon=False)

        # Prevent an overly tight y-axis when all values are nearly identical.
        ymin, ymax = ax.get_ylim()
        if math.isclose(ymin, ymax):
            ax.set_ylim(ymin - pad, ymax + pad)

    metric_label = (
        "final round"
        if SUMMARY_MODE == "final"
        else f"mean of last {TAIL_K} rounds"
    )
    fig.suptitle(
        "Effect of Number of Clients on PFLAD Performance "
        f"({metric_label}; gap relative to N={baseline['num_clients']})",
        fontsize=12,
        fontweight="bold",
        y=1.02,
    )

    fig.tight_layout()
    fig.savefig(EFFECT_PDF)
    fig.savefig(EFFECT_PNG)
    plt.close(fig)


# ---------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------
def main():
    if not DATA_DIR.exists():
        raise FileNotFoundError(
            f"Data directory does not exist: {DATA_DIR}\n"
            "Place this script next to client_scaling_20261031_224126/ "
            "or change DATA_DIR."
        )

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    INDIVIDUAL_DIR.mkdir(parents=True, exist_ok=True)

    csv_paths = sorted(
        p for p in DATA_DIR.glob("*.csv")
        if p.name != SUMMARY_CSV.name
    )

    if not csv_paths:
        raise FileNotFoundError(f"No CSV files found under: {DATA_DIR}")

    if len(csv_paths) != EXPECTED_CSV_COUNT:
        print(
            f"[Warning] Expected {EXPECTED_CSV_COUNT} CSV files, "
            f"but found {len(csv_paths)}."
        )

    runs = []
    summary_rows = []

    print(f"Found {len(csv_paths)} CSV file(s).")
    for csv_path in csv_paths:
        run = read_run(csv_path)
        runs.append(run)

        out_pdf = plot_individual_run(run)
        summary = summarize_run(run)
        summary_rows.append(summary)

        print(
            f"[OK] {csv_path.name}: "
            f"N={run['num_clients']}, selected={run['num_selected']}, "
            f"final AUROC={summary['final_auroc']:.6f}, "
            f"final AUPRC={summary['final_auprc']:.6f}"
        )
        print(f"     saved: {out_pdf}")

    summary_rows.sort(key=lambda r: (r["num_clients"], r["file"]))
    write_summary_csv(summary_rows)

    aggregated = aggregate_by_num_clients(summary_rows)
    plot_effect(aggregated)

    print()
    print(f"Saved summary CSV: {SUMMARY_CSV}")
    print(f"Saved effect PDF : {EFFECT_PDF}")
    print(f"Saved effect PNG : {EFFECT_PNG}")


if __name__ == "__main__":
    main()
'''

path = Path("/mnt/data/plot_client_scaling_batch.py")
path.write_text(code, encoding="utf-8")
py_compile.compile(str(path), doraise=True)

print(f"Created and syntax-checked: {path}")