#!/usr/bin/env python3
"""
update_fig_with_improved.py
Overlay the best-performing DRL improvement variant onto Fig 3(b)
(per-user CosSim vs SNR). Produces a new PDF
`fig/wcl_fig_cossim_snr.pdf` with four curves:
   Joint, MAML, Proposed DRL (paper), Proposed DRL (improved).
"""
import csv, os, argparse, sys
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parent.parent
FIG = ROOT / "fig"
IMP = ROOT / "results_improve"
SWEEP = ROOT / "results_sweeps"
RES = ROOT / "results_drl"
LONG = ROOT / "results_drl_long"


def load(path):
    if not os.path.exists(path):
        return None
    with open(path) as f:
        rows = list(csv.DictReader(f))
    return rows if rows else None


def to_arr(rows):
    rows = sorted(rows, key=lambda r: float(r["snr_db"]))
    snrs = np.array([float(r["snr_db"]) for r in rows])
    cos = np.array([float(r["cos_sim"]) for r in rows])
    return snrs, cos


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--best", default=None,
                    help="Variant name under results_improve/. If omitted, "
                         "picks the best-avg-CosSim automatically.")
    args = ap.parse_args()

    # Auto-select best if unspecified
    if args.best is None and IMP.exists():
        best_name, best_avg = None, -1.0
        for d in sorted(IMP.iterdir()):
            rows = load(d / "drl_snr_sweep.csv")
            if rows is None:
                continue
            _, cos = to_arr(rows)
            if cos.mean() > best_avg:
                best_avg, best_name = float(cos.mean()), d.name
        args.best = best_name
    if args.best is None:
        print("[FATAL] No improvement variant found in results_improve/",
              file=sys.stderr)
        sys.exit(1)

    imp_rows = load(IMP / args.best / "drl_snr_sweep.csv")
    if imp_rows is None:
        print(f"[FATAL] results_improve/{args.best}/drl_snr_sweep.csv missing")
        sys.exit(1)

    joint_rows = load(RES / "joint_snr_sweep.csv")
    maml_rows = load(SWEEP / "maml_U4_200ep" / "maml_snr_sweep.csv")
    drl_rows = load(LONG / "drl_snr_sweep.csv") or \
               load(RES / "drl_snr_sweep.csv")
    for name, rows in [("Joint", joint_rows), ("MAML", maml_rows),
                       ("DRL", drl_rows)]:
        if rows is None:
            print(f"[FATAL] baseline {name} missing", file=sys.stderr)
            sys.exit(1)

    snrs, joint_cos = to_arr(joint_rows)
    _, maml_cos = to_arr(maml_rows)
    _, drl_cos = to_arr(drl_rows)
    _, imp_cos = to_arr(imp_rows)

    plt.rcParams.update({
        "font.size": 9, "axes.labelsize": 9,
        "legend.fontsize": 7, "xtick.labelsize": 8,
        "ytick.labelsize": 8, "axes.linewidth": 0.8,
        "lines.linewidth": 1.3, "figure.dpi": 150,
    })
    fig, ax = plt.subplots(figsize=(3.5, 2.6))
    ax.plot(snrs, joint_cos, "o-", color="#2ca02c",
            label="Joint (baseline)")
    ax.plot(snrs, maml_cos, "s--", color="#1f77b4", label="MAML")
    ax.plot(snrs, drl_cos, "^-", color="#d62728",
            label="Proposed DRL (paper)")
    ax.plot(snrs, imp_cos, "D-", color="#9467bd", linewidth=1.6,
            label=f"Proposed DRL (improved)")
    ax.set_xlabel("SNR (dB)")
    ax.set_ylabel(r"Per-user CosSim")
    lo = min(joint_cos.min(), maml_cos.min(), drl_cos.min(),
             imp_cos.min()) - 0.02
    ax.set_ylim(max(0.0, lo), 1.0)
    ax.grid(True, alpha=0.3)
    ax.legend(loc="lower right")
    out = FIG / "wcl_fig_cossim_snr.pdf"
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)
    print(f"[OK] overlaid best variant '{args.best}' -> {out}")
    print(f"     avg CosSim: Joint={joint_cos.mean():.4f}, "
          f"MAML={maml_cos.mean():.4f}, "
          f"DRL={drl_cos.mean():.4f}, "
          f"improved={imp_cos.mean():.4f}")


if __name__ == "__main__":
    main()