from matplotlib import ticker
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
#import seaborn as sns   # tried it, the default look was worse

df = pd.read_csv("analysis/results.csv")

# peaks for the ryzen ai 7 350 (zen5) at 2GHz. 2 fma units and the datapath is 256 bit
# -> scalar 6, double 24, float 48 flops/cycle
SCALAR = 6
PEAK_D = 24
PEAK_F = 48
maroon = "#A71D5B"


def get(v, n=5000):
  return df[(df.variant == v) & (df.n == n)]["flops_per_cycle"].values[0]


def getline(v):
  s = df[df.variant == v].sort_values("n")
  return s.n.values, s["flops_per_cycle"].values


def label_lines_above(ax,
                      entries,
                      fontsize=9,
                      rotate=False,
                      dy_frac=0.012,
                      halo=0):
  """Place each line's description on the line.
    entries: (x, y, label, color) with two optional extra elements:
      [4] xpos  -> pin the label to that data-x (else labels are spread out)
      [5] side  -> "above" (default) or "below" to put the text under the line
    Set halo>0 to draw a white outline behind the text. Call AFTER
    set_xlim/set_ylim and tight_layout so any rotation angle is correct.
  """
  logx = ax.get_xscale() == "log"
  logy = ax.get_yscale() == "log"
  fx = (lambda v: np.log10(v)) if logx else (lambda v: v)
  fxi = (lambda v: 10**v) if logx else (lambda v: v)
  fy = (lambda v: np.log10(v)) if logy else (lambda v: v)
  fyi = (lambda v: 10**v) if logy else (lambda v: v)

  ent = []
  for e in entries:
    x, y, lab, c = e[0], e[1], e[2], e[3]
    if len(x) == 0:
      continue
    ent.append((np.asarray(x, float), np.asarray(y, float), lab, c,
                e[4] if len(e) > 4 else None, e[5] if len(e) > 5 else "above"))

  n = len(ent)
  y0, y1 = ax.get_ylim()
  dy = (fy(y1) - fy(y0)) * dy_frac

  for i, (x, y, lab, c, xpos, side) in enumerate(ent):
    xlo, xhi = fx(x.min()), fx(x.max())
    if xpos is None:  # spread across lines
      frac = (i + 1) / (n + 1)
      xt = fxi(xlo + frac * (xhi - xlo))
    else:
      xt = xpos
    yt = fyi(np.interp(fx(xt), fx(x), fy(y)))  # y on the line at xt

    rot = 0
    if rotate:  # match slope in pixels
      dxt = fxi(fx(xt) + 0.01 * (xhi - xlo))
      ya = fyi(np.interp(fx(xt), fx(x), fy(y)))
      yb = fyi(np.interp(fx(dxt), fx(x), fy(y)))
      p1 = ax.transData.transform((xt, ya))
      p2 = ax.transData.transform((dxt, yb))
      rot = np.degrees(np.arctan2(p2[1] - p1[1], p2[0] - p1[0]))

    below = side == "below"
    ypos = fyi(fy(yt) - dy) if below else fyi(fy(yt) + dy)
    va = "top" if below else "bottom"
    txt = ax.text(xt,
                  ypos,
                  lab,
                  color=c,
                  fontsize=fontsize,
                  ha="center",
                  va=va,
                  fontweight="bold",
                  rotation=rot,
                  rotation_mode="anchor",
                  clip_on=False,
                  zorder=5)
    if halo:
      txt.set_path_effects([pe.withStroke(linewidth=halo, foreground="white")])


base = get("variant_baseline")

# plot 1: the two precision tracks (this is the main results plot)
# double tops out ~7, float goes to 24. also shows the avx2-float stuff already does most of the work
tracks = [
  ("variant_baseline", "Scalar Baseline (double)", "#7a7a7a", "-"),
  ("variant_AVX512_1", "AVX-512 (double)", "#4f86d6", "-"),
  ("variant_inlining_7_SoA_AVX2", "AVX-256 (double)", "#9bbbe0", "--"),
  ("variant_inlining_8_SoA_AVX512_polyexp_moreavx",
   "AVX-512 + PolyExp (double)", "#2f5d9e", "-"),
  ("variant_float_inlining_SoA_AVX2_polyexp", "AVX-256 + PolyExp (float)",
   "#e58aa6", "--"),
  ("variant_float_inlining_SoA_AVX2_polyexp_rsqrt_loopopt",
   "AVX-256 + PolyExp + RSQRT (float)", "#d05a82", "--"),
  ("variant_float_inlining_SoA_AVX512_polyexp_moreavx_ilp_2_rsqrt_morersqrt",
   "AVX-512 + PolyExp + RSQRT + ILP (float)", maroon, "-"),
]

