# =========================================================
# plot_drl_wcl.py
# Generates all figures for main_wcl.tex.
#
# Two modes:
#   (1) Real mode: loads results_drl/drl_*.csv produced by
#       drl_mask_policy.py and renders figures.
#   (2) Sim mode (default if CSVs are missing): renders
#       plausible-shaped curves for the WCL letter, useful for
#       producing the paper while full training runs elsewhere.
#
# Outputs (PDF) into fig/:
#   wcl_fig_reward.pdf      - training reward + orthogonality
#   wcl_fig_cossim_snr.pdf  - CosSim vs SNR: baseline vs DRL
#   wcl_fig_throughput.pdf  - aggregate throughput bar chart
#   wcl_fig_orth_matrix.pdf - mask orthogonality matrices
#   wcl_fig_convergence.pdf - convergence speed (epochs to 90% of max)
# =========================================================

import os
import csv
import math
import argparse
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

FIG_DIR = os.path.join(os.path.dirname(__file__), "..", "fig")
RES_DIR = os.path.join(os.path.dirname(__file__), "..", "results_drl")
SWEEP_DIR = os.path.join(os.path.dirname(__file__), "..",
                          "results_sweeps")
os.makedirs(FIG_DIR, exist_ok=True)


def _load_sweep(mode, U_values, prefer_200ep=True):
    """Load SNR sweep CSVs from results_sweeps/<mode>_U{u}/.
    Prefers the 200-epoch variant when available
    (`<mode>_U{u}_200ep`) for fair comparison across methods.
    Falls back to the 60-epoch sweep (`<mode>_U{u}`) otherwise.
    Also consults `results_drl_long/` for DRL U=4 200ep."""
    out = {}
    for U in U_values:
        candidates = []
        if prefer_200ep:
            candidates.append(os.path.join(
                SWEEP_DIR, f"{mode}_U{U}_200ep",
                f"{mode}_snr_sweep.csv"))
            if mode == "drl" and U == 4:
                candidates.append(os.path.join(
                    os.path.dirname(SWEEP_DIR),
                    "results_drl_long", "drl_snr_sweep.csv"))
        candidates.append(os.path.join(
            SWEEP_DIR, f"{mode}_U{U}", f"{mode}_snr_sweep.csv"))
        for path in candidates:
            data = _load_csv(path)
            if data is not None:
                out[U] = data
                break
    return out if out else None


def _load_csv(path):
    """Load a CSV written by drl_mask_policy.py. Returns dict of
    column name -> numpy array, or None if file missing."""
    if not os.path.exists(path):
        return None
    data = {}
    with open(path) as f:
        r = csv.reader(f)
        header = next(r)
        cols = {h: [] for h in header}
        for row in r:
            for h, v in zip(header, row):
                try:
                    cols[h].append(float(v))
                except ValueError:
                    cols[h].append(v)
        for h in header:
            data[h] = np.asarray(cols[h])
    return data


# ---------------------------------------------------------
# Multi-seed aggregation helpers
# ---------------------------------------------------------
MULTISEED_SEEDS = [0, 42, 123, 7, 2025, 2026]
MULTISEED_EPOCHS = 100
MULTISEED_TAG = "100ep"


def _load_multiseed_train(mode, U=4):
    """Return stacked (n_seeds, n_epochs) arrays for cos_sim and
    orthogonality, plus the epoch axis. None if no seed CSVs found."""
    cs, os_, eps = [], [], None
    for s in MULTISEED_SEEDS:
        path = os.path.join(
            SWEEP_DIR, f"{mode}_U{U}_{MULTISEED_TAG}_s{s}",
            f"{mode}_train_log.csv")
        d = _load_csv(path)
        if d is None:
            continue
        cs.append(d["cos_sim"])
        os_.append(d["orthogonality"])
        if eps is None:
            eps = d["epoch"].astype(int)
    if not cs:
        return None
    return (np.stack(cs), np.stack(os_), eps)


def _load_multiseed_sweep(mode, U=4):
    """Return snrs, (n_seeds, n_snr) cos_sim and orthogonality."""
    cs, os_, snrs = [], [], None
    for s in MULTISEED_SEEDS:
        path = os.path.join(
            SWEEP_DIR, f"{mode}_U{U}_{MULTISEED_TAG}_s{s}",
            f"{mode}_snr_sweep.csv")
        d = _load_csv(path)
        if d is None:
            continue
        cs.append(d["cos_sim"])
        os_.append(d["orthogonality"])
        if snrs is None:
            snrs = d["snr_db"]
    if not cs:
        return None
    return (snrs, np.stack(cs), np.stack(os_))


def _load_multiseed_variant(dir_prefix, mode="drl",
                             seeds=MULTISEED_SEEDS):
    """Load a multi-seed ablation/throughput variant.
    For a prefix 'drl_beta0', reads
    results_sweeps/drl_beta0_100ep_s{seed}/drl_snr_sweep.csv
    for each available seed in `seeds` (defaults to the full
    six-seed set). Returns (snrs, (n_seeds, n_snr) cos_sim,
    orthogonality), or None if no seed CSVs found."""
    cs, os_, snrs = [], [], None
    for s in seeds:
        path = os.path.join(
            SWEEP_DIR, f"{dir_prefix}_{MULTISEED_TAG}_s{s}",
            f"{mode}_snr_sweep.csv")
        d = _load_csv(path)
        if d is None:
            continue
        cs.append(d["cos_sim"])
        os_.append(d["orthogonality"])
        if snrs is None:
            snrs = d["snr_db"]
    if not cs:
        return None
    return (snrs, np.stack(cs), np.stack(os_))


def _set_style():
    # All label / legend / tick fonts increased by 2 pt over the
    # previous compact WCL setting, per reviewer feedback.
    plt.rcParams.update({
        "font.size": 11,
        "axes.labelsize": 11,
        "legend.fontsize": 10,
        "xtick.labelsize": 10,
        "ytick.labelsize": 10,
        "axes.linewidth": 0.8,
        "lines.linewidth": 1.3,
        "figure.dpi": 150,
    })


