"""Scientific plots from saved tables/probabilities only; no model access."""
import argparse
import csv
import json
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from .protocol import LAYERS,PAIRS


def read_csv(path):
    with path.open() as f:return list(csv.DictReader(f))


def main():
    p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);args=p.parse_args();root=args.root;out=root/'figures';out.mkdir(exist_ok=True)
    plt.rcParams.update({'font.size':9,'pdf.fonttype':42,'svg.fonttype':'none','savefig.bbox':'tight'})
    def save(fig,name):
        for ext in ('png','pdf','svg'):fig.savefig(out/f'{name}.{ext}',dpi=180)
        plt.close(fig)
    status=json.loads((root/'status.json').read_text());summary=json.loads((root/'calibration-summary.json').read_text())
    matrices=[]
    for mode in ('qk_head','identity','matched'):
        matrix=np.full((6,6),np.nan)
        for i,j in PAIRS:
            rows=[r for r in summary if r['i']==i and r['j']==j and ((mode=='qk_head' and r['kind']=='qk_head') or (mode!='qk_head' and r['id']==f'qk_map_{i}_{j}_{mode}'))]
            value=min(r['worst_mean_js'] for r in rows)
            a,b=LAYERS.index(i),LAYERS.index(j);matrix[a,b]=matrix[b,a]=value
        matrices.append(matrix)
    fig,axes=plt.subplots(1,3,figsize=(13,4.7));vmax=max(np.nanmax(m) for m in matrices)
    for ax,m,title in zip(axes,matrices,['Best single head pair','All heads: original order','All heads: calibration bijection']):
        im=ax.imshow(m,cmap='YlOrRd',vmin=0,vmax=vmax)
        for a in range(6):
            ax.add_patch(Rectangle((a-.5,a-.5),1,1,facecolor='#eeeeee',edgecolor='#aaaaaa',hatch='///',lw=.3))
            for b in range(6):
                if a!=b:ax.text(b,a,f'{m[a,b]:.3f}',ha='center',va='center',fontsize=7,color='black' if m[a,b]<vmax*.65 else 'white')
        ax.set_xticks(range(6),LAYERS);ax.set_yticks(range(6),LAYERS);ax.set_xlabel('Layer index');ax.set_ylabel('Layer index');ax.set_title(title)
    fig.suptitle('QK retrieval divergence: worst-text mean JS over both source input sets\nCalibration only; single-head panel minimizes over 64 cross-layer head pairs; hatched diagonal unused')
    fig.subplots_adjust(top=.77,bottom=.18,right=.9,wspace=.35);cax=fig.add_axes([.92,.2,.012,.5]);fig.colorbar(im,cax=cax,label='JS (natural log); lower is closer');save(fig,'cross_layer_qk_divergence')
    order=['full_attention','qkv_head','qk_map','qk_head','v_group'];labels=['Whole attention','QKV+gate/O head','QK map transplant','QK head transplant','V group transplant']
    fig,ax=plt.subplots(figsize=(10,4.9));rng=np.random.default_rng(0)
    for idx,(kind,label) in enumerate(zip(order,labels)):
        rows=[r for r in summary if r['kind']==kind and r['output_verified']]
        if rows:
            x=idx+rng.uniform(-.15,.15,len(rows));y=[r['worst_output_error'] for r in rows]
            ax.scatter(x,y,s=12,alpha=.35,color='#606060')
            ax.text(idx,max(y)*1.1,f'n={len(rows)}',ha='center',fontsize=8)
        else:ax.text(idx,.3,'not output-probed\ninitial gate failed',ha='center',va='center',fontsize=8)
    ax.axhline(.05,color='#b2182b',ls='--',label='Strict output error = 0.05')
    ax.axhline(.2,color='#d95f02',ls=':',label='Approximate output error = 0.20')
    ax.set_yscale('log');ax.set_xlim(-.5,4.5);ax.set_xticks(range(5),labels,rotation=20,ha='right');ax.set_ylabel('Worst directional relative L2, queries 7–31')
    ax.set_title('Primary protocol: calibrated cross-layer output comparisons\nEach point is one unit; no candidate passed all predeclared gates')
    ax.legend(loc='upper left',bbox_to_anchor=(0,-.28),fontsize=8);fig.tight_layout();save(fig,'cross_layer_output_errors')
    values=read_csv(root/'calibration-value-metrics.csv');unit={}
    for r in values:unit.setdefault(r['candidate'],[]).append(r)
    xx=[np.mean([float(r['centered_value_gram_cka']) for r in rows]) for rows in unit.values()]
    yy=[max(max(float(r['error_a']),float(r['error_b'])) for r in rows) for rows in unit.values()]
    controls=read_csv(root/'controls.csv');random_cka=[float(r['gram_cka']) for r in controls if r['control']=='value_gram_random_projection_pair']
    fig,ax=plt.subplots(figsize=(7,4.5));ax.scatter(xx,yy,s=24,color='#377eb8',label='60 cross-layer KV-head pairs')
    ax.axvspan(min(random_cka),max(random_cka),color='gray',alpha=.15,label='Random-projection Gram CKA range')
    ax.axvline(np.median(random_cka),color='gray',ls=':')
    ax.axhline(.2,color='#b2182b',ls='--',label='Approximate direct V error = 0.20')
    ax.set_xlabel('Mean centered token-Gram CKA');ax.set_ylabel('Worst direct V relative L2');ax.set_title('V geometry is a diagnostic, not substitution evidence')
    ax.legend(loc='upper left',bbox_to_anchor=(0,-.19),fontsize=8);fig.tight_layout();save(fig,'value_geometry_and_direct_error')
    conditional_path=root/'conditional-status.json'
    cond=json.loads(conditional_path.read_text()) if conditional_path.exists() else {}
    if (root/'conditional-function-metrics.csv').exists():
        rows=read_csv(root/'conditional-function-metrics.csv');labels=[r['sample'] for r in rows];x=np.arange(len(rows))
        fig,axes=plt.subplots(1,2,figsize=(13,5))
        axes[0].plot(x,[float(r['top5']) for r in rows],'o-',label='Unweighted top-5 overlap')
        axes[0].plot(x,[1-float(r['total_variation']) for r in rows],'o-',label='Probability-mass overlap (1−TV)')
        axes[0].set_ylim(0,1.03);axes[0].set_ylabel('Overlap');axes[0].legend(loc='lower left',fontsize=8)
        axes[1].plot(x,[max(float(r['output_error_a']),float(r['output_error_b'])) for r in rows],'o-',color='#238b45')
        axes[1].set_ylim(0,max(.05,max(max(float(r['output_error_a']),float(r['output_error_b'])) for r in rows)*1.15))
        axes[1].set_ylabel('Head readout relative L2');axes[1].set_title('Fixed layer23 V/gate/O',fontsize=9)
        for ax in axes:
            for pos in (3.5,7.5):ax.axvline(pos,color='gray',lw=.6)
            ax.set_xticks(x,labels,rotation=50,ha='right');ax.set_xlabel('Calibration / old validation / fresh tests, 4 each')
        fig.suptitle('Exploratory fixed case: layer3 head7 vs layer23 head6, both applied to H23\nFresh JS / tail-mass limits not fully passed; network substitution not run')
        fig.tight_layout();save(fig,'conditional_probability_mass_and_readout')
    if (root/'conditional-fresh-probabilities.npz').exists():
        p=np.load(root/'conditional-fresh-probabilities.npz');a=p['target_layer23_head6'][0,0];b=p['donor_layer3_head7_on_H23'][0,0];d=p['donor_layer3_head7_on_H3'][0,0]
        fig,axes=plt.subplots(2,2,figsize=(9,8));images=[a,b,np.abs(a-b),d]
        titles=['Layer23 head6 on H23','Layer3 head7 on the same H23','Absolute same-input probability difference','Layer3 head7 on its actual H3 (different input)']
        for idx,(ax,z,title) in enumerate(zip(axes.flat,images,titles)):
            im=ax.imshow(z,cmap='magma',vmin=0,vmax=(max(float(z.max()),1e-8) if idx==2 else 1.));fig.colorbar(im,ax=ax,fraction=.045,pad=.03)
            ax.set_xticks([0,7,15,23,31]);ax.set_yticks([0,7,15,23,31]);ax.set_xlabel('Key position');ax.set_ylabel('Query position');ax.set_title(title,fontsize=9)
        fig.suptitle('One fresh Chinese text: cross-layer conditional self-token retrieval\nRows 0–6 shown for context; primary metrics use queries 7–31')
        fig.tight_layout();save(fig,'conditional_cross_layer_probability_maps')
    text=['# Cross-layer QKV CPU study','',
          'Primary protocol and exploratory addendum are reported separately; no post-hoc success relabeling.',
          f'- Primary state: {status["state"]}; selected candidate: {status.get("selected",[])}.',
          f'- Candidate counts: {status.get("candidate_counts")}.',
          f'- Passing counts (approximate tier, all gates): {status.get("passing_counts")}.',
          f'- Primary computation seconds: {status.get("elapsed_seconds")}; peak RSS KiB: {status.get("peak_rss_kib")}.','',
          '## What was measured','',
          'Six full-attention layers, indices 3/7/11/15/19/23; all 15 cross-layer pairs. No within-layer Q/K/V interchange.',
          'Same normalized hidden inputs from both source layers; frozen Q/K norm, partial RoPE, GQA, V, output gate and O.',
          'QK head probabilities, whole-map injection in original/frozen matched head order, direct V functions, complete head contributions in residual coordinates, and complete branches were separated.',
          'Primary metrics use queries 7–31. Discrete matching uses calibration only; no continuous alignment, scale fitting or training.',
          'No primary candidate passed, so primary validation/fresh/network intervention stages were not entered.','']
    for kind in order:
        group=[r for r in summary if r['kind']==kind];outrows=[r for r in group if r['worst_output_error'] is not None];p=[r for r in group if r['worst_mean_js'] is not None];v=[r for r in group if r['worst_value_error'] is not None]
        text += [f'- {kind}: n={len(group)}, best worst output error={min((r["worst_output_error"] for r in outrows),default=None)}, minimum worst-text mean JS={min((r["worst_mean_js"] for r in p),default=None)}, best worst direct V error={min((r["worst_value_error"] for r in v),default=None)}.']
    if cond:
        text+=['','## Exploratory metric audit: frozen one-source candidate','',
               'After inspecting primary calibration results, one candidate was frozen: layer3 head7 vs layer23 head6 on H23.',
               'This addendum audited the instability of unweighted top-k tail ranking; original primary results were not changed.',
               'Four old validation and four previously unused fresh texts were evaluated with frozen extra gates and no candidate fallback.',
               f'- Addendum fresh function pass: {cond.get("fresh_function_pass")}.',
               f'- Overall addendum state: {cond.get("state")}; primary unchanged/negative: {cond.get("primary_protocol_unchanged_and_negative")}.',
               'Fresh probability-mass and fixed V/gate/O readout measurements show conditional behavioral overlap; strict fresh mean-JS, row-JS-P95 and row-TV-P95 limits still failed.',
               'No full-network substitution was run after that failure. V similarity, bidirectionality and complete QKV-head equivalence are not established.',
               'The actual-H3 cache probability is a separate different-input diagnostic, not the same-input function comparison.']
    text+=['','## Limits and retention','',
           'This is a short text-only pilot, not proof of absence of repetition elsewhere. Exact arithmetic, matrix coordinates, retrieval behavior and final residual output are distinct questions.',
           'Centered value Gram CKA can be high even for independent Gaussian projections of common inputs; it is not substitution evidence.',
           'All weights and large profile caches remain on the authorized server. Git delivery includes all plot tables, probabilities, protocols and cache location/hash metadata.',
           'No GPU, new checkpoint download, model training, continuous R fitting or automatic weight/cache deletion.']
    (root/'REPORT.md').write_text('\n'.join(text)+'\n')


if __name__=='__main__':main()