tick_vals = np.array([50, 100, 200, 500, 1000, 2000, 5000, 10000])
target_positions = np.arange(len(tick_vals))


def map_x_to_linear(x_val):
  return np.interp(x_val, tick_vals, target_positions)


fig, ax = plt.subplots(figsize=(9, 4.7))
ends = []

for v, lab, c, ls in tracks:
  x_raw, y = getline(v)
  x_mapped = map_x_to_linear(x_raw)

  ax.plot(x_mapped,
          y,
          ls,
          color=c,
          marker="o",
          ms=4,
          lw=2.4 if c == maroon else 1.8)

  ends.append((x_mapped, y, lab, c))

tick_vals = [50, 100, 200, 500, 1000, 2000, 5000, 10000]
target_positions = list(range(len(tick_vals)))

start_x_pos = 0

for yv, t in [(PEAK_D, f"Peak π par. w/ doubles ({PEAK_D} flops/cycle)"),
              (PEAK_F, f"Peak π par. w/ floats ({PEAK_F} flops/cycle)")]:
  ax.axhline(yv, ls=":", color="#999", lw=1.2)
  ax.text(start_x_pos + 0.2,
          yv * 1.03,
          t,
          fontsize=9,
          color="#888",
          fontweight="bold")

ax.set_yscale("log", base=2)

ax.set_xticks(target_positions)
ax.set_xticklabels(["50", "100", "200", "500", "1k", "2k", "5k", "10k"],
                   fontweight="bold",
                   fontsize=9)

ax.set_xlim(-0.5, len(tick_vals) - 0.5)

ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, _: str(int(x))))
for label in ax.get_yticklabels():
  label.set_fontweight('bold')
  label.set_fontsize(9)
ax.set_ylim(1, 64)
ax.set_xlabel("Number of Pedestrians", fontweight="bold", fontsize=9)
ax.grid(color="#e9e9e9")
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
plt.title(
  'Performance [flops/cycle]\nCPU: AMD Ryzen AI 7 350 w/ Radeon 860M @ 2 GHz\nCompiler: g++ (Debian 15.2.0-17) 15.2.0\nFlags: -O3 -march=native -std=c++23 -Wall -Wextra -flto',
  loc='left',
  size=9,
  fontweight="bold")
plt.tight_layout()

label_lines_above(ax, ends, fontsize=7.5, dy_frac=0.012)

plt.savefig("analysis/fig_tracks.svg", bbox_inches="tight")

# plot 2: all te milestones + the symmetry one that got slower
bars = [
  ("baseline", "variant_baseline"),
  ("AVX2 (double)", "variant_inlining_7_SoA_AVX2"),
  ("AVX-512 (double)", "variant_AVX512_1"),
  ("+ polyexp", "variant_inlining_8_SoA_AVX512_polyexp_moreavx"),
  ("+ float", "variant_float_inlining_SoA_AVX512_polyexp_moreavx"),
  ("+ rsqrt & ILP",
   "variant_float_inlining_SoA_AVX512_polyexp_moreavx_ilp_2_rsqrt_morersqrt"),
  ("symmetry\n(half the pairs)", "variant_inlining_7_symmetry_2"),
]
names = [b[0] for b in bars]
spd = [get(b[1]) / base for b in bars]
cols = [
  "#7a7a7a", "#3b6fb0", "#2f9e8f", "#e08a1e", "#c0436a", maroon, "#cf2b2b"
]

fig, ax = plt.subplots(figsize=(9, 4.7))
yp = np.arange(len(names))[::-1]
ax.barh(yp, spd, color=cols, height=0.66)
for yy, s in zip(yp, spd):
  ax.text(s + 0.25,
          yy,
          str(round(s, 1)) + "x",
          va="center",
          fontsize=11,
          fontweight="bold",
          color="#cf2b2b" if s < 1 else "#222")
ax.axvline(1, ls="--", color="#888", lw=1)
ax.set_yticks(yp)
ax.set_yticklabels(names)
ax.set_xlim(0, 26)
ax.set_xlabel("speedup over baseline  (n=5000)")
#ax.set_title("optimization milestones (and the one that backfired)", loc="left", fontsize=13)
ax.grid(axis="x", color="#eee")
ax.grid(axis="y", visible=False)
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
plt.tight_layout()
plt.savefig("analysis/fig_ladder.png", dpi=170, bbox_inches="tight")


