"""Publication figures from saved observations only. No inference or LLM judge."""
import argparse
import csv
import json
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

from analysis import interval

FAMILIES=['openai','anthropic','google']
LABELS={'openai':'Astra / mini','anthropic':'Fable / Haiku','google':'Pro / Flash'}
CONDITIONS=['hh','ll','hl','hl_aid_l','hl_blind_l']
TITLES=['H + H','L + L','H + L','H + L\nL: calculation table','H + L\nL: direct feed omitted']
COLORS=['#275d84','#dd7b32','#426a52','#8264a5','#a14b65']


def saved(path):return json.loads(Path(path).read_text(encoding='utf-8'))
def finish(fig,out,name):
 for ext in ['png','svg']:fig.savefig(Path(out)/(name+'.'+ext),dpi=180,bbox_inches='tight',facecolor='white')
 plt.close(fig)


def trajectories(run,out,name):
 fig,axes=plt.subplots(3,5,figsize=(14,8),sharex=True,sharey=True)
 for row,family in enumerate(FAMILIES):
  for col,condition in enumerate(CONDITIONS):
   ax=axes[row,col];ax.axhspan(145,150,color='#bdc7d2',alpha=.4,zorder=0);ax.axhline(190,color='#a9a9b3',ls=':',lw=1)
   episodes=[]
   for p in sorted((Path(run)/'episodes').glob('*/status.json')):
    s=saved(p);j=s['job']
    if j['study']=='market' and j['family']==family and j['condition']==condition:episodes.append((p.parent,s))
   completed=0
   for d,s in episodes:
    history=saved(d/'market-history.json') if (d/'market-history.json').exists() else []
    complete=s['status']=='completed' and s.get('metrics',{}).get('valid') and len(history)==s['job']['rounds']
    completed+=bool(complete)
    if history:
     prices=[sum(h['prices'])/2 for h in history]
     ax.plot(range(1,len(prices)+1),prices,color=COLORS[col],alpha=.8 if len(episodes)==1 else .45,lw=1.7 if complete else 1.2,ls='-' if complete else '--')
     if not complete:ax.scatter([len(prices)],[prices[-1]],marker='x',color='#912a42',s=35,zorder=5)
   if not any((d/'market-history.json').exists() for d,s in episodes):ax.text(.5,.5,'No executed\nprice pair',ha='center',va='center',transform=ax.transAxes,fontsize=9,color='#676775')
   ax.text(.03,.04,f'Complete: {completed}/{len(episodes)}',transform=ax.transAxes,fontsize=8,color='#555560')
   if row==0:ax.set_title(TITLES[col],fontsize=11,pad=14)
   if col==0:ax.set_ylabel(LABELS[family]+'\nMean executed price',fontsize=10)
   if row==2:ax.set_xlabel('Period')
   ax.set_ylim(100,245);ax.set_xlim(.8,12.3);ax.set_xticks([1,4,8,12]);ax.grid(axis='y',alpha=.16)
   ax.spines[['top','right']].set_visible(False)
 fig.suptitle(('Pilot' if name.startswith('pilot') else 'Independent replication')+' | every observed market trajectory',fontsize=18,x=.065,ha='left',y=1.025)
 fig.text(.065,-.03,'Solid: complete episode. Dashed / x: invalid or interrupted prefix. Gray band: two competitive grid equilibria (145, 150). Dotted: joint-optimal common price (190).\nH / L identify configurations, not established intelligence ranks. Price paths alone do not establish collusion.',fontsize=9,color='#555560')
 fig.tight_layout(h_pad=2,w_pad=1)
 finish(fig,out,name)


