plot_heatmap.py

#!/home/aris/.cs204_venv/bin/python3
"""
plot_heatmap.py — CS204 Per-Set Cache Eviction Heatmap
IIT Ropar | 2025-26

Reads one or more heatmap CSVs produced by:
  ./cache_sim ... --heatmap <policy>_heatmap.csv

Usage:
  # Single policy heatmap
  python3 plot_heatmap.py --input LRU_heatmap.csv

  # Compare all policies side-by-side (pass multiple files)
  python3 plot_heatmap.py --input LRU_heatmap.csv FIFO_heatmap.csv \
                                  LFU_heatmap.csv RANDOM_heatmap.csv \
                                  OPT_heatmap.csv

  # Save to file instead of showing interactively
  python3 plot_heatmap.py --input LRU_heatmap.csv --out heatmap.png
"""

import argparse
import os
import sys
import csv

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.gridspec import GridSpec

# ── helpers ──────────────────────────────────────────────────────────────────

def load_csv(path: str) -> dict:
    """Return dict with arrays: set_id, evictions, accesses, hits, hit_rate."""
    data = {"set_id": [], "evictions": [], "accesses": [], "hits": [], "hit_rate": []}
    with open(path, newline="") as f:
        reader = csv.DictReader(f)
        for row in reader:
            data["set_id"].append(int(row["set_id"]))
            data["evictions"].append(int(row["evictions"]))
            data["accesses"].append(int(row["accesses"]))
            data["hits"].append(int(row["hits"]))
            data["hit_rate"].append(float(row["hit_rate"]))
    return {k: np.array(v) for k, v in data.items()}


def policy_from_filename(path: str) -> str:
    """Infer policy name from filename like 'LRU_heatmap.csv'."""
    base = os.path.basename(path)
    for p in ("LRU", "FIFO", "LFU", "RANDOM", "OPT"):
        if p in base.upper():
            return p
    return base.replace("_heatmap.csv", "").upper()


# ── single-policy plot ────────────────────────────────────────────────────────

def plot_single(data: dict, policy: str, out: str | None):
    """
    Three-panel figure:
      Top    — eviction count per set (bar chart, heat-coloured)
      Middle — access count per set
      Bottom — hit rate per set (line)
    """
    sets      = data["set_id"]
    evictions = data["evictions"]
    accesses  = data["accesses"]
    hit_rate  = data["hit_rate"]
    n         = len(sets)

    # Normalise evictions 0→1 for colour mapping
    ev_norm = evictions / (evictions.max() if evictions.max() > 0 else 1)
    cmap    = plt.cm.YlOrRd
    colors  = [cmap(v) for v in ev_norm]

    fig, axes = plt.subplots(3, 1, figsize=(max(12, n * 0.3), 10),
                             sharex=True, constrained_layout=True)
    fig.suptitle(f"Per-Set Cache Heatmap — Policy: {policy}", fontsize=14, fontweight="bold")

    # ── Evictions bar ────────────────────────────────────────────
    axes[0].bar(sets, evictions, color=colors, width=0.8, edgecolor="none")
    axes[0].set_ylabel("Evictions")
    axes[0].set_title("Eviction Count per Set  (hot = more evictions)")
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=mcolors.Normalize(0, evictions.max()))
    sm.set_array([])
    fig.colorbar(sm, ax=axes[0], orientation="vertical", fraction=0.02, pad=0.01,
                 label="Eviction intensity")

    # Annotate top-5 hottest sets
    top5_idx = np.argsort(evictions)[-5:][::-1]
    for idx in top5_idx:
        axes[0].annotate(f"S{sets[idx]}",
                         xy=(sets[idx], evictions[idx]),
                         xytext=(0, 4), textcoords="offset points",
                         ha="center", fontsize=7, color="darkred")

    # ── Accesses bar ─────────────────────────────────────────────
    axes[1].bar(sets, accesses, color="steelblue", width=0.8, edgecolor="none", alpha=0.8)
    axes[1].set_ylabel("Accesses")
    axes[1].set_title("Total Accesses per Set")

    # ── Hit rate line ─────────────────────────────────────────────
    axes[2].plot(sets, hit_rate, color="green", linewidth=1.2, alpha=0.9)
    axes[2].fill_between(sets, hit_rate, alpha=0.15, color="green")
    axes[2].set_ylabel("Hit Rate (%)")
    axes[2].set_xlabel("Set Index")
    axes[2].set_title("Hit Rate per Set")
    axes[2].set_ylim(0, 105)
    axes[2].axhline(hit_rate.mean(), color="gray", linestyle="--",
                    linewidth=0.8, label=f"Mean {hit_rate.mean():.1f}%")
    axes[2].legend(fontsize=8)

    if out:
        fig.savefig(out, dpi=150, bbox_inches="tight")
        print(f"Saved → {out}")
    else:
        plt.show()