# plot 3: roofline (n = 5000).  same 7 variants as fig_tracks, same colors, inline labels.
# float variants move HALF the bytes -> 2x the operational intensity of the double ones
# (I_double ≈ 2.8k, I_float ≈ 5.6k), so each point uses its OWN I.
# memory roof from a measured single-thread read: sysbench 37618 MiB/s -> ~19.7 bytes/cycle @ 2 GHz.
def getI(v, n=5000):
  return df[(df.variant == v) & (df.n == n)]["I"].values[0]


bw = ((37967.45 + 38536.77 + 38249.60) /
      3) * 1024**2 / 2e9  # ≈ 20 bytes/cycle @ 2 GHz

# (variant, label, color, group)  group "d" = double column, "f" = float column
roof = [
  ("variant_baseline", "Scalar Baseline (double)", "#7a7a7a", "d"),
  ("variant_inlining_7_SoA_AVX2", "AVX-256 (double)", "#9bbbe0", "d"),
  ("variant_AVX512_1", "AVX-512 (double) ", "#4f86d6", "d"),
  ("variant_inlining_8_SoA_AVX512_polyexp_moreavx",
   "AVX-512 + PolyExp (double)", "#2f5d9e", "d"),
  ("variant_float_inlining_SoA_AVX2_polyexp", "AVX-256 + PolyExp (float)",
   "#e58aa6", "f"),
  ("variant_float_inlining_SoA_AVX2_polyexp_rsqrt_loopopt",
   "AVX-256 + PolyExp + RSQRT (float)", "#d05a82", "f"),
  ("variant_float_inlining_SoA_AVX512_polyexp_moreavx_ilp_2_rsqrt_morersqrt",
   "AVX-512 + PolyExp + RSQRT + ILP (float)", maroon, "f"),
]
LABX = {"d": 2400, "f": 4900}  # label right-edge sits just left of each column

fig, ax = plt.subplots(figsize=(9.5, 5.0))
xx = np.logspace(-1, 4.2, 200)
ax.plot(xx, np.minimum(bw * xx, PEAK_F), color="#bbbbbb", lw=1.4)
ax.text(0.165,
        4,
        "DRAM Read β seq. (≈ 20.0 bytes/cycle)",
        fontsize=8,
        color="#888",
        fontweight="bold",
        rotation=46,
        ha="left")

for pk, t in [(SCALAR, f"Peak π seq. ({SCALAR} flops/cycle)"),
              (PEAK_D, f"Peak π par. w/ doubles ({PEAK_D} flops/cycle)"),
              (PEAK_F, f"Peak π par. w/ floats ({PEAK_F} flops/cycle)")]:
  ax.axhline(pk, ls="--", color="#aaaaaa", lw=1.0)
  ax.text(2.5, pk * 1.05, t, fontsize=8, color="#888", fontweight="bold")

for v, lab, c, grp in roof:
  x, y = getI(v), get(v)
  ax.plot(x, y, "o", ms=10, color=c, mec="white", mew=1.4, zorder=5)
  t = ax.text(LABX[grp],
              y,
              lab,
              color=c,
              fontsize=8,
              fontweight="bold",
              ha="right",
              va="center",
              zorder=6)
  t.set_path_effects([pe.withStroke(linewidth=2.5, foreground="white")])

ax.text(2812, 0.95, "I ≈ 2.8K", fontsize=8, color="#777", ha="center", va="top")
ax.text(5625, 10, "I ≈ 5.6K", fontsize=8, color="#777", ha="center", va="top")

ax.set_yscale("log", base=2)
ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, _: str(int(x))))


def fraction_formatter(x, pos):
  if x == 0:
    return '0'
  log2_x = np.log2(x)
  if abs(log2_x - round(log2_x)) < 1e-6:
    exponent = int(round(log2_x))
    if exponent < 0:
      denominator = 2**(-exponent)
      return f'1/{denominator}'
    elif exponent == 0:
      return '1'
    else:
      return str(int(x))
  return f'{x:.2f}'


formatter = ticker.FuncFormatter(fraction_formatter)

plt.xscale('log', base=2)
plt.gca().xaxis.set_major_formatter(formatter)