def market_comparison(analysis_dir,out):
 rows=list(csv.DictReader((Path(analysis_dir)/'market-episodes.csv').open(encoding='utf-8')))
 fig,axes=plt.subplots(1,3,figsize=(12.5,5),sharey=True)
 rng=np.random.default_rng(82041)
 for ax,family in zip(axes,FAMILIES):
  ax.axhspan(145,150,color='#bdc7d2',alpha=.5);ax.axhline(190,color='#a9a9b3',ls=':',lw=1)
  for i,condition in enumerate(CONDITIONS):
   subset=[r for r in rows if r['family']==family and r['condition']==condition]
   values=[float(r['tail_mean_price']) for r in subset if r['complete_valid']=='True']
   if values:
    ax.scatter(i+rng.uniform(-.16,.16,len(values)),values,s=31,alpha=.7,color=COLORS[i],edgecolor='white',linewidth=.4)
    stats=interval(values)
    ax.plot([i-.23,i+.23],[stats['mean']]*2,c='#151523',lw=2)
    if stats['ci95']:
     lo,hi=stats['ci95'];ax.plot([i,i],[lo,hi],c='#151523',lw=1.4)
   ax.text(i,106,f'{len(values)}/{len(subset)}',ha='center',fontsize=9,color='#535364')
  ax.set_title(LABELS[family],loc='left',fontsize=13,pad=12);ax.set_xticks(range(5),['HH','LL','HL','HL+aid','HL-feed'],rotation=25,ha='right');ax.set_ylim(100,240);ax.set_xlim(-.45,4.45);ax.grid(axis='y',alpha=.15);ax.spines[['top','right']].set_visible(False)
 axes[0].set_ylabel('Mean price in periods 7–12')
 fig.suptitle('Independent replication | prices across complete episodes',x=.065,ha='left',fontsize=17,y=1.025)
 fig.text(.065,-.10,'Each dot is one fresh market episode, not one round. Black bars: mean and episode bootstrap 95% interval. Counts: complete / planned.\nIncomplete or invalid episodes are shown in the trajectory figure and logs; their missing outcomes are not imputed. H / L are profile labels.',fontsize=9,color='#555560')
 fig.tight_layout(w_pad=2)
 finish(fig,out,'replication-comparison')


def primary_effects(analysis_dir,out):
 data=saved(Path(analysis_dir)/'primary-contrasts.json')['primary_tests']
 fig,ax=plt.subplots(figsize=(9,3.8));ax.axvline(0,color='#9494a5',lw=1);rng=np.random.default_rng(640)
 for i,row in enumerate(data):
  values=[p['difference'] for p in row['pairs'] if p['difference'] is not None]
  ax.scatter(values,i+rng.uniform(-.09,.09,len(values)),s=40,color=COLORS[i*2],alpha=.65)
  if row['mean'] is not None:
   ax.scatter([row['mean']],[i],marker='D',color='#191923',s=47,zorder=4)
   if row['ci95']:ax.plot(row['ci95'],[i,i],color='#191923',lw=2)
  ptext=f'Holm p={row["holm_adjusted_p"]:.3g}' if row['two_sided_sign_flip_p'] is not None else 'test unavailable'
  ax.text(.99,i+.18,f'n={row["n"]}/{row["planned_pairs"]}; {ptext}',transform=ax.get_yaxis_transform(),ha='right',fontsize=10)
 ax.set_yticks(range(2),[LABELS[r['family']] for r in data]);ax.set_ylim(-.45,1.5);ax.set_xlabel('Tail price difference: direct feed omitted − full feed');ax.grid(axis='x',alpha=.16);ax.spines[['top','right']].set_visible(False)
 fig.suptitle('The two prospectively specified comparisons',x=.04,ha='left',fontsize=16)
 fig.text(.04,-.1,'Dots: matched episode blocks. Diamond: mean. Line: 95% paired-block bootstrap interval.\nThis intervention removes a direct display; own quantities and private messages can still reveal rival prices.',fontsize=9,color='#555560')
 fig.tight_layout();finish(fig,out,'primary-effects')


def calibration(analysis_dir,out):
 data=saved(Path(analysis_dir)/'summary.json')['calibration'];keys=['astra','mini','fable','haiku','pro','flash'];labels=['Astra','mini','Fable','Haiku','Pro','Flash']
 fig,axes=plt.subplots(1,2,figsize=(10,3.9))
 for ax in axes:ax.spines[['top','right']].set_visible(False);ax.grid(axis='y',alpha=.15);ax.set_xticks(range(6),labels)
 counts=[data[k]['exact_optimal_items'] for k in keys];axes[0].bar(range(6),counts,color='#426a52');axes[0].set_ylim(0,14);axes[0].set_title('Exactly optimal price choices / 12',loc='left',fontsize=12)
 for i,n in enumerate(counts):axes[0].text(i,n+.3,str(n),ha='center',fontsize=10)
 regrets=[data[k]['normalized_regret']['mean']*100 for k in keys];axes[1].bar(range(6),regrets,color='#8264a5');axes[1].set_title('Mean normalized regret (%)',loc='left',fontsize=12);axes[1].set_ylim(0,max(regrets)*1.4)
 for i,n in enumerate(regrets):axes[1].text(i,n+.08,f'{n:.3f}',ha='center',fontsize=9)
 fig.suptitle('Pilot calibration | task-specific arithmetic, with a ceiling',x=.04,ha='left',fontsize=16,y=1.04)
 fig.text(.04,-.08,'Two independent prompts per profile, six shared items per prompt. The 12 items are not 12 independent model invocations.\nThis calibration does not establish a general intelligence ranking or negotiation ability.',fontsize=9,color='#555560')
 fig.tight_layout(w_pad=3);finish(fig,out,'calibration')


