import matplotlib.pyplot as plt
import numpy as np

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5), dpi=300)

layers = [f'L{i}' for i in range(6)]
partial_r = [-0.307, -0.319, -0.338, -0.364, -0.240, +0.056]

# Left: Partial correlation across layers
ax1.plot(layers, partial_r, marker='o', linewidth=2.5, markersize=8, color='#ef4444', label='Partial r (frequency | ||W_U||)')
ax1.axhline(0, color='gray', linestyle='--', alpha=0.7)
ax1.set_ylabel('Partial Correlation', fontsize=11, fontweight='bold')
ax1.set_xlabel('Transformer Layer', fontsize=11, fontweight='bold')
ax1.set_title('Layer Dynamics Component\n(W_U Regressed Out)', fontsize=12, fontweight='bold', pad=12)
ax1.set_ylim(-0.5, 0.2)
ax1.grid(True, linestyle=':', alpha=0.6)
ax1.legend(loc='lower left', frameon=True)

for i, txt in enumerate(partial_r):
    ax1.annotate(f'{txt:+.3f}', (layers[i], partial_r[i]), textcoords="offset points", xytext=(0,10 if txt>0 else -15), ha='center', fontweight='bold', fontsize=9)

# Right: Static geometry W_U emergence (Init vs Trained)
stages = ['Init (Seed 0)', 'Init (Seed 1)', 'Init (Seed 2)', 'Trained Model']
wu_r = [+0.002, -0.189, -0.166, -0.606]
colors = ['#94a3b8', '#94a3b8', '#94a3b8', '#6366f1']

bars = ax2.bar(stages, wu_r, color=colors, width=0.55, alpha=0.9)
ax2.axhline(0, color='gray', linestyle='--', alpha=0.7)
ax2.set_ylabel('r( ||W_U[k]||, Token Frequency )', fontsize=11, fontweight='bold')
ax2.set_title('Static Unembedding Geometry\n(Learned vs Initialized)', fontsize=12, fontweight='bold', pad=12)
ax2.set_ylim(-0.8, 0.2)
ax2.grid(True, linestyle=':', alpha=0.6)

for bar in bars:
    yval = bar.get_height()
    ax2.text(bar.get_x() + bar.get_width()/2.0, yval - 0.05 if yval < 0 else yval + 0.02, f'{yval:+.3f}', ha='center', va='top' if yval < 0 else 'bottom', fontweight='bold', fontsize=9)

plt.suptitle('Decomposition of Frequency Confound: Static Geometry vs. Layer Dynamics', fontsize=14, fontweight='bold', y=0.98)
plt.tight_layout(rect=[0, 0, 1, 0.93])
plt.savefig('docs/assets/fig2_wu_decomposition.png', dpi=300)
print("Saved Fig 2")