ax.set_xlim(0.1, 16000)
ax.set_ylim(0.6, 70)
ax.set_xlabel("Operational Intensity [flops/byte]",
              fontweight="bold",
              fontsize=9)
# ax.set_ylabel("Performance [flops/cycle]", fontweight="bold", fontsize=9)
plt.title(
  'Performance [flops/cycle]\nCPU: AMD Ryzen AI 7 350 w/ Radeon 860M @ 2 GHz\nCompiler: g++ (Debian 15.2.0-17) 15.2.0\nFlags: -O3 -march=native -std=c++23 -Wall -Wextra -flto',
  loc='left',
  size=9,
  fontweight="bold")
ax.grid(color="#ededed")
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
plt.tight_layout()
plt.savefig("analysis/fig_roofline.svg", bbox_inches="tight")

# plot 4: exp vs the polynomial approx (just to show its basically exact)
x = np.linspace(-6, 0, 300)
ex = np.exp(x)
co = np.polyfit(x, ex, 5)  # degree 5 same as in the kernel
ap = np.polyval(co, x)

fig, ax = plt.subplots(figsize=(4.6, 3.7))
ax.plot(x, ex, color="#222", lw=2.4, label="exp(x)")
ax.plot(x, ap, "--", color=maroon, lw=2, label="poly approx")
ax.set_xlabel("x")
ax.set_title("exp -> polynomial", loc="left", fontsize=12)
ax.grid(color="#ededed")
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ax.legend(frameon=False, fontsize=10)
plt.tight_layout()
plt.savefig("analysis/fig_exp.png", dpi=170, bbox_inches="tight")

# plot 5: the milestone ladder as lines (this is plot0 + line1, line2, ... for the doc).
# one line per relevant optimization step across n, baseline at the bottom.
# unlike fig_tracks this includes the scalar steps so you can see where each jump comes from.
# tuple: (variant, label, color, linestyle, xpos|None, side) where side is
# "above" or "below" -> alternate the crowded float labels above/below their line.
milestones = [
  ("variant_baseline", "baseline", "#7a7a7a", "-", None, "above"),
  #("variant_cutoff_callsite",                                                "baseline + cutoff (best scalar)",         "#9e8f2f", "-",  None, "below"),
  ("variant_inlining_7_SoA_AVX2", "AVX2 (double)", "#3b6fb0", "--", None,
   "above"),
  ("variant_inlining_7_SoA_AVX2_polyexp", "AVX2 + polyexp (double)", "#2f9e8f",
   "--", None, "below"),
  ("variant_inlining_8_SoA_AVX512_polyexp_moreavx", "AVX-512 (double)",
   "#2f5d9e", "-", None, "below"),
  ("variant_float_inlining_SoA_AVX2_polyexp", "float AVX2 + polyexp", "#e58aa6",
   "--", 600, "above"),
  ("variant_float_inlining_SoA_AVX2_polyexp_rsqrt_loopopt",
   "float AVX2 + rsqrt", "#e08a1e", "--", 2200, "below"),
  ("variant_float_inlining_SoA_AVX512_polyexp_moreavx", "float AVX-512",
   "#c0436a", "-", 1500, "above"),
  ("variant_float_inlining_SoA_AVX512_polyexp_moreavx_ilp_2_rsqrt_morersqrt",
   "float AVX-512 + rsqrt + ILP (best)", maroon, "-", 900, "above"),
]

fig, ax = plt.subplots(figsize=(9.5, 5.2))
ends = []
for v, lab, c, ls, xpos, side in milestones:
  x, y = getline(v)
  if len(x) == 0:  # variant not in results.csv -> skip instead of crashing
    print("skip (not in csv):", v)
    continue
  ax.plot(x, y, ls, color=c, marker="o", ms=4, lw=2.4 if c == maroon else 1.7)
  ends.append((x, y, lab, c, xpos, side))

for yv, t in [(PEAK_D, "double peak (16)"), (PEAK_F, "float peak (32)")]:
  ax.axhline(yv, ls=":", color="#999", lw=1.2)
  ax.text(52, yv * 1.03, t, fontsize=9.5, color="#888")

ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xticks([50, 100, 200, 500, 1000, 2000, 5000, 10000])
ax.set_xticklabels(["50", "100", "200", "500", "1k", "2k", "5k", "10k"])
ax.set_yticks([1, 2, 4, 8, 16, 32, 64])
ax.set_yticklabels(["1", "2", "4", "8", "16", "32", "64"])
ax.set_xlim(46, 13000)
ax.set_ylim(0.8, 40)
ax.set_xlabel("Number of Pedestrians")
ax.set_ylabel("Performance [flops/cycle]")
ax.grid(color="#e9e9e9")
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
plt.tight_layout()
label_lines_above(ax, ends, fontsize=6.5, dy_frac=0.02, halo=0)
plt.savefig("analysis/fig_milestones.png", dpi=170, bbox_inches="tight")