def followup_comparison(run,analysis_dir,out):
 data=saved(Path(analysis_dir)/'followup-contrast.json')['descriptive_contrast']
 fig,axes=plt.subplots(1,3,figsize=(12.5,4.5));colors=['#426a52','#a14b65']
 for ax,condition,color in zip(axes[:2],['hl','hl_blind_l'],colors):
  subset=[];complete_count=0
  ax.axhspan(145,150,color='#bdc7d2',alpha=.4);ax.axhline(190,color='#a9a9b3',ls=':',lw=1)
  for path in sorted((Path(run)/'episodes').glob('*/status.json')):
   state=saved(path)
   if state['job']['condition']!=condition:continue
   history=saved(path.parent/'market-history.json') if (path.parent/'market-history.json').exists() else []
   complete=state['status']=='completed' and state.get('metrics',{}).get('valid') and len(history)==12
   complete_count+=bool(complete);subset.append(state)
   if history:
    prices=[sum(h['prices'])/2 for h in history];ax.plot(range(1,len(prices)+1),prices,color=color,alpha=.55,ls='-' if complete else '--')
    if not complete:ax.scatter([len(prices)],[prices[-1]],marker='x',color='#912a42',s=40)
  ax.set_title(('Full direct feed' if condition=='hl' else 'L direct feed omitted')+f' | {complete_count}/{len(subset)} complete',loc='left',fontsize=11)
  ax.set_xlabel('Period');ax.set_xlim(.8,12.3);ax.set_ylim(100,240);ax.set_xticks([1,4,8,12]);ax.spines[['top','right']].set_visible(False);ax.grid(axis='y',alpha=.15)
 axes[0].set_ylabel('Mean executed price')
 ax=axes[2];ax.axhline(0,c='#aaaab5',lw=1)
 values=[b['difference'] for b in data['pairs'] if b['difference'] is not None]
 ax.scatter([b['rep'] for b in data['pairs'] if b['difference'] is not None],values,c='#275d84',s=45)
 if data['mean'] is not None:
  ax.axhline(data['mean'],c='#191923',ls='--',lw=1.3)
  if data['ci95']:ax.axhspan(*data['ci95'],color='#275d84',alpha=.1)
 ax.set_title(f'Paired tail-price differences | n={data["n"]}/8',loc='left',fontsize=11);ax.set_ylabel('Omitted minus full feed');ax.set_xlabel('Matched block ID (0–7)');ax.set_xticks(range(8));ax.spines[['top','right']].set_visible(False);ax.grid(axis='y',alpha=.15)
 fig.suptitle('Finite runtime follow-up | Astra + mini',x=.065,ha='left',fontsize=17,y=1.04)
 fig.text(.065,-.12,'Left: every observed trajectory; dashed / x marks an incomplete prefix. Right: block differences, mean and 95% paired-bootstrap band.\nAdaptively selected after first-replication results and host diagnostics. Descriptive only; no new significance test or pooled estimate.',fontsize=9,color='#555560')
 fig.tight_layout(w_pad=2);finish(fig,out,'followup-comparison')


if __name__=='__main__':
 ap=argparse.ArgumentParser();ap.add_argument('--pilot',required=True);ap.add_argument('--pilot-analysis',default='analysis/pilot');ap.add_argument('--replication');ap.add_argument('--replication-analysis',default='analysis/replication');ap.add_argument('--followup');ap.add_argument('--followup-analysis',default='analysis/followup');ap.add_argument('--out',required=True);args=ap.parse_args();Path(args.out).mkdir(parents=True,exist_ok=True)
 plt.rcParams.update({'font.family':'DejaVu Sans','font.size':10,'svg.fonttype':'none'})
 trajectories(args.pilot,args.out,'pilot-trajectories');calibration(args.pilot_analysis,args.out)
 if args.replication:
  trajectories(args.replication,args.out,'replication-trajectories');market_comparison(args.replication_analysis,args.out);primary_effects(args.replication_analysis,args.out)
 if args.followup:followup_comparison(args.followup,args.followup_analysis,args.out)
