from pathlib import Path
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

ROOT=Path(__file__).resolve().parent
rows=json.loads((ROOT/'per_instance.json').read_text())
transitions=json.loads((ROOT/'transitions.json').read_text())
variants=['human','codegen_success','mix_h50_c50']
labels=['Human','Codegen','Mix 50:50']
colors=['#177c76','#ce5835','#6956a7']
fig,axes=plt.subplots(3,3,figsize=(14,8),sharex=True,sharey=True)
for col,instance in enumerate([303,304,312]):
    for row,(v,label,color) in enumerate(zip(variants,labels,colors)):
        ax=axes[row,col]
        r=next(r for r in rows if r['variant']==v and r['instance_id']==instance)
        tr=transitions[f'{v}/{instance}']
        x=[0]+[c['step'] for c in tr]+[r['steps']]
        y=[0]+[c['after'] for c in tr]+[tr[-1]['after'] if tr else 0]
        ax.step(x,y,where='post',color=color,lw=1.5)
        resets=[c for c in tr if c['kind']=='Stage corrected']
        ax.scatter([c['step'] for c in resets],[c['after'] for c in resets],color='#be2432',s=17,zorder=5)
        ax.scatter([r['steps']],[y[-1]],marker='o' if r['success'] else 'x',s=55,color=color,zorder=6)
        outcome='SUCCESS' if r['success'] else 'TIMEOUT'
        ax.text(.03,.96,f"{label}: {outcome}\n4→2 resets: {r['stage4_resets']} | action overrides: {r['action_modified']}",transform=ax.transAxes,va='top',fontsize=9,bbox=dict(facecolor='white',alpha=.9,edgecolor='none'))
        ax.set_ylim(-.2,5.6);ax.set_yticks(range(5));ax.set_xlim(0,3330)
        ax.grid(alpha=.16)
        if row==0:ax.set_title(f'Instance {instance}',weight='bold')
        if col==0:ax.set_ylabel('Wrapper stage')
        if row==2:ax.set_xlabel('Simulator steps')
fig.suptitle('Task 00: stage cycling and action intervention in recorded runs',fontsize=15,weight='bold')
fig.text(.5,.025,'Red dots: heuristic stage reset 4→2. Curves reconstructed from every logged stage transition.\nAction overrides replace the predicted chunk with current state and opened gripper(s).',ha='center',fontsize=10)
fig.tight_layout(rect=[0,.08,1,.95])
fig.savefig(ROOT/'stage_timelines.png',dpi=170)
fig.savefig(ROOT/'stage_timelines.pdf')