import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
import numpy as np

ram_beta = 1.91
l1d_l2_beta = 64
l3_beta = 32

pi_scalar = 6.0
pi_vector = 48.0

I_vals = np.logspace(-50, 8, 500, base=2)

plt.figure(figsize=(10, 7))

plt.gca().set_facecolor('lightgrey')

plt.plot(I_vals, ram_beta * I_vals, 'k-', label='_nolegend_')
plt.plot(I_vals, l1d_l2_beta * I_vals, 'k-', label='_nolegend_')
plt.plot(I_vals, l3_beta * I_vals, 'k-', label='_nolegend_')

ram_bound_x_pos = 0.25
ram_bound_y_pos = ram_beta * ram_bound_x_pos + 0.175

plt.text(
  ram_bound_x_pos,
  ram_bound_y_pos,
  f'Read β mem. ({ram_beta} bytes/cycle)',
  color='black',
  va='center',
  ha='center',
  fontsize=8,
  rotation=20.5,
)

l1d_l2_bound_x_pos = 0.25
l1d_l2_bound_y_pos = l1d_l2_beta * l1d_l2_bound_x_pos + 3.5

plt.text(
  l1d_l2_bound_x_pos,
  l1d_l2_bound_y_pos,
  f'Read β L1D/L2 ({l1d_l2_beta} bytes/cycle)',
  color='black',
  va='center',
  ha='center',
  fontsize=8,
  rotation=20.5,
)

l3_bound_x_pos = 0.375
l3_bound_y_pos = l3_beta * l3_bound_x_pos + 3.5

plt.text(
  l3_bound_x_pos,
  l3_bound_y_pos,
  f'Read β L3 ({l3_beta} bytes/cycle)',
  color='black',
  va='center',
  ha='center',
  fontsize=8,
  rotation=20.5,
)

plt.plot(I_vals, [pi_scalar] * len(I_vals), 'k-', label='_nolegend_')
plt.text(
  85,
  pi_scalar + 1.5,
  f'Peak π seq. ({pi_scalar} flops/cycle)',
  color='black',
  va='center',
  ha='center',
  fontsize=8,
)

plt.plot(I_vals, [pi_vector] * len(I_vals), 'k-', label='_nolegend_')
plt.text(
  85,
  pi_vector + 12.5,
  f'Peak π par. ({pi_vector} flops/cycle)',
  color='black',
  va='center',
  ha='center',
  fontsize=8,
)


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:.2f}'
    elif exponent == 0:
      return '1'
    else:
      return str(int(x))
  return f'{x:.2f}'


formatter = FuncFormatter(fraction_formatter)

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

plt.xlabel('Operational Intensity [flops/byte]', fontsize=10)
plt.title("""Roofline Plot for the Social Force Model
CPU: AMD Ryzen AI 7 350 w/ Radeon 860M @ 2 GHz
Compiler: g++ (Debian 15.2.0-16) 15.2.0
Command: ./main 100 10000

Performance [flops/cycle]""",
          loc='left',
          size=10)
plt.scatter([53], [0.0936126],
            c='black',
            marker='o',
            edgecolors='black',
            label='baseline (-O0)')
plt.text(
  53,
  0.0936126 - 0.03,
  f'0.09 flops/cycle',
  color='black',
  va='center',
  ha='center',
  fontsize=8,
)
plt.scatter([53], [0.627801],
            c='black',
            marker='s',
            edgecolors='black',
            label='baseline (-O3 -march=native)')
plt.text(
  53,
  0.627801 - 0.2,
  f'0.63 flops/cycle',
  color='black',
  va='center',
  ha='center',
  fontsize=8,
)

plt.legend(loc='lower right',
           fontsize=11,
           frameon=True,
           facecolor='white',
           edgecolor='black')
plt.grid(True, which="both", ls="-", color='white', alpha=0.8, linewidth=1)

plt.xlim(0.05, 200)
plt.ylim(0.0005, 100)

plt.tight_layout()
plt.show()