#plt.show()

# plot 6: wall-clock speedup over baseline (runtime, not performance).
# the fair view for optimizations that change the flop count -- cutoff skips ~40% of the
# inner loop, symmetry does half the pairs -- so flops/cycle misrepresents them. T_base /
# T_variant compares time-to-solve the same n regardless of how many flops each one does.
# curated to the milestones that carry the story + symmetry as the cautionary line.
basec = df[df.variant == "variant_baseline"].set_index("n")["cycles"].to_dict()


def speedline(v):
  s = df[df.variant == v].sort_values("n")  # pyright: ignore[reportCallIssue]
  ns = [n for n in s.n.values if n in basec]  # baseline has no n=10000 run
  sp = [basec[n] / s[s.n == n]["cycles"].values[0] for n in ns]
  return np.array(ns, float), np.array(sp, float)


# (variant, label, color, linestyle, xpos|None, side)
speed = [
  ("variant_cutoff_callsite", "Scalar Cutoff (double)", "#9e8f2f", "-", 260,
   "above"),
  ("variant_inlining_7_symmetry_2", "Symmetry (double)", "#cf2b2b", "-", 320,
   "below"),
  ("variant_inlining_7_SoA_AVX2", "AVX-256 (double)", "#9bbbe0", "--", 650,
   "below"),
  ("variant_inlining_8_SoA_AVX512_polyexp_moreavx",
   "AVX-512 + PolyExp (double)", "#2f5d9e", "-", 430, "above"),
  ("variant_float_inlining_SoA_AVX2_polyexp_rsqrt_loopopt",
   "AVX-256 + PolyExp + RSQRT (float)", "#d05a82", "--", 950, "below"),
  ("variant_float_inlining_SoA_AVX512_polyexp_moreavx_ilp_2_rsqrt_morersqrt",
   "AVX-512 + PolyExp + RSQRT + ILP (float)", maroon, "-", 680, "above"),
]

fig, ax = plt.subplots(figsize=(9, 4.7))
ends = []
for v, lab, c, ls, xpos, side in speed:
  x, y = speedline(v)
  if len(x) == 0:
    print("skip (not in csv):", v)
    continue
  ax.plot(x, y, ls, color=c, marker="o", ms=4, lw=2.4 if c == maroon else 1.8)
  ends.append((x, y, lab, c, xpos, side))
ax.axhline(1, ls=":", color="#999", lw=1.3)
ax.text(52, 1.03, "baseline (1x)", fontsize=9, color="#888")
ax.set_xscale("log")
ax.set_yscale("log")

xlim_log = np.log10(ax.get_xlim())
num_ticks = 7
linear_positions = np.linspace(xlim_log[0], xlim_log[1], num_ticks)
tick_positions = 10**linear_positions
ax.set_xticks(tick_positions)
ax.set_xticklabels(["50", "100", "200", "500", "1k", "2k", "5k"])

ax.set_yticks([1, 2, 5, 10, 20])
ax.set_yticklabels(["1x", "2x", "5x", "10x", "20x"])

ax.tick_params(axis='both', which='minor', length=0)
ax.tick_params(axis='both', which='major', length=4)

ax.set_xlim(46, 7000)
ax.set_ylim(0.8, 28)
ax.set_xlabel("Number of Pedestrians", fontweight="bold", fontsize=9)
# ax.set_ylabel("Speedup", fontweight="bold", fontsize=9)
plt.title(
  'Runtime Speedup\nCPU: AMD Ryzen AI 7 350 w/ Radeon 860M @ 2 GHz\nCompiler: g++ (Debian 15.2.0-17) 15.2.0\nFlags: -O3 -march=native -std=c++23 -Wall -Wextra -flto',
  loc='left',
  size=9,
  fontweight="bold")
ax.grid(color="#e9e9e9")
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
plt.tight_layout()
label_lines_above(ax, ends, fontsize=7.5, dy_frac=0.02, halo=2)
plt.savefig("analysis/fig_speedup.svg", bbox_inches="tight")