import sys
import os
import numpy as np

sys.path.insert(0, 'tools')
from viz_theme import apply_monochrome, save_svg, load_csv

apply_monochrome()
import matplotlib.pyplot as plt

def read_data(rows):
    ops = []
    refnv = []; refnv_err = []
    ref = []; ref_err = []
    avx2 = []; avx2_err = []
    for row in rows:
        ops.append(row['op'])
        refnv.append(float(row['refnv_sp']))
        refnv_err.append([float(row['refnv_elo']), float(row['refnv_ehi'])])
        ref.append(float(row['ref_sp']))
        ref_err.append([float(row['ref_elo']), float(row['ref_ehi'])])
        avx2.append(float(row['avx2_sp']))
        avx2_err.append([float(row['avx2_elo']), float(row['avx2_ehi'])])
    
    refnv_err = np.array(refnv_err).T
    ref_err = np.array(ref_err).T
    avx2_err = np.array(avx2_err).T
    return ops, refnv, refnv_err, ref, ref_err, avx2, avx2_err

params = [("ML-KEM-512", "decomp_mlkem512.csv"),
          ("ML-KEM-768", "decomp_mlkem768.csv"),
          ("ML-KEM-1024", "decomp_mlkem1024.csv")]

fig, axes = plt.subplots(1, 3, figsize=(12, 4), sharey=True)

bar_width = 0.25
colors = ['#333333', '#777777', '#bbbbbb']
labels = ['O3 (no auto-vec)', 'O3 + auto-vec', 'O3 + hand SIMD']

for i, (title, filepath) in enumerate(params):
    ops, refnv, refnv_err, ref, ref_err, avx2, avx2_err = read_data(load_csv(filepath))
    ax = axes[i]
    x = np.arange(len(ops))
    
    ax.bar(x - bar_width, refnv, bar_width, label=labels[0], color=colors[0], yerr=refnv_err, edgecolor='none')
    ax.bar(x, ref, bar_width, label=labels[1], color=colors[1], yerr=ref_err, edgecolor='none')
    ax.bar(x + bar_width, avx2, bar_width, label=labels[2], color=colors[2], yerr=avx2_err, edgecolor='none')
    
    ax.set_title(title, pad=10)
    ax.set_xticks(x)
    
    # Format xticklabels to replace iDec, iEnc, iKeypair, gena
    display_ops = [op.replace('gena', 'gen_a') for op in ops]
    ax.set_xticklabels(display_ops, rotation=45, ha='right')
    
    ax.set_yscale('log')
    if i == 0:
        ax.set_ylabel("Speedup over -O0 ($\\times$)")
    
    # Tick formatting
    ax.set_ylim(bottom=1, top=500)
    import matplotlib.ticker as ticker
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda y, pos: f"${int(y)}\\times$"))

axes[-1].legend(loc='upper right', frameon=False, fontsize='small')
save_svg(
    fig,
    alt=(
        "Grouped bar chart of cumulative speedup over unoptimized reference "
        "code, for nine ML-KEM operations at each of the three parameter "
        "sets."
    ),
    desc=(
        "Three panels, one per parameter set. Each operation has three bars "
        "on a logarithmic axis: compiler O3 without auto-vectorization, O3 "
        "with it, and hand-written AVX2. The first two bars are nearly the "
        "same height everywhere -- around 3 to 4 times -- so auto- "
        "vectorization adds almost nothing. The third bar is far taller for "
        "the polynomial operations, reaching roughly 200 times for INVNTT, "
        "and the gap narrows for the sampling operations at the right of "
        "each panel."
    ),
)