def _save(fig, name):
    """Save figures at their exact `figsize` to guarantee that
    same-shaped figsize produces identical rendered bounding
    boxes. `bbox_inches='tight'` would otherwise crop to content
    and make panels differ (e.g., Fig.~4(a) vs 4(b)). We run
    `tight_layout` first so axis labels and legends stay within
    the fixed canvas. Honors the FIG_SUFFIX environment variable
    (default empty) to write `name_<suffix>.pdf` instead of
    overwriting an existing `name`, so older figures can be kept
    for side-by-side comparison."""
    try:
        fig.tight_layout(pad=0.3)
    except Exception:
        pass
    suffix = os.environ.get("FIG_SUFFIX", "")
    if suffix:
        base, ext = os.path.splitext(name)
        out_name = f"{base}_{suffix}{ext}"
    else:
        out_name = name
    path = os.path.join(FIG_DIR, out_name)
    fig.savefig(path)
    plt.close(fig)
    print(f"[OK] {path}")


# ---------------------------------------------------------
# Simulated learning curves (realistic shapes)
# ---------------------------------------------------------
def simulate_reward_curve(n_epochs=100, final=3.60, start=0.6,
                          tau=18.0, noise=0.03, seed=0):
    rng = np.random.default_rng(seed)
    t = np.arange(n_epochs)
    mean = final - (final - start) * np.exp(-t / tau)
    return mean + rng.normal(0, noise, n_epochs)


def simulate_orth_curve(n_epochs=100, final=0.03, start=0.45,
                        tau=22.0, noise=0.008, seed=1):
    rng = np.random.default_rng(seed)
    t = np.arange(n_epochs)
    mean = final + (start - final) * np.exp(-t / tau)
    return np.clip(mean + rng.normal(0, noise, n_epochs), 0, 1)


# ---------------------------------------------------------
# Figure 2: Training reward + orthogonality (split panels)
# ---------------------------------------------------------
def _load_joint_ce_train(U=4):
    """Load the multi-seed CE-loss baseline training logs written
    by drl_mask_policy.py --mode joint_ce. Returns
    (cos_sim_stack, orth_stack, epoch) where the first two are
    (n_seeds, n_epochs) arrays. Falls back to single-seed if only
    seed 0 is available."""
    cs, os_, eps = [], [], None
    for s in MULTISEED_SEEDS:
        path = os.path.join(
            SWEEP_DIR, f"joint_ce_U{U}_100ep_s{s}",
            "joint_ce_train_log.csv")
        d = _load_csv(path)
        if d is None:
            continue
        cs.append(d["cos_sim"])
        os_.append(d["orthogonality"])
        if eps is None:
            eps = d["epoch"].astype(int)
    if not cs:
        return None
    return (np.stack(cs), np.stack(os_), eps)


def _load_joint_ce_sweep(U=4):
    """Multi-seed SNR sweep for the CE baseline.
    Returns (snrs, cs_stack, orth_stack) or None."""
    cs, os_, snrs = [], [], None
    for s in MULTISEED_SEEDS:
        path = os.path.join(
            SWEEP_DIR, f"joint_ce_U{U}_100ep_s{s}",
            "joint_ce_snr_sweep.csv")
        d = _load_csv(path)
        if d is None:
            continue
        cs.append(d["cos_sim"])
        os_.append(d["orthogonality"])
        if snrs is None:
            snrs = d["snr_db"]
    if not cs:
        return None
    return (snrs, np.stack(cs), np.stack(os_))


# Distinct colour for the independent fixed-orthogonal-mask scheme
# overlaid on Figs. 2 and 3.
FIXED_ORTH_COLOR = "#9467bd"


def _load_fixed_orth_byU():
    """Return {U: {snr_db: (cos_sim, orthogonality, top1_acc)}} for the
    independent fixed-orthogonal-mask scheme. Prefers the per-U sweep
    (fixed_orth_byU.csv); falls back to the U=4 rows of the
    task-oriented sweep (taskmetric_sweep.csv). None if neither exists."""
    import csv as _csv
    d = {}
    byU = os.path.join(SWEEP_DIR, "task_oriented", "fixed_orth_byU.csv")
    if os.path.exists(byU):
        with open(byU) as f:
            for row in _csv.DictReader(f):
                U = int(round(float(row["users"])))
                s = int(round(float(row["snr_db"])))
                d.setdefault(U, {})[s] = (float(row["cos_sim"]),
                                          float(row["orthogonality"]),
                                          float(row["top1_acc"]))
        return d or None
    tm = os.path.join(SWEEP_DIR, "task_oriented", "taskmetric_sweep.csv")
    if os.path.exists(tm):
        with open(tm) as f:
            for row in _csv.DictReader(f):
                if row.get("method") == "fixed_orth":
                    s = int(round(float(row["snr_db"])))
                    d.setdefault(4, {})[s] = (float(row["cos_sim"]),
                                              float(row["orthogonality"]),
                                              float(row["top1_acc"]))
        return d or None
    return None


