"""Render scientific pilot figures using saved data; never reopen model weights."""
import argparse
import csv
import json
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def main():
    p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);a=p.parse_args();root=a.root
    out=root/'figures';out.mkdir(exist_ok=True)
    plt.rcParams.update({'font.size':10,'savefig.bbox':'tight','pdf.fonttype':42,'svg.fonttype':'none'})
    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/'functional-status.json').read_text())
    if (root/'candidate-summary.json').exists():
        c=json.loads((root/'candidate-summary.json').read_text())
        fig,ax=plt.subplots(figsize=(10,4.5))
        ax.bar(np.arange(len(c)),[max(v['maximum_symmetric_relative_error'],1e-8) for v in c],color=['#238b45' if v['passed'] else '#9e9e9e' for v in c])
        ax.axhline(.05,color='#b2182b',linestyle='--',label='Relative L2 screen = 0.05; other gates also apply')
        ax.set_yscale('log');ax.set_xticks(np.arange(len(c)),[v['id'] for v in c],rotation=45,ha='right')
        ax.set_ylabel('Worst calibration directional relative L2')
        ax.set_title('Fixed candidates; same-input function comparison\n4 calibration texts, 32 tokens; no training')
        ax.legend(loc='upper left',bbox_to_anchor=(0,-.38),fontsize=8);fig.tight_layout();save(fig,'candidate_function_errors')
    if (root/'interventions.json').exists():
        interventions=json.loads((root/'interventions.json').read_text())
        fig,ax=plt.subplots(figsize=(8,4.5))
        for idx,r in enumerate(interventions):
            values=r['sequence_nll_deltas'];x=np.arange(len(values))
            ax.plot(x,values,'o-',label=r['direction'])
        ax.axhline(0,color='gray',linewidth=.7);ax.axhline(.01,color='#b2182b',linestyle='--',label='Mean-delta budget (not per-sequence bound)')
        ax.set_xticks(np.arange(4),['heldout zh1','heldout zh2','heldout en1','heldout en2'])
        ax.set_ylabel('Token NLL change (nat/token)');ax.set_title('Single-site interventions; frozen candidate\n4 heldout texts, 32 tokens; limited pilot evidence')
        ax.legend(loc='upper left',bbox_to_anchor=(0,-.18),fontsize=8);fig.tight_layout();save(fig,'heldout_intervention_nll')
    norm_path=root/'norm-repetition-summary.json'
    repetition=json.loads(norm_path.read_text()) if norm_path.exists() else {}
    if (root/'norm-repetition-function-metrics.csv').exists():
        with (root/'norm-repetition-function-metrics.csv').open() as f:rows=list(csv.DictReader(f))
        sample_ids=list(dict.fromkeys(row['sample_id'] for row in rows))
        fig,ax=plt.subplots(figsize=(9,4.8))
        for source in sorted({row['source_layer'] for row in rows},key=int):
            rr=[r for r in rows if r['source_layer']==source]
            ax.plot(range(len(rr)),[max(float(r['relative_l2']),float(r['reverse_relative_l2'])) for r in rr],'o-',label=f'Input captured at layer {source}')
        ax.axhline(.05,color='#b2182b',ls='--',label='Heldout empirical screen = 0.05')
        chosen=repetition.get('selected')
        if chosen:ax.axhline(chosen['worst_uniform_bound'],color='#6a51a3',ls=':',label='Proven uniform all-input bound')
        ax.axvline(3.5,color='gray',linewidth=.8)
        ax.set_xticks(range(len(sample_ids)),sample_ids,rotation=30,ha='right');ax.set_ylabel('Worst of two directional relative L2 errors')
        ax.set_title('Learned input RMSNorm at layer indices 12/13: same-input comparison\nCalibration (left) / independent heldout texts (right), 32 tokens each')
        ax.legend(loc='upper left',bbox_to_anchor=(0,-.25),fontsize=8);fig.tight_layout();save(fig,'learned_norm_function_errors')
    if repetition.get('interventions'):
        fig,ax=plt.subplots(figsize=(8,4.5))
        for v in repetition['interventions']:
            ax.plot(range(4),v['sequence_nll_deltas'],'o-',label=f'Layer {v["target"]} uses {v["donor"]}; mean={v["mean_nll_delta"]:.5g}')
        ax.axhline(0,color='gray',linewidth=.8);ax.set_xticks(range(4),['zh1','zh2','en1','en2'])
        ax.set_xlabel('Independent heldout text (32 tokens)');ax.set_ylabel('Mean token NLL change (nat/token)')
        ax.set_title('Bidirectional single-site substitution of existing learned RMSNorm modules')
        ax.legend(loc='upper left',bbox_to_anchor=(0,-.18),fontsize=8);fig.tight_layout();save(fig,'learned_norm_network_substitution')
    if (root/'gain-factorization-score-metrics.csv').exists():
        with (root/'gain-factorization-score-metrics.csv').open() as f:rr=list(csv.DictReader(f))
        layers=sorted({int(v['layer']) for v in rr});fig,axes=plt.subplots(1,2,figsize=(10,4))
        for ax,key,label in zip(axes,['score_max_abs_error','probability_max_abs_error'],['Maximum absolute score error','Maximum absolute probability error']):
            for direction in ('q_to_k','k_to_q'):
                ax.plot(layers,[max(float(v[key]) for v in rr if int(v['layer'])==layer and v['direction']==direction) for layer in layers],'o-',label=direction)
            ax.set_xlabel('Full-attention layer index');ax.set_xticks(layers);ax.set_ylabel(label);ax.ticklabel_format(axis='y',style='sci',scilimits=(0,0));ax.legend(loc='upper left',fontsize=8)
        fig.suptitle('Supplement: algebraic reparameterization, separate from learned-function repetition');fig.tight_layout();save(fig,'algebraic_gain_identity')
    neuron_path=root/'neuron-reuse-summary.json'
    neuron=json.loads(neuron_path.read_text()) if neuron_path.exists() else {}
    text=['# CPU learned-function repetition pilot','',
          'Primary question: do two distinct trained parts implement similar functions? Cost is not a success condition.',
          f'- Existing learned RMSNorm function positive: {repetition.get("learned_function_repetition_positive",False)}',
          f'- Bidirectional network quality screen passed: {repetition.get("bidirectional_network_quality_pass",False)}',
          f'- Learned scalar MLP neuron search positive: {neuron.get("empirical_repeated_function_positive",False)}',
          f'- Original 11 candidate run state: {status["state"]}; selected: {status.get("selected_candidate",[])}',
          '- All model weights are retained on the server until the user explicitly requests deletion.','']
    if repetition.get('selected'):
        c=repetition['selected'];text += ['## Existing learned module evidence','',
          f'Role: {c["role"]}; layers {c["a"]} and {c["b"]}; each maps 2048-dimensional input to output.',
          f'Uniform real-arithmetic relative-error bound, first-to-second: {c["a_to_b_uniform_bound"]:.12g}; reverse: {c["b_to_a_uniform_bound"]:.12g}.',
          f'Maximum heldout directional output relative L2: {repetition.get("worst_heldout_relative_l2")}; minimum cosine: {repetition.get("worst_heldout_cosine")}.',
          'Same epsilon and nonzero gains: the common RMS denominator cancels. For any nonzero input x,',
          '`||f_a(x)-f_b(x)|| / ||f_a(x)|| <= max_j |a_j-b_j|/|a_j|`; swap a/b for the reverse direction.',
          'Selection used only the tight all-input bound across 552 same-role pairs, before inspecting inputs.',
          'Real probes use four calibration and four heldout experiment-authored texts, 32 tokens each, and both source layers.',
          'These are learned, nonconstant gain vectors. The claim concerns RMSNorm modules; it does not establish attention/MLP-body duplication.','']
        for v in repetition.get('interventions',[]):text += [f'- Layer {v["target"]} uses {v["donor"]}: mean heldout NLL delta {v["mean_nll_delta"]:.12g}; sequence P95 delta {v["sequence_p95_nll_delta"]:.12g}; screen pass={v["passed_quality"]}.']
    text += ['', '## Negative results and separate algebraic result','',
             'The original whole-attention, whole-MLP and within-layer Q/K norm candidates did not pass their gates.',
             'The scalar neuron search excludes low-RMS/near-constant profiles and selects at most one pair without heldout feedback; consult its search summary for coverage.',
             'Q/K non-RoPE gain consolidation is a separately verified real-arithmetic identity, not evidence that two existing learned modules duplicate functions.',
             'Earlier local CPU timing records are preserved for provenance; performance is outside the current success criterion.','',
             '## Limits and provenance','',
             'Original BF16 weights were decoded into streaming FP32 CPU operators; no training, scale fitting, R fitting or GPU.',
             'The numerical reference includes 22 independent official Torch comparisons. Early adapter errors and two resource-guard stops are preserved as historical logs.',
             'The persistent computation ledger includes a conservative 180-second reserve for early checks; no budget reset after that reserve.',
             'The text sample is tiny and short. Network NLL checks are not broad task-accuracy validation.','']
    (root/'FUNCTIONAL_REPORT.md').write_text('\n'.join(text))


if __name__=='__main__':main()