from pathlib import Path
import csv
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

source=Path('/home/eai/code/training_recipe/2026-09-21_b1k_task00_pi05_codegen_success_sft_h20/config/validation_history.csv')
out=Path(__file__).parent
rows=list(csv.DictReader(source.open()))
steps=[int(r['step']) for r in rows]
val=[float(r['validation_loss']) for r in rows]
train=[float(r['train_eval_loss']) for r in rows]
best=min(range(len(val)),key=val.__getitem__)
plt.rcParams.update({'font.family':'DejaVu Sans','font.size':10,'axes.spines.top':False,'axes.spines.right':False,'axes.labelcolor':'#334155','text.color':'#0f172a','xtick.color':'#475569','ytick.color':'#475569','axes.edgecolor':'#cbd5e1'})
fig,axs=plt.subplots(1,2,figsize=(13.4,5.6),gridspec_kw={'width_ratios':[1,1.2]})
fig.subplots_adjust(left=.075,right=.976,top=.77,bottom=.21,wspace=.25)
fig.suptitle('Pi0.5 · turning_on_radio',x=.075,y=.975,ha='left',fontsize=20,fontweight='bold')
fig.text(.075,.9,'Best checkpoint: step 6,000  |  Training stopped: step 10,000',fontsize=12,color='#475569')
for ax in axs:
 ax.plot(steps,val,color='#2563eb',lw=2.5,marker='o',ms=3.6,label='Validation · 26 held-out instances')
 ax.plot(steps,train,color='#d97706',lw=2.3,marker='o',ms=3.3,label='Training subset · 26 instances')
 ax.axvline(steps[best],color='#64748b',lw=1.15,ls='--')
 ax.set_xlabel('Training step',labelpad=9)
 ax.set_ylabel('Mean flow loss · 23 action dimensions',labelpad=9)
 ax.xaxis.set_major_formatter(FuncFormatter(lambda x,p:f'{x:,.0f}'))
 ax.grid(axis='y',color='#e2e8f0',lw=.8)
 ax.set_axisbelow(True)
axs[0].set_title('Full run · logarithmic loss scale',loc='left',fontsize=12,pad=14)
axs[0].set_yscale('log');axs[0].set_xlim(-150,10300);axs[0].set_ylim(.0035,.8)
axs[0].set_xticks([0,2000,4000,6000,8000,10000])
axs[1].set_title('Detail after step 1,000 · linear loss scale',loc='left',fontsize=12,pad=14)
axs[1].set_xlim(1000,10300);axs[1].set_ylim(.0035,.0155)
axs[1].set_xticks([1000,3000,6000,8000,10000])
axs[1].axvspan(6000,10000,color='#fee2e2',alpha=.32,zorder=-1)
axs[1].scatter([6000],[val[best]],s=80,facecolor='white',edgecolor='#2563eb',lw=2,zorder=5)
axs[1].annotate(f'Best: {val[best]:.6f}',xy=(6000,val[best]),xytext=(4100,.01125),fontsize=10,color='#1d4ed8',arrowprops={'arrowstyle':'-','color':'#2563eb','lw':1.1})
axs[1].annotate(f'Final: {val[-1]:.6f}\n+{(val[-1]/val[best]-1)*100:.1f}% vs. best',xy=(10000,val[-1]),xytext=(7850,.0124),fontsize=10,color='#1d4ed8',arrowprops={'arrowstyle':'-','color':'#2563eb','lw':1.1})
handles,labels=axs[0].get_legend_handles_labels()
fig.legend(handles,labels,loc='lower left',bbox_to_anchor=(.07,.075),frameon=False,ncol=2,fontsize=10)
fig.text(.075,.035,'Comparable evaluation losses: no augmentation, fixed noise, episode averages. Points every 500 steps; no smoothing.',fontsize=9,color='#64748b')
fig.savefig(out/'loss_curve.png',dpi=180,facecolor='white')
fig.savefig(out/'loss_curve.svg',facecolor='white')
(out/'validation_history.csv').write_bytes(source.read_bytes())
print(out/'loss_curve.png')