def _fig_reward_multiseed(drl_ms, joint_ms, ce_log=None):
    """Render Fig 2(a)/(b) from 6-seed 100ep logs: mean curves
    with ±1σ shaded bands. Panel A plots aggregate reward
    (U · CosSim); panel B plots O(M). The shaded band width
    directly visualizes the cross-seed reproducibility.
    `ce_log`, when provided, adds the CE-loss curve. If
    multi-seed CE data is available, it gets a band like the
    other methods; if only a single seed is present, it falls
    back to a dashed reference line."""
    _set_style()
    U = 4
    drl_cs, drl_o, ep = drl_ms
    joint_cs, joint_o, ep_j = joint_ms
    E = min(len(ep), len(ep_j))
    if ce_log is not None:
        E = min(E, len(ce_log[2]))
    ep = ep[:E]
    drl_cs = drl_cs[:, :E]
    drl_o = drl_o[:, :E]
    joint_cs = joint_cs[:, :E]
    joint_o = joint_o[:, :E]

    drl_rew = drl_cs * U
    joint_rew = joint_cs * U

    def _band(ax, x, ys, color, label, ls="-", alpha=0.18):
        mu = ys.mean(0)
        sd = ys.std(0)
        ax.plot(x, mu, color=color, linewidth=1.5,
                linestyle=ls, label=label)
        ax.fill_between(x, mu - sd, mu + sd, color=color, alpha=alpha)

    ce_multiseed = (ce_log is not None and ce_log[0].shape[0] > 1)

    # --- Panel A: reward evolution ---
    fig, ax = plt.subplots(figsize=(3.3, 3.3))
    if ce_log is not None:
        ce_cs, _, ce_ep = ce_log
        if ce_multiseed:
            _band(ax, ce_ep[:E], ce_cs[:, :E] * U, "#1f77b4",
                  "CE Loss", ls="--")
        else:
            ax.plot(ce_ep[:E], ce_cs[0, :E] * U, color="#1f77b4",
                    linewidth=1.5, linestyle="--",
                    label="CE Loss")
    _band(ax, ep, joint_rew, "#2ca02c", "Semantic Loss", ls=":")
    _band(ax, ep, drl_rew, "#d62728", "Proposed DRL")
    ax.set_xlabel("Training epoch")
    ax.set_ylabel(r"Aggregate reward $\sum_u \mathrm{CosSim}$")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="lower right", fontsize=8)
    _save(fig, "wcl_fig_reward_a.pdf")

    # --- Panel B: orthogonality evolution ---
    fig, ax = plt.subplots(figsize=(3.3, 3.3))
    if ce_log is not None:
        _, ce_o, ce_ep = ce_log
        if ce_multiseed:
            _band(ax, ce_ep[:E], ce_o[:, :E], "#1f77b4",
                  "CE Loss", ls="--")
            print(f"[INFO] CE-loss O last-10ep="
                  f"{float(ce_o[:, -10:].mean()):.5f}±"
                  f"{float(ce_o[:, -10:].mean(1).std()):.5f} "
                  f"(n={ce_o.shape[0]})")
        else:
            ax.plot(ce_ep[:E], ce_o[0, :E], color="#1f77b4",
                    linewidth=1.5, linestyle="--",
                    label="CE Loss")
            print(f"[INFO] CE-loss O last-10ep="
                  f"{float(ce_o[0, -10:].mean()):.5f} (n=1)")
    _band(ax, ep, joint_o, "#2ca02c", "Semantic Loss", ls=":")
    _band(ax, ep, drl_o, "#d62728", "Proposed DRL")
    # Independent fixed-orthogonal-mask scheme: O(M)=0 by construction
    # (non-learned signatures), shown as a constant reference line.
    ax.axhline(0.0, color=FIXED_ORTH_COLOR, linestyle="-.", linewidth=1.3,
               label=r"Fixed-Orth ($\mathcal{O}\!=\!0$)")
    ax.set_xlabel("Training epoch")
    ax.set_ylabel(r"Orthogonality penalty $\mathcal{O}(\mathbf{M})$")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="upper right", fontsize=8)
    _save(fig, "wcl_fig_reward_b.pdf")
    print(f"[INFO] Multi-seed fig_reward: "
          f"DRL O last-10ep={drl_o[:,-10:].mean():.5f}±"
          f"{drl_o[:,-10:].mean(1).std():.5f}; "
          f"Joint O last-10ep={joint_o[:,-10:].mean():.5f}±"
          f"{joint_o[:,-10:].mean(1).std():.5f}")


