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_main = ['RTX 40 Series¹', 'RTX 40 Series²', 'RTX 40 Series²', '4 TET-FPGA', '4 TET-FPGA'][::-1]
subtexts = ['60FPS', '60FPS', '360FPS', '60FPS (Worst-Case)', '60FPS'][::-1]
note = ['competitor', 'competitor', 'competitor', 'ours', 'ours'][::-1]

invert_faster = False
invert_ours = True

# Combine main titles and subtexts with newline characters
labels = [f"{main}\n{sub}" for main, sub in zip(labels_main, subtexts)]

values = [57, 35, 17, 24.8, 12.8][::-1]          # Base values (e.g., Delay Part 1)
values_extra = [0, 0, 0, 25, 17][::-1]            # Extra add-on values (e.g., Delay Part 2)

# Define your unit here (add a space at the start so it looks clean)
unit_label = "Delay in [ms]"  

# Style choices
brand_color = '#00E5FF'
competitor_color = '#00FF00'
extra_color = '#FF0000'                          # Color for extra segments (ours)
extra_color_extra = "#0E0E0E"                    # Color for extra segments (competitor)

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

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

# Create color lists dynamically based on the 'note' array
colors = [brand_color if n == 'ours' else competitor_color for n in note]
extra_colors = [extra_color if n == 'ours' else extra_color_extra for n in note]

# Plot base bars and stacked extra bars
bars_base = ax.barh(labels, values, color=colors, height=0.5, zorder=3)
bars_extra = ax.barh(labels, values_extra, left=values, color=extra_colors, height=0.5, zorder=3)

# Calculate total combined values for limits and bracket positioning
total_values = [v + e for v, e in zip(values, values_extra)]
max_val = max(total_values)

# 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='x', colors='white', labelsize=12)
ax.tick_params(axis='y', colors='white', labelsize=18)

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 and text inside/outside the extra segment
for bar_base, bar_ext, v_base, v_ext in zip(bars_base, bars_extra, values, values_extra):
    base_width = bar_base.get_width()
    extra_width = bar_ext.get_width()
    total_width = base_width + extra_width
    
    if v_ext > 0:
        # 1. Label at the very end of the total bar (just base+extra numbers)
        end_txt = ax.text(total_width + (max_val * 0.02), 
                          bar_base.get_y() + bar_base.get_height()/2, 
                          f"{format_val(total_width)} ({format_val(v_base)}+{format_val(v_ext)})", 
                          ha='left', 
                          va='center', color='white', fontsize=10, fontweight='bold')
        end_txt.set_path_effects(black_outline)
        
        # 2. Text "[Peripherals]" centered horizontally inside the red/extra bar segment
        mid_extra_x = base_width + (extra_width / 2.0)
        in_txt = ax.text(mid_extra_x, 
                         bar_base.get_y() + bar_base.get_height()/2, 
                         "[Peripherals]", 
                         ha='center', 
                         va='center', color='white', fontsize=9, fontweight='bold')
        in_txt.set_path_effects(black_outline)
    else:
        # Standard label for bars without an extra segment
        end_txt = ax.text(total_width + (max_val * 0.02), 
                          bar_base.get_y() + bar_base.get_height()/2, 
                          format_val(total_width), 
                          ha='left', 
                          va='center', color='white', fontsize=10, fontweight='bold')
        end_txt.set_path_effects(black_outline)

# ==========================================
# 3. DYNAMIC BRACKET (First vs Last based on totals)
# ==========================================
y_start = bars_base[0].get_y() + bars_base[0].get_height()/2
y_end = bars_base[-1].get_y() + bars_base[-1].get_height()/2
x_line, x_tick = max_val * 1.15, max_val * 1.12 

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)

start_val, end_val = values[0], values[-1]
multiplier = end_val / start_val

if (invert_faster):
    mult_txt = ax.text(x_line + (max_val * 0.03), (y_start+y_end)/2, 
                   f'{(1/multiplier):.1f}x\nFASTER', 
                   ha='left', va='center', color=brand_color, fontsize=18, fontweight='heavy')
else:
    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)

# Set x-axis limit to comfortably fit the stacked bars, labels, and bracket
ax.set_xlim(0, max_val * 1.40)

plt.tight_layout()
plt.subplots_adjust(bottom=0.2)
plt.savefig('dynamic_comparison.png', dpi=300, transparent=True)
plt.show()