from pathlib import Path
import csv,json,sys
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
ROOT=Path(__file__).resolve().parents[1]
plt.rcParams.update({'font.size':9,'axes.spines.top':False,'axes.spines.right':False,'savefig.dpi':180})
out=ROOT/'figures_v3';out.mkdir(exist_ok=True)
r=json.loads((ROOT/'results_v3/innovation_report.json').read_text())
rows=list(csv.DictReader((ROOT/'results_v3/innovation_raw.csv').open()))
names={'epoch_eliminate':'Reset on geometry change','v2_guarded':'v2 contradiction guard','certificate_iid':'Certificates, iid draws','certificate_eliminate':'Certificates, elimination'}
colors=['#a96633','#757c89','#4989a1','#116e58']
fig,axes=plt.subplots(1,2,figsize=(9.4,3.3))
for method,color in zip(names,colors):
    sub=[x for x in r['summary'] if x['method']==method]
    phases=list(r['protocol']['phases'])
    axes[0].plot(range(len(phases)),[x['expected_mse'] for x in sub],marker='o',label=names[method],color=color)
    axes[1].plot(range(len(phases)),[x['rays_per_receiver'] for x in sub],marker='o',color=color)
for ax in axes:
    ax.set_xticks(range(7),['Cold','Warm','Revisit','Relight','Smooth','Jump','Return'],rotation=30,ha='right');ax.grid(alpha=.18)
axes[0].set_yscale('log');axes[0].set_ylabel('Conditional expected linear RGB MSE');axes[1].set_ylabel('Physical queries / receiver / frame')
handles,labels=axes[0].get_legend_handles_labels();fig.legend(handles,labels,loc='lower center',ncol=2,fontsize=7,frameon=False);fig.tight_layout(rect=[0,.14,1,1])
for ext in ('png','pdf'):fig.savefig(out/f'innovation_results.{ext}')
plt.close(fig)

fig,axes=plt.subplots(1,2,figsize=(9.4,3.1))
for method,color in zip(names,colors):
    frames=sorted(set(int(x['frame']) for x in rows))
    mean=[np.mean([float(x['expected_mse']) for x in rows if x['method']==method and int(x['frame'])==f]) for f in frames]
    axes[0].plot(frames,mean,label=names[method],color=color)
axes[0].set_yscale('log');axes[0].set_xlabel('Observed frame');axes[0].set_ylabel('Expected linear RGB MSE')
for p in (40,56,64):axes[0].axvline(p,color='gray',lw=.6,ls='--')
axes[0].text(41,.004,'Smooth',fontsize=8);axes[0].text(57,.004,'Jump',fontsize=8);axes[0].text(65,.004,'Return',fontsize=8)
q=json.loads((ROOT/'results_v3/queries_report.json').read_text())
axes[1].barh(['Fresh shared visibility','Certificate reuse'],[q['shared_fresh_baseline_queries'],q['total_queries_including_initialization']],color=['#757c89','#116e58'])
axes[1].set_xlabel('Physical queries, including certificate initialization');axes[1].ticklabel_format(style='sci',axis='x',scilimits=(6,6))
axes[1].text(q['total_queries_including_initialization']+30000,1,f"{100*q['amortized_query_reduction_vs_shared']:.1f}% fewer",va='center',fontsize=8)
fig.tight_layout()
for ext in ('png','pdf'):fig.savefig(out/f'innovation_dynamics.{ext}')
plt.close(fig)

data=np.load(ROOT/'results_v3/queries_frames.npz')
fig,axes=plt.subplots(3,3,figsize=(9.4,4.8))
for i,t in enumerate((0.,.5,1.)):
    for j in range(3):
        im=data[f't{t}_appearance{j}'];axes[i,j].imshow(np.clip(im,0,1)**(1/2.2));axes[i,j].set_axis_off()
        axes[i,j].set_title(f'Known time {t:g}, appearance {j+1}',fontsize=8)
fig.suptitle('One coarse world state: finer spatial queries, prescribed motion, three readouts',fontsize=10)
fig.tight_layout()
for ext in ('png','pdf'):fig.savefig(out/f'query_readouts.{ext}')
plt.close(fig)
print('Saved three scientific figures as PNG and PDF')