"""Plot all training seeds; pending conditions have no invented value."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def plot(directory):
    data=json.loads((directory/'summary.json').read_text())
    plt.rcParams.update({'font.size':10,'font.family':'DejaVu Sans','axes.spines.top':False,
                         'axes.spines.right':False,'svg.fonttype':'none'})
    fig,(ax,pair)=plt.subplots(1,2,figsize=(12,4.8),gridspec_kw={'width_ratios':[1.6,1]})
    colors={'0':'#287b78','10':'#b66d32'}
    names=[('mujoco_cpu','MuJoCo CPU'),('newton_mjwarp','Newton'),('physx_gpu','PhysX GPU'),('genesis','Genesis')]
    for i,(engine,label) in enumerate(names):
        for weight,shift in [('0',-.15),('10',.15)]:
            condition=data['engines'][engine].get(weight)
            if condition is None:
                ax.text(i+.15,47,'Pending',ha='center',va='center',rotation=90,color='#666',fontsize=9)
                continue
            values=condition['success_rate_across_seeds']['new_50']['values']
            xs=[i+shift+(seed-2)*.028 for seed in range(5)]
            ax.scatter(xs,[v*100 for v in values],color=colors[weight],s=39,alpha=.85,zorder=3)
            mean=sum(values)/len(values)*100
            ax.plot([i+shift-.10,i+shift+.10],[mean,mean],color=colors[weight],lw=2.3)
    ax.set_xticks(range(4),[n for _,n in names]);ax.set_ylim(0,105)
    ax.set_ylabel('Success on 50 new starts (%)');ax.set_title('One point per training seed',loc='left',fontweight='bold',pad=14)
    ax.grid(axis='y',alpha=.18);ax.set_axisbelow(True)
    ax.legend([plt.Line2D([],[],marker='o',linestyle='',color=colors[w]) for w in ['0','10']],
              ['Baseline','Force penalty'],loc='lower left',frameon=False)
    physx=data['engines']['physx_gpu']
    palette=['#287b78','#8273a6','#6c8195','#b28b38','#b64d48']
    for seed,color in enumerate(palette):
        values=[physx[w]['seeds'][seed]['subsets']['new_50']['success_rate']*100 for w in ['0','10']]
        pair.plot([0,1],values,marker='o',color=color,lw=2.5 if seed==4 else 1.4,label=f'Seed {seed}')
    pair.set_xticks([0,1],['Baseline','Force penalty']);pair.set_xlim(-.12,1.12);pair.set_ylim(0,105)
    pair.set_title('PhysX: the gain is mostly seed 4',loc='left',fontweight='bold',pad=14)
    pair.grid(axis='y',alpha=.18);pair.set_axisbelow(True)
    pair.legend(loc='lower right',frameon=False,ncol=2,fontsize=9)
    fig.text(.5,.015,'Same 50 layouts per seed. Earlier eight cases excluded. Lines pair training-seed IDs, not identical trained policies.',ha='center',fontsize=9,color='#555')
    fig.tight_layout(rect=[0,.055,1,1])
    for suffix in ['png','svg']:
        fig.savefig(directory/f'seed-results.{suffix}',dpi=170,facecolor='#faf8f3')
    plt.close(fig)


if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('directory',type=Path)
    plot(parser.parse_args().directory)