import matplotlib.pyplot as plt
import numpy as np

# Set aesthetic styling
plt.style.use('seaborn-v0_8-darkgrid' if 'seaborn-v0_8-darkgrid' in plt.style.available else 'default')
fig, ax = plt.subplots(figsize=(9, 5), dpi=300)

layers = [f'L{i}' for i in range(6)]
proxy_r = [-0.661, -0.673, -0.653, -0.648, -0.562, -0.665]
faithful_r = [-0.643, -0.668, -0.672, -0.685, -0.637, -0.606]
spearman_r = [-0.844, -0.851, -0.825, -0.769, -0.716, -0.808]

x = np.arange(len(layers))
width = 0.25

rects1 = ax.bar(x - width, faithful_r, width, label='Faithful Lens (Pearson r)', color='#6366f1', alpha=0.9)
rects2 = ax.bar(x, proxy_r, width, label='Proxy Lens (Pearson r)', color='#a855f7', alpha=0.7)
rects3 = ax.bar(x + width, spearman_r, width, label='Faithful Lens (Spearman ρ)', color='#06b6d4', alpha=0.8)

ax.set_ylabel('Correlation with Token Frequency', fontsize=12, fontweight='bold', labelpad=10)
ax.set_title('Token Frequency Anti-Correlation Across Layers (nanoGPT 10.65M)', fontsize=14, fontweight='bold', pad=15)
ax.set_xticks(x)
ax.set_xticklabels(layers, fontsize=11)
ax.axhline(0, color='black', linewidth=0.8, linestyle='--')
ax.set_ylim(-1.0, 0.1)
ax.legend(frameon=True, facecolor='#ffffff', edgecolor='none', fontsize=10, loc='lower left')

# Add values above/below bars
for rect in rects1:
    h = rect.get_height()
    ax.annotate(f'{h:.2f}', xy=(rect.get_x() + rect.get_width() / 2, h),
                xytext=(0, 6), textcoords="offset points", ha='center', va='bottom', fontsize=8, color='white', fontweight='bold')

for rect in rects3:
    h = rect.get_height()
    ax.annotate(f'{h:.2f}', xy=(rect.get_x() + rect.get_width() / 2, h),
                xytext=(0, 6), textcoords="offset points", ha='center', va='bottom', fontsize=8, color='white', fontweight='bold')

plt.tight_layout()
plt.savefig('docs/assets/fig1_layer_correlation.png', dpi=300)
print("Saved Fig 1")