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 4, double 16, float 32 flops/cycle
SCALAR = 4
PEAK_D = 16
PEAK_F = 32
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", "baseline (scalar)", "#7a7a7a", "-"),
 ("variant_AVX512_1", "double, AVX-512 ", "#4f86d6", "-"),
 ("variant_inlining_7_SoA_AVX2", "double, AVX-256", "#9bbbe0", "--"),
 ("variant_inlining_8_SoA_AVX512_polyexp_moreavx", "double, AVX-512 + polyexp", "#2f5d9e", "-"),
 ("variant_float_inlining_SoA_AVX2_polyexp", "float, AVX-256 + polyexp", "#e58aa6", "--"),
 ("variant_float_inlining_SoA_AVX2_polyexp_rsqrt_loopopt", "float, AVX-256 + polyexp + rsqrt", "#d05a82", "--"),
 ("variant_float_inlining_SoA_AVX512_polyexp_moreavx_ilp_2_rsqrt_morersqrt", "float, AVX-512 + rsqrt + ILP (best)", maroon, "-"),
]
fig, ax = plt.subplots(figsize=(9,4.7))
ends = []
for v,lab,c,ls in tracks:
    x,y = getline(v)
    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))

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]); ax.set_yticklabels(["1","2","4","8","16","32"])
ax.set_xlim(46,13000); ax.set_ylim(1,40)
ax.set_xlabel("number of pedestrians  n")
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=7.5, dy_frac=0.012)
plt.savefig("analysis/fig_tracks.png", dpi=170, 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 = 37618.01 * 1024**2 / 2e9          # MiB/s -> bytes/cycle at 2 GHz  (≈ 19.72)

# (variant, label, color, group)  group "d" = double column, "f" = float column
roof = [
 ("variant_baseline", "baseline (scalar)", "#7a7a7a", "d"),
 ("variant_inlining_7_SoA_AVX2", "double, AVX-256", "#9bbbe0", "d"),
 ("variant_AVX512_1", "double, AVX-512 ", "#4f86d6", "d"),
 ("variant_inlining_8_SoA_AVX512_polyexp_moreavx", "double, AVX-512 + polyexp", "#2f5d9e", "d"),
 ("variant_float_inlining_SoA_AVX2_polyexp", "float, AVX-256 + polyexp", "#e58aa6", "f"),
 ("variant_float_inlining_SoA_AVX2_polyexp_rsqrt_loopopt", "float, AVX2 + polyexp + rsqrt", "#d05a82", "f"),
 ("variant_float_inlining_SoA_AVX512_polyexp_moreavx_ilp_2_rsqrt_morersqrt", "float, AVX-512 + rsqrt + ILP (best)", 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.115, 0.85, "mem roof: DRAM read ≈ 37.6 GiB/s (≈19.7 B/cycle)", fontsize=8.3, color="#9a9a9a", ha="left")

for pk,t in [(SCALAR,"scalar (4)"), (PEAK_D,"double (16)"), (PEAK_F,"float (32)")]:
    ax.axhline(pk, ls="--", color="#aaaaaa", lw=1.0)
    ax.text(0.11, pk*1.05, t+" peak", fontsize=9, color="#888")

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.74, "I ≈ 2.8k", fontsize=8, color="#777", ha="center", va="top")
ax.text(5625, 0.74, "I ≈ 5.6k", fontsize=8, color="#777", ha="center", va="top")

ax.set_xscale("log"); ax.set_yscale("log")
ax.set_xlim(0.1, 16000); ax.set_ylim(0.6, 46)
ax.set_xlabel("operational intensity  I  [flops/byte]")
ax.set_ylabel("performance [flops/cycle]")
ax.grid(color="#ededed")
ax.spines["top"].set_visible(False); ax.spines["right"].set_visible(False)
plt.tight_layout()
plt.savefig("analysis/fig_roofline.png", dpi=170, 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]); ax.set_yticklabels(["1","2","4","8","16","32"])
ax.set_xlim(46,13000); ax.set_ylim(0.8,40)
ax.set_xlabel("number of pedestrians  n")
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()
print("done")

# 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")
    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 (approx)", "#9e8f2f", "-",  260, "above"),
 ("variant_inlining_7_symmetry_2", "symmetry (half the pairs)", "#cf2b2b", "-",  320, "below"),
 ("variant_inlining_7_SoA_AVX2", "double, AVX2", "#9bbbe0", "--", 650, "below"),
 ("variant_inlining_8_SoA_AVX512_polyexp_moreavx", "double, AVX-512 + polyexp", "#2f5d9e", "-",  430, "above"),
 ("variant_float_inlining_SoA_AVX2_polyexp_rsqrt_loopopt", "float, AVX2 + polyexp + rsqrt", "#d05a82", "--", 950, "below"),
 ("variant_float_inlining_SoA_AVX512_polyexp_moreavx_ilp_2_rsqrt_morersqrt", "float, AVX-512 + rsqrt + ILP (best)", 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")
ax.set_xticks([50,100,200,500,1000,2000,5000]); 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.set_xlim(46,7000); ax.set_ylim(0.8,28)
ax.set_xlabel("number of pedestrians  n")
ax.set_ylabel("speedup over baseline  (wall-clock cycles)")
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.png", dpi=170, bbox_inches="tight")