import matplotlib.pyplot as plt
import matplotlib.patheffects as path_effects
import matplotlib.ticker as ticker

# ==========================================
# 0. HELPER FUNCTION
# ==========================================
def format_val(x, pos=None):
    """Formats numbers to k/M format for axes and labels"""
    if x >= 1_000_000:
        return f"{x/1_000_000:.1f}".replace('.0', '') + 'M'
    elif x >= 1000:
        return f"{x/1000:.1f}".replace('.0', '') + 'k'
    return str(int(x))

# ==========================================
# 1. EDIT YOUR DATA HERE
# ==========================================
labels = ['Tetris AI test CPU', 'tet_cpu_per_row', '4 tet_cpu_per_row']
values = [41700, 909000, 909000*4]  # e.g., Speed, FPS, Transactions per second (Higher is better)

# Define your unit here (add a space at the start so it looks clean)
unit_label = "Moves per Second [Moves/sec]"  

# Style choices
brand_color = '#00E5FF'
competitor_color = '#888888'

max_val = max(values)
# Bracket will compare the first and last (or you can change these indices)
start_val, end_val = values[0], values[-1]
multiplier = end_val / start_val

black_outline = [path_effects.withStroke(linewidth=3, foreground='black')]

# ==========================================
# 2. SETUP & DYNAMIC PLOTTING
# ==========================================
fig, ax = plt.subplots(figsize=(10, 6)) # Increased height for more labels
fig.patch.set_alpha(0.0)
ax.patch.set_alpha(0.0)

# Create color list: grey for all, brand_color for the last one
colors = [competitor_color] * (len(values) - 1) + [brand_color]
bars = ax.barh(labels, values, color=colors, height=0.5, zorder=3)

# Configure Axis
ax.xaxis.set_major_formatter(ticker.FuncFormatter(format_val))
for spine in ax.spines.values(): spine.set_visible(False)
ax.spines['bottom'].set_visible(True)
ax.spines['bottom'].set_color('white')
ax.spines['bottom'].set_path_effects(black_outline)

ax.tick_params(axis='both', colors='white', labelsize=12)
for label in ax.get_xticklabels() + ax.get_yticklabels():
    label.set_path_effects(black_outline)

ax.xaxis.grid(True, linestyle='--', color='white', alpha=0.3, zorder=0)
xlabel = ax.set_xlabel(unit_label, color='white', fontsize=12, labelpad=12, fontfamily='monospace')
xlabel.set_path_effects(black_outline)

# Add Labels for each bar
for bar, val in zip(bars, values):
    width = bar.get_width()
    txt = ax.text(width - (max_val * 0.02) if width >= (max_val*0.15) else width + (max_val*0.02), 
                  bar.get_y() + bar.get_height()/2, 
                  format_val(val), 
                  ha='right' if width >= (max_val*0.15) else 'left', 
                  va='center', color='white', fontsize=12, fontweight='bold')
    txt.set_path_effects(black_outline)

# ==========================================
# 3. DYNAMIC BRACKET (First vs Last)
# ==========================================
y_start, y_end = bars[0].get_y() + bars[0].get_height()/2, bars[-1].get_y() + bars[-1].get_height()/2
x_line, x_tick = max_val * 1.10, max_val * 1.07 # Shifted further right to accommodate N items

l1, = ax.plot([x_tick, x_line], [y_end, y_end], color=brand_color, lw=2)
l2, = ax.plot([x_tick, x_line], [y_start, y_start], color=brand_color, lw=2)
l3, = ax.plot([x_line, x_line], [y_start, y_end], color=brand_color, lw=2)
for l in [l1, l2, l3]: l.set_path_effects(black_outline)

mult_txt = ax.text(x_line + (max_val * 0.03), (y_start+y_end)/2, 
                   f'{multiplier:.1f}x\nFASTER', 
                   ha='left', va='center', color=brand_color, fontsize=18, fontweight='heavy')
mult_txt.set_path_effects(black_outline)

ax.set_xlim(0, max_val * 1.5)
plt.tight_layout()
plt.subplots_adjust(bottom=0.2)
plt.savefig('dynamic_comparison.png', dpi=300, transparent=True)
plt.show()