import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(figsize=(8, 5), dpi=300)

seeds = ['Seed 0', 'Seed 1', 'Seed 2', 'Cross-Seed Mean']

# q norm values
q_2x = [0.0163, 0.0150, 0.0161, 0.0158]
control = [0.0174, 0.0171, 0.0158, 0.0168]
ctrl_rand = [0.0152, 0.0161, 0.0169, 0.0161]

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

rects1 = ax.bar(x - width, q_2x, width, label='q (2x Loss Weight)', color='#ef4444', alpha=0.85)
rects2 = ax.bar(x, control, width, label='Plain Control', color='#3b82f6', alpha=0.85)
rects3 = ax.bar(x + width, ctrl_rand, width, label='Random Token (2x Loss Weight)', color='#10b981', alpha=0.85)

ax.set_ylabel("Faithful J-Lens Norm of 'q' (L2-L4 Mean)", fontsize=11, fontweight='bold')
ax.set_title("Loss-Reweighting Causal Test Across 3 Seeds", fontsize=13, fontweight='bold', pad=15)
ax.set_xticks(x)
ax.set_xticklabels(seeds, fontsize=11, fontweight='bold')
ax.set_ylim(0.010, 0.022)
ax.grid(True, linestyle=':', alpha=0.6)
ax.legend(loc='upper left', frameon=True)

# Annotate cross-seed mean ratios
ax.annotate('Ratio q/Control: 0.944\n(95% CI: [0.881, 1.019])', xy=(3 - width/2, 0.019),
            ha='center', fontsize=9, fontweight='bold', bbox=dict(boxstyle='round', facecolor='#fee2e2', edgecolor='none'))

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