# ── multi-policy comparison ───────────────────────────────────────────────────

def plot_multi(files: list[str], out: str | None):
    """
    One row per policy: eviction heatmap as a 1-D colour strip + bar.
    All strips share the same colour scale (global max).
    """
    datasets = [(policy_from_filename(f), load_csv(f)) for f in files]
    n_pol    = len(datasets)
    n_sets   = len(datasets[0][1]["set_id"])

    # Global colour scale
    global_max = max(d["evictions"].max() for _, d in datasets)
    cmap  = plt.cm.YlOrRd
    norm  = mcolors.Normalize(vmin=0, vmax=global_max)

    fig, axes = plt.subplots(n_pol, 1,
                             figsize=(max(14, n_sets * 0.25), 3 * n_pol),
                             constrained_layout=True)
    if n_pol == 1:
        axes = [axes]

    fig.suptitle("Per-Set Eviction Heatmap Comparison", fontsize=14, fontweight="bold")

    for ax, (policy, data) in zip(axes, datasets):
        evictions = data["evictions"]
        sets      = data["set_id"]

        colors = [cmap(norm(v)) for v in evictions]
        bars   = ax.bar(sets, evictions, color=colors, width=0.9, edgecolor="none")

        ax.set_ylabel("Evictions", fontsize=9)
        ax.set_title(f"{policy}  —  total evictions: {evictions.sum():,}  "
                     f"  hottest set: {sets[evictions.argmax()]} ({evictions.max():,})",
                     fontsize=10, loc="left")
        ax.set_xlim(-0.5, n_sets - 0.5)

    axes[-1].set_xlabel("Set Index")

    # Shared colorbar
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    fig.colorbar(sm, ax=axes, orientation="vertical", fraction=0.015, pad=0.01,
                 label="Evictions (shared scale)")

    if out:
        fig.savefig(out, dpi=150, bbox_inches="tight")
        print(f"Saved → {out}")
    else:
        plt.show()


# ── summary table ────────────────────────────────────────────────────────────

def print_summary(files: list[str]):
    print(f"\n{'Policy':<10} {'TotalEv':>10} {'HottestSet':>12} "
          f"{'HottestEv':>12} {'AvgHitRate':>12}")
    print("-" * 58)
    for f in files:
        policy = policy_from_filename(f)
        data   = load_csv(f)
        ev     = data["evictions"]
        hr     = data["hit_rate"]
        hot    = data["set_id"][ev.argmax()]
        print(f"{policy:<10} {ev.sum():>10,} {hot:>12} "
              f"{ev.max():>12,} {hr.mean():>11.2f}%")
    print()


# ── main ─────────────────────────────────────────────────────────────────────

def main():
    p = argparse.ArgumentParser(description="CS204 Cache Heatmap Plotter")
    p.add_argument("--input", nargs="+", required=True, metavar="CSV",
                   help="Heatmap CSV file(s) from cache_sim --heatmap")
    p.add_argument("--out",   default=None, metavar="FILE",
                   help="Save figure to file instead of showing (e.g. heatmap.png)")
    args = p.parse_args()

    for path in args.input:
        if not os.path.exists(path):
            sys.exit(f"Error: file not found: {path}")

    print_summary(args.input)

    if len(args.input) == 1:
        policy = policy_from_filename(args.input[0])
        data   = load_csv(args.input[0])
        plot_single(data, policy, args.out)
    else:
        plot_multi(args.input, args.out)


if __name__ == "__main__":
    main()