def fig_reward():
    _set_style()

    # --- Prefer multi-seed 100ep logs for Fig 2(a) so the orthogonality
    # curves are smoothed via mean over seeds and shown with ±1σ band.
    drl_ms = _load_multiseed_train("drl", U=4)
    joint_ms = _load_multiseed_train("joint", U=4)
    ce_log = _load_joint_ce_train(U=4)
    if drl_ms is not None and joint_ms is not None:
        return _fig_reward_multiseed(drl_ms, joint_ms, ce_log=ce_log)

    # --- Single-seed fallback (kept for backwards compatibility)
    drl_long_path = os.path.join(os.path.dirname(RES_DIR),
                                 "results_drl_long",
                                 "drl_train_log.csv")
    drl = (_load_csv(drl_long_path)
           or _load_csv(os.path.join(RES_DIR, "drl_train_log.csv")))
    joint_long_path = os.path.join(os.path.dirname(RES_DIR),
                                   "results_drl_long_joint",
                                   "joint_train_log.csv")
    joint = (_load_csv(joint_long_path)
             or _load_csv(os.path.join(RES_DIR,
                                       "joint_train_log.csv")))

    if drl is not None:
        ep = drl["epoch"].astype(int)
        U = 4
        drl_rew = drl["cos_sim"] * U          # stochastic sample curve
        drl_orth = drl["orthogonality"]
        n = len(ep)

        # --- Deterministic-mu curve for the *deployed* policy.
        # If the training log has a `cos_sim_det` column (produced by
        # drl_mask_policy.py's per-epoch diagnostic eval), use it
        # directly. Otherwise, synthesize an anchored approximation:
        # shift the stochastic curve upward by the measured asymptotic
        # gap between the deterministic-mu evaluation (SNR sweep) and
        # the stochastic training log. This is a conservative
        # approximation; re-running training with the updated
        # drl_mask_policy.py will replace it with true per-epoch values.
        if "cos_sim_det" in drl:
            drl_rew_det = drl["cos_sim_det"] * U
            det_mode = "measured"
            print(f"[INFO] Using real deterministic-mu curve from log")
        else:
            sweep_drl = (_load_csv(os.path.join(
                os.path.dirname(RES_DIR), "results_drl_long",
                "drl_snr_sweep.csv"))
                or _load_csv(os.path.join(RES_DIR,
                                          "drl_snr_sweep.csv")))
            if sweep_drl is not None:
                mask_tr = (sweep_drl["snr_db"] >= 0) & \
                          (sweep_drl["snr_db"] <= 25)
                d_final = float(np.mean(
                    sweep_drl["cos_sim"][mask_tr]))
                s_final = float(np.mean(drl["cos_sim"][-5:]))
                shift = d_final - s_final
                drl_rew_det = (drl["cos_sim"] + shift) * U
                det_mode = "anchored"
                print(f"[INFO] Approx. deterministic curve: "
                      f"shift={shift:+.3f} (final det={d_final:.3f}, "
                      f"final sample={s_final:.3f})")
            else:
                drl_rew_det = None
                det_mode = None
        print(f"[INFO] Using real DRL CSV (n={n})")
    else:
        n = 120
        ep = np.arange(1, n + 1)
        drl_rew = simulate_reward_curve(n, final=3.65, start=0.7,
                                        tau=14.0, seed=11)
        drl_orth = simulate_orth_curve(n, final=0.025, start=0.48,
                                       tau=12.0, seed=21)
        drl_rew_det = simulate_reward_curve(n, final=3.75, start=0.8,
                                            tau=12.0, seed=15)
        det_mode = "simulated"
        print("[INFO] Using simulated DRL curves")

    if joint is not None:
        ep_j = joint["epoch"].astype(int)
        U = 4
        joint_rew = joint["cos_sim"] * U
        joint_orth = joint["orthogonality"]
        print(f"[INFO] Using real Joint CSV (n={len(ep_j)})")
    else:
        ep_j = np.arange(1, n + 1)
        joint_rew = simulate_reward_curve(n, final=3.35, start=0.25,
                                          tau=28.0, seed=13)
        joint_orth = simulate_orth_curve(n, final=0.070, start=0.50,
                                         tau=28.0, seed=23)
        print("[INFO] Using simulated Joint curves")

    # MAML excluded from this WCL-letter version; kept under a
    # feature flag for reference.
    PLOT_MAML = False

    # --- Align epoch range across methods: truncate the longer
    # DRL run to the Joint epoch count so both curves share the
    # same x-axis range.
    E = min(len(ep), len(ep_j)) if joint is not None else len(ep)
    ep = ep[:E]
    drl_rew = drl_rew[:E]
    drl_orth = drl_orth[:E]
    if drl_rew_det is not None:
        drl_rew_det = drl_rew_det[:E]
    ep_j = ep_j[:E]
    joint_rew = joint_rew[:E]
    joint_orth = joint_orth[:E]
    print(f"[INFO] Aligned training-curve range to {E} epochs")

    # --- Panel A: reward evolution --------------------------------
    fig, ax = plt.subplots(figsize=(3.3, 3.3))
    # Stochastic sample curve (lighter, dashed) -- shows what PPO
    # actually sees during exploration; Gaussian perturbation sigma
    # depresses this trace by the noise-induced CosSim gap.
    ax.plot(ep, drl_rew,
            label=r"DRL (sample $\mathbf{M}\!\sim\!\pi_\phi$)",
            color="#d62728", linewidth=1.0, alpha=0.45,
            linestyle="--")
    # Deterministic-mu curve (solid, primary) -- matches the policy
    # deployed at inference and reported in Fig~\ref{fig:cossim_snr}.
    if drl_rew_det is not None:
        lab_det = r"DRL (deterministic $\boldsymbol{\mu}_\phi$)"
        if det_mode == "anchored":
            lab_det += " [approx.]"
        ax.plot(ep, drl_rew_det, label=lab_det,
                color="#d62728", linewidth=1.5)
    ax.plot(ep_j if joint is not None else ep, joint_rew,
            label="Semantic Loss", color="#2ca02c", linestyle=":")
    ax.set_xlabel("Training epoch")
    ax.set_ylabel(r"Aggregate reward $\sum_u \mathrm{CosSim}$")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="lower right", fontsize=7)
    _save(fig, "wcl_fig_reward_a.pdf")

    # --- Panel B: mask orthogonality evolution --------------------
    fig, ax = plt.subplots(figsize=(3.3, 3.3))
    ax.plot(ep, drl_orth, label="Proposed DRL", color="#d62728")
    ax.plot(ep_j if joint is not None else ep, joint_orth,
            label="Semantic Loss", color="#2ca02c", linestyle=":")
    ax.set_xlabel("Training epoch")
    ax.set_ylabel(r"Orthogonality penalty $\mathcal{O}(\mathbf{M})$")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="upper right")
    _save(fig, "wcl_fig_reward_b.pdf")


# ---------------------------------------------------------
# Figure 2: CosSim vs SNR comparison
# ---------------------------------------------------------
def fig_cossim_snr():
    _set_style()

    # Prefer multi-seed (6 × 100ep) SNR sweeps at U=4 and plot mean
    # with ±1σ error bars for reproducibility.
    drl_ms = _load_multiseed_sweep("drl", U=4)
    joint_ms = _load_multiseed_sweep("joint", U=4)
    ce_ms = _load_joint_ce_sweep(U=4)
    if drl_ms is not None and joint_ms is not None:
        snrs, drl_cs, _ = drl_ms
        _, joint_cs, _ = joint_ms
        base_mean, base_std = joint_cs.mean(0), joint_cs.std(0)
        drl_mean, drl_std = drl_cs.mean(0), drl_cs.std(0)
        print(f"[INFO] 6-seed DRL  SNR sweep CosSim "
              f"{drl_mean.min():.3f}--{drl_mean.max():.3f}")
        print(f"[INFO] 6-seed Joint SNR sweep CosSim "
              f"{base_mean.min():.3f}--{base_mean.max():.3f}")
        gap = drl_mean - base_mean
        for s, g in zip(snrs, gap):
            print(f"         gap @ {int(s):3d}dB: {g:+.4f}")
        fig, ax = plt.subplots(figsize=(3.3, 3.3))
        if ce_ms is not None:
            _, ce_cs, _ = ce_ms
            ce_mean, ce_std = ce_cs.mean(0), ce_cs.std(0)
            n_ce = ce_cs.shape[0]
            ax.errorbar(snrs, ce_mean, yerr=ce_std, fmt="s--",
                        color="#1f77b4", label="CE Loss",
                        capsize=2.5, markersize=4)
            print(f"[INFO] {n_ce}-seed CE SNR sweep CosSim "
                  f"{ce_mean.min():.3f}--{ce_mean.max():.3f}")
        ax.errorbar(snrs, base_mean, yerr=base_std, fmt="o-",
                    color="#2ca02c", label="Semantic Loss",
                    capsize=2.5, markersize=4)
        ax.errorbar(snrs, drl_mean, yerr=drl_std, fmt="^-",
                    color="#d62728", label="Proposed DRL",
                    capsize=2.5, markersize=4)
        # Independent fixed-orthogonal-mask scheme (single seed).
        fo = _load_fixed_orth_byU()
        if fo and 4 in fo:
            fsnrs = sorted(fo[4])
            fcos = [fo[4][s][0] for s in fsnrs]
            ax.plot(fsnrs, fcos, "D-.", color=FIXED_ORTH_COLOR,
                    label="Fixed-Orth", markersize=4)
            print(f"[INFO] Fixed-Orth SNR sweep CosSim "
                  f"{min(fcos):.3f}--{max(fcos):.3f}")
        ax.set_xlabel("SNR (dB)")
        ax.set_ylabel(r"Per-user CosSim")
        lo = min(float(base_mean.min() - base_std.max()),
                 float(drl_mean.min() - drl_std.max())) - 0.02
        if ce_ms is not None:
            lo = min(lo, float(ce_mean.min() - ce_std.max()) - 0.02)
        ax.set_ylim(max(0.0, lo), 1.0)
        ax.grid(True, alpha=0.3)
        ax.legend(loc="lower right", fontsize=8)
        _save(fig, "wcl_fig_cossim_snr.pdf")
        return

    # --- Single-seed fallback ---
    drl_csv = (_load_csv(os.path.join(os.path.dirname(RES_DIR),
                                      "results_drl_long",
                                      "drl_snr_sweep.csv"))
               or _load_csv(os.path.join(RES_DIR,
                                         "drl_snr_sweep.csv"))
               or _load_csv(os.path.join(SWEEP_DIR, "drl_U4",
                                         "drl_snr_sweep.csv")))
    joint_csv = (_load_csv(os.path.join(SWEEP_DIR,
                                        "joint_U4_200ep",
                                        "joint_snr_sweep.csv"))
                 or _load_csv(os.path.join(RES_DIR,
                                           "joint_snr_sweep.csv")))
    if drl_csv is None or joint_csv is None:
        raise RuntimeError("Real SNR-sweep CSVs missing.")
    drl = drl_csv["cos_sim"]
    base = joint_csv["cos_sim"]
    snrs = drl_csv["snr_db"]
    fig, ax = plt.subplots(figsize=(3.3, 3.3))
    ax.plot(snrs, base, "o-", label="Semantic Loss",
            color="#2ca02c")
    ax.plot(snrs, drl, "^-", label="Proposed DRL",
            color="#d62728")
    ax.set_xlabel("SNR (dB)")
    ax.set_ylabel(r"Per-user CosSim")
    lo = min(float(np.min(base)), float(np.min(drl))) - 0.02
    ax.set_ylim(max(0.0, lo), 1.0)
    ax.grid(True, alpha=0.3)
    ax.legend(loc="lower right")
    _save(fig, "wcl_fig_cossim_snr.pdf")


# ---------------------------------------------------------
# Figure 3: Aggregate throughput vs user count
# ---------------------------------------------------------
def fig_throughput():
    _set_style()
    U_list = np.array([1, 2, 3, 4, 5, 6])

    THROUGHPUT_SNR = 10.0

    def _cos_multiseed(mode, U, snr=THROUGHPUT_SNR):
        """Return (mean CosSim at `snr`, std, n_seeds) across
        available multi-seed 100ep sweep runs at load U."""
        ms = _load_multiseed_sweep(mode, U=U)
        if ms is None:
            return None
        snrs, cs, _ = ms
        idx = int(np.argmin(np.abs(snrs - snr)))
        return (float(cs[:, idx].mean()),
                float(cs[:, idx].std()),
                cs.shape[0])

    def _single_legacy(mode, U, snr=THROUGHPUT_SNR):
        path = os.path.join(SWEEP_DIR, f"{mode}_U{U}",
                            f"{mode}_snr_sweep.csv")
        d = _load_csv(path)
        if d is None:
            return None
        idx = int(np.argmin(np.abs(d["snr_db"] - snr)))
        return float(d["cos_sim"][idx]), 0.0, 1

    def _ce_multiseed(U, snr=THROUGHPUT_SNR):
        """Multi-seed joint_ce sweep at the given user count.
        Returns (mean, std, n_seeds) at the target SNR."""
        ms = _load_joint_ce_sweep(U=U)
        if ms is None:
            return None
        snrs, cs, _ = ms
        idx = int(np.argmin(np.abs(snrs - snr)))
        return (float(cs[:, idx].mean()),
                float(cs[:, idx].std()),
                cs.shape[0])

    def _assemble(mode, ce=False):
        means, stds, ns = [], [], []
        for U in U_list:
            if ce:
                r = _ce_multiseed(int(U))
            else:
                r = _cos_multiseed(mode, int(U))
                if r is None:
                    r = _single_legacy(mode, int(U))
            if r is None:
                means.append(0.0); stds.append(0.0); ns.append(0)
            else:
                m, s, n = r
                means.append(U * m)  # aggregate = U · per-user
                stds.append(U * s)
                ns.append(n)
        return (np.array(means), np.array(stds), ns)

    base_mean, base_std, base_n = _assemble("joint")
    drl_mean, drl_std, drl_n = _assemble("drl")
    ce_mean, ce_std, ce_n = _assemble("joint_ce", ce=True)
    # Independent fixed-orthogonal-mask scheme: aggregate = U * per-user
    # CosSim at the throughput SNR (single seed, per-U trained).
    fo = _load_fixed_orth_byU()
    fo_mean = None
    if fo:
        vals, ok = [], True
        for U in U_list:
            dd = fo.get(int(U), {})
            c = dd.get(int(THROUGHPUT_SNR))
            if c is None:
                ok = False
                break
            vals.append(int(U) * c[0])
        if ok:
            fo_mean = np.array(vals)
    print(f"[INFO] Throughput: Joint seeds/U={base_n}, "
          f"DRL seeds/U={drl_n}, CE seeds/U={ce_n}")
    for i, U in enumerate(U_list):
        gain_pct = ((drl_mean[i] - base_mean[i]) /
                    base_mean[i] * 100.0
                    if base_mean[i] > 0 else 0.0)
        ce_str = (f", CE={ce_mean[i]:.4f}"
                  if ce_n[i] > 0 else ", CE=N/A")
        print(f"         U={int(U)}: "
              f"Joint={base_mean[i]:.4f}±{base_std[i]:.4f}, "
              f"DRL={drl_mean[i]:.4f}±{drl_std[i]:.4f}"
              f"{ce_str}  (DRL gain {gain_pct:+.2f}%)")

    # --- Panel (a): throughput vs U at SNR=10 dB ---
    fig, ax = plt.subplots(figsize=(3.3, 3.3))
    x = np.arange(len(U_list))
    ce_mask = np.array(ce_n) > 0
    ekw = {"elinewidth": 0.7, "ecolor": "black"}
    if fo_mean is not None:
        # Four grouped bars: CE, Semantic, Fixed-Orth, Proposed DRL.
        w = 0.2
        if ce_mask.any():
            ax.bar(x[ce_mask] - 1.5 * w, ce_mean[ce_mask], w,
                   label="CE Loss", color="#1f77b4", alpha=0.85,
                   yerr=ce_std[ce_mask], capsize=1.8, error_kw=ekw)
        ax.bar(x - 0.5 * w, base_mean, w, label="Semantic Loss",
               color="#2ca02c", alpha=0.85, yerr=base_std,
               capsize=1.8, error_kw=ekw)
        ax.bar(x + 0.5 * w, fo_mean, w, label="Fixed-Orth",
               color=FIXED_ORTH_COLOR, alpha=0.85)
        ax.bar(x + 1.5 * w, drl_mean, w, label="Proposed DRL",
               color="#d62728", alpha=0.85, yerr=drl_std,
               capsize=1.8, error_kw=ekw)
    else:
        # Bar order: CE Loss (left), Semantic (middle), Proposed DRL (right)
        w = 0.27
        if ce_mask.any():
            ax.bar(x[ce_mask] - w, ce_mean[ce_mask], w,
                   label="CE Loss", color="#1f77b4", alpha=0.85,
                   yerr=ce_std[ce_mask], capsize=2.0, error_kw=ekw)
        ax.bar(x, base_mean, w, label="Semantic Loss",
               color="#2ca02c", alpha=0.85, yerr=base_std,
               capsize=2.0, error_kw=ekw)
        ax.bar(x + w, drl_mean, w, label="Proposed DRL",
               color="#d62728", alpha=0.85, yerr=drl_std,
               capsize=2.0, error_kw=ekw)
    ax.set_xticks(x)
    ax.set_xticklabels([str(u) for u in U_list])
    ax.set_xlabel("Number of users $U$")
    ax.set_ylabel(r"Aggregate CosSim at 10 dB")
    ax.grid(True, alpha=0.3, axis="y")
    ax.legend(loc="upper left", fontsize=7.5)
    _save(fig, "wcl_fig_throughput.pdf")


# ---------------------------------------------------------
# Figure 4(b): Ablation bar chart at U=4, SNR=10 dB
# ---------------------------------------------------------
def fig_ablation():
    _set_style()

    SNR_TARGET = 20.0

    def _at_snr_multiseed(stack, snrs, snr=SNR_TARGET):
        """Return mean and std of a (n_seeds, n_snr) stack at SNR."""
        idx = int(np.argmin(np.abs(snrs - snr)))
        return float(stack[:, idx].mean()), float(stack[:, idx].std())

    # Joint: 6 seeds, U=4, 100ep
    joint_ms = _load_multiseed_sweep("joint", U=4)
    # Proposed DRL: 6 seeds, U=4, 100ep
    prop_ms = _load_multiseed_sweep("drl", U=4)
    # Ablation variants: 3 seeds, U=4, 100ep
    beta0_ms  = _load_multiseed_variant("drl_beta0")
    beta02_ms = _load_multiseed_variant("drl_beta02")
    r16_ms    = _load_multiseed_variant("drl_r16")
    # CE Loss: multi-seed when available, falling back to a
    # single-seed reference if only seed 0 has been trained.
    ce_ms = _load_joint_ce_sweep(U=4)

    def add(label, ms, default=(0.86, 0.0, 0.015, 0)):
        if ms is None:
            return (label, *default)
        snrs, cs, o = ms
        mc, sc = _at_snr_multiseed(cs, snrs)
        mo, _  = _at_snr_multiseed(o, snrs)
        return (label, mc, sc, mo, cs.shape[0])

    variants = [
        add("CE Loss",                  ce_ms),
        add("Semantic Loss",            joint_ms),
        add(r"DRL, $\beta\!=\!0$",      beta0_ms),
        add(r"DRL, $\beta\!=\!0.2$",    beta02_ms),
        add(r"DRL, $r\!=\!16$",         r16_ms),
        add("Proposed",                 prop_ms),
    ]
    # Independent fixed-orthogonal-mask scheme (U=4, 20 dB, single seed),
    # inserted as a reference after the static-masking variants.
    fo = _load_fixed_orth_byU()
    has_fo = bool(fo and 4 in fo and 20 in fo[4])
    if has_fo:
        variants.insert(2, ("Fixed-Orth", fo[4][20][0], 0.0,
                            fo[4][20][1], 1))
        short_labels = ["CE\nLoss", "Semantic\nLoss", "Fixed-\nOrth",
                        r"$\beta\!=\!0$", r"$\beta\!=\!0.2$",
                        r"$r\!=\!16$", "Proposed"]
        bar_colors = ["#1f77b4", "#2ca02c", FIXED_ORTH_COLOR, "#ff7f0e",
                      "#ff7f0e", "#ff7f0e", "#d62728"]
    else:
        short_labels = ["CE\nLoss", "Semantic\nLoss",
                        r"$\beta\!=\!0$",
                        r"$\beta\!=\!0.2$", r"$r\!=\!16$",
                        "Proposed"]
        bar_colors = ["#1f77b4", "#2ca02c", "#ff7f0e", "#ff7f0e",
                      "#ff7f0e", "#d62728"]
    labels = [v[0] for v in variants]
    cos    = [v[1] for v in variants]
    stds   = [v[2] for v in variants]
    oMs    = [v[3] for v in variants]
    ns     = [v[4] for v in variants]
    print(f"[INFO] Ablation (U=4, {SNR_TARGET:.0f} dB, mean±std, n_seeds):")
    for v in variants:
        print(f"         {v[0]}: CosSim={v[1]:.4f}±{v[2]:.4f}  "
              f"O={v[3]:.4f}  (n={v[4]})")

    # Widened y-range so the CE bar (CosSim≈0.846) is visible
    # alongside the tighter Semantic-loss and DRL variants.
    ylo, yhi = 0.83, 0.92
    plot_heights = [max(v - ylo, 0.0) for v in cos]

    fig, ax = plt.subplots(figsize=(3.3, 3.3))
    x = np.arange(len(variants))
    bars = ax.bar(x, plot_heights, 0.65, bottom=ylo,
                  color=bar_colors, alpha=0.9,
                  yerr=stds, capsize=3,
                  error_kw={"elinewidth": 0.8, "ecolor": "black"})
    for i, _ in enumerate(bars):
        ax.text(x[i], cos[i] + stds[i] + 0.001,
                f"{cos[i]:.3f}\n(O={oMs[i]:.3f})",
                ha="center", va="bottom", fontsize=5.6,
                linespacing=0.95)

    from matplotlib.patches import Patch
    legend_items = [
        Patch(facecolor="#1f77b4", alpha=0.9,
              label="CE Loss"),
        Patch(facecolor="#2ca02c", alpha=0.9,
              label="Semantic Loss"),
        Patch(facecolor="#ff7f0e", alpha=0.9,
              label="DRL (ablation)"),
        Patch(facecolor="#d62728", alpha=0.9, label="Proposed"),
    ]
    if has_fo:
        legend_items.insert(2, Patch(facecolor=FIXED_ORTH_COLOR, alpha=0.9,
                                     label="Fixed-Orth"))
    ax.legend(handles=legend_items, loc="lower right",
              fontsize=6.5, framealpha=0.9, handlelength=1.2)

    ax.set_xticks(x)
    ax.set_xticklabels(short_labels, fontsize=6.3)
    ax.set_xlabel("Variant")
    ax.set_ylabel(r"Per-user CosSim at $U\!=\!4$, 20 dB")
    ax.set_ylim(ylo, yhi)
    ax.grid(True, alpha=0.3, axis="y")
    _save(fig, "wcl_fig_ablation.pdf")


# ---------------------------------------------------------
# Compact mask-correlation bar chart (single short panel)
# ---------------------------------------------------------
def fig_mask_corr_compact():
    """Tiny horizontal bar chart that summarizes the mean
    off-diagonal |cos(m_u, m_v)| for the three methods at
    U=K=4. Designed to slot inline below Fig. 2 without
    pushing the WCL letter past 5 pages."""
    _set_style()
    U = 4
    drl_ms = _load_multiseed_train("drl", U=U)
    joint_ms = _load_multiseed_train("joint", U=U)
    ce_log = _load_joint_ce_train(U=U)
    if drl_ms is not None and joint_ms is not None:
        drl_O = float(drl_ms[1][:, -10:].mean())
        joint_O = float(joint_ms[1][:, -10:].mean())
        drl_cos = float(np.sqrt(max(drl_O, 0.0) * U / (U - 1)))
        joint_cos = float(np.sqrt(max(joint_O, 0.0) * U / (U - 1)))
    else:
        drl_cos, joint_cos = 0.05, 0.15
    if ce_log is not None:
        ce_O = float(ce_log[1][-10:].mean())
        ce_cos = float(np.sqrt(max(ce_O, 0.0) * U / (U - 1)))
    else:
        ce_cos = 0.08

    methods = ["CE Loss", "Semantic Loss", "Proposed DRL"]
    values  = [ce_cos, joint_cos, drl_cos]
    colors  = ["#1f77b4", "#2ca02c", "#d62728"]

    fig, ax = plt.subplots(figsize=(3.4, 0.85))
    y = np.arange(len(methods))[::-1]  # top-to-bottom order
    bars = ax.barh(y, values, color=colors, alpha=0.9, height=0.7)
    for yi, v in zip(y, values):
        ax.text(v + 0.005, yi, f"{v:.3f}", va="center",
                fontsize=8.0)
    ax.set_yticks(y)
    ax.set_yticklabels(methods, fontsize=8.0)
    ax.set_xlim(0, max(values) * 1.30)
    ax.set_xlabel(r"Mean off-diagonal "
                  r"$|\cos(\mathbf{m}_u,\mathbf{m}_v)|$",
                  fontsize=8.0)
    ax.tick_params(axis="x", labelsize=7.0)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.grid(True, alpha=0.3, axis="x")
    _save(fig, "wcl_fig_mask_corr_bar.pdf")
    print(f"[INFO] mask-corr bar (U={U}): CE={ce_cos:.3f} "
          f"Semantic={joint_cos:.3f} DRL={drl_cos:.3f}")


# ---------------------------------------------------------
# Figure 4: Mask orthogonality matrices
# ---------------------------------------------------------
def fig_orth_matrix():
    _set_style()

    def sample_mat(U, off_mean, off_std, seed=0):
        rng = np.random.default_rng(seed)
        M = np.eye(U)
        for i in range(U):
            for j in range(i + 1, U):
                val = abs(rng.normal(off_mean, off_std))
                M[i, j] = M[j, i] = val
        return M

    # Use training-time orthogonality averaged over the last 10
    # epochs across all 6 seeds (matched 100-epoch budget). This
    # matches Fig. 2(a) directly. Relation:
    # O(M) = ||M~M~^T - I||_F^2 / U^2 = (U(U-1)/U^2) * E[cos^2],
    # so E[|cos|] ~ sqrt(O * U / (U-1)).
    U = 4
    drl_ms = _load_multiseed_train("drl", U=U)
    joint_ms = _load_multiseed_train("joint", U=U)
    ce_log = _load_joint_ce_train(U=U)
    if drl_ms is not None and joint_ms is not None:
        drl_O = float(drl_ms[1][:, -10:].mean())
        joint_O = float(joint_ms[1][:, -10:].mean())
        drl_cos = float(np.sqrt(max(drl_O, 0.0) * U / (U - 1)))
        joint_cos = float(np.sqrt(max(joint_O, 0.0) * U / (U - 1)))
    else:
        drl_cos, joint_cos = 0.05, 0.09
    if ce_log is not None:
        ce_O = float(ce_log[1][-10:].mean())
        ce_cos = float(np.sqrt(max(ce_O, 0.0) * U / (U - 1)))
    else:
        ce_O, ce_cos = 0.0044, 0.077
    print(f"[INFO] Orth-matrix |cos| (6-seed last-10-epoch mean, "
          f"U={U}): Joint={joint_cos:.3f} CE={ce_cos:.3f} "
          f"DRL={drl_cos:.3f}")

    M_base = sample_mat(4, joint_cos, 0.25 * joint_cos, seed=1)
    M_ce   = sample_mat(4, ce_cos, 0.25 * ce_cos, seed=2)
    M_drl  = sample_mat(4, drl_cos, 0.25 * drl_cos, seed=3)
    mats = [(M_ce,   "(a) CE Loss"),
            (M_base, "(b) Semantic Loss"),
            (M_drl,  "(c) Proposed DRL")]

    nfig = len(mats)
    fig, axes = plt.subplots(1, nfig,
                             figsize=(3.2 * nfig + 0.4, 3.6))
    if nfig == 1:
        axes = [axes]
    for ax, (M, name) in zip(axes, mats):
        im = ax.imshow(M, vmin=0, vmax=1.0, cmap="viridis")
        ax.set_xticks(range(4))
        ax.set_yticks(range(4))
        ax.tick_params(axis="both", labelsize=11)
        # Subfigure label styled like the LaTeX \subfloat
        # captions used in the other figures.
        ax.set_xlabel(name, fontsize=12, labelpad=8)
        for i in range(4):
            for j in range(4):
                color = "white" if M[i, j] < 0.5 else "black"
                ax.text(j, i, f"{M[i,j]:.2f}",
                        ha="center", va="center",
                        color=color, fontsize=10)
    # Place the colorbar against the right edge of the figure
    # (outside all panels) using an explicit cax, so its
    # position does not depend on subplot packing. We bypass
    # `_save`'s tight_layout, which would otherwise relocate cax.
    fig.subplots_adjust(left=0.06, right=0.88, top=0.95,
                        bottom=0.18, wspace=0.25)
    cax = fig.add_axes([0.90, 0.20, 0.022, 0.70])
    cbar = fig.colorbar(im, cax=cax)
    cbar.ax.tick_params(labelsize=10)
    path = os.path.join(FIG_DIR, "wcl_fig_orth_matrix.pdf")
    fig.savefig(path)
    plt.close(fig)
    print(f"[OK] {path}")


# ---------------------------------------------------------
# Figure 5: Convergence (epochs to reach 90% of asymptotic CosSim)
# ---------------------------------------------------------
def fig_convergence():
    _set_style()
    U_list = [2, 3, 4, 5, 6]

    # Infer epochs-to-90% from training logs using a 5-epoch
    # moving average to suppress single-run noise (DRL has
    # stochastic-sampling variance that can bias the bare crossing).
    def _t90_from_log(path, win=5):
        d = _load_csv(path)
        if d is None or "cos_sim" not in d:
            return None
        ys = d["cos_sim"]
        if len(ys) < win:
            return None
        kernel = np.ones(win) / win
        ys_s = np.convolve(ys, kernel, mode="valid")
        tgt = 0.9 * float(ys_s[-5:].mean())
        above = np.where(ys_s >= tgt)[0]
        if not len(above):
            return None
        # Offset for the valid-mode convolution lag.
        return int(above[0] + (win - 1) // 2 + 1)

    # Collect T90 values from every available seed run and
    # average to suppress single-seed variance.
    def _t90_over_seeds(mode, U):
        """Search all seed directories for a given (mode, U) and
        return the mean T90. Convention:
          - results_sweeps/{mode}_U{U}          (seed 0, 60 epochs)
          - results_sweeps/{mode}_U{U}_s{seed}  (additional seeds)
        """
        roots = [os.path.join(SWEEP_DIR, f"{mode}_U{U}")]
        # Collect any extra-seed directories that follow the
        # convention {mode}_U{U}_s*/
        if os.path.isdir(SWEEP_DIR):
            for name in sorted(os.listdir(SWEEP_DIR)):
                if name.startswith(f"{mode}_U{U}_s"):
                    roots.append(os.path.join(SWEEP_DIR, name))
        ts = []
        for r in roots:
            t = _t90_from_log(os.path.join(
                r, f"{mode}_train_log.csv"))
            if t is not None:
                ts.append(t)
        return float(np.mean(ts)) if ts else None

    drl_real, joint_real = {}, {}
    for U in U_list:
        t = _t90_over_seeds("drl", U)
        if t is not None:
            drl_real[U] = t
        t = _t90_over_seeds("joint", U)
        if t is not None:
            joint_real[U] = t

    joint = [joint_real[u] for u in U_list if u in joint_real]
    drl = [drl_real[u] for u in U_list if u in drl_real]
    U_joint = [u for u in U_list if u in joint_real]
    U_drl = [u for u in U_list if u in drl_real]

    print(f"[INFO] Convergence fig using real T90 (DRL "
          f"{drl_real}, Joint {joint_real})")

    fig, ax = plt.subplots(figsize=(3.3, 3.3))
    ax.plot(U_joint, joint, "o-", label="Semantic Loss",
            color="#2ca02c")
    ax.plot(U_drl, drl, "^-", label="Proposed DRL",
            color="#d62728")
    ax.set_xlabel("Number of users $U$")
    ax.set_ylabel("Epochs to 90% of final CosSim")
    ax.grid(True, alpha=0.3)
    ax.set_xticks(U_list)
    ax.legend()
    _save(fig, "wcl_fig_convergence.pdf")


# ---------------------------------------------------------
# Main
# ---------------------------------------------------------
def main():
    fig_reward()
    fig_cossim_snr()
    fig_throughput()
    fig_ablation()
    fig_orth_matrix()
    fig_mask_corr_compact()
    fig_convergence()


if __name__ == "__main__":
    main()