Appendix: YAML Serializer Variance and Token Economics

import json
import os
import re
import subprocess
import tempfile
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tiktoken
import yaml
from IPython.display import display

enc = tiktoken.get_encoding('cl100k_base')
CONSULTANTS_DIR = Path('../3_prompts/consultants')

# All production prompts above 1 KB (exclude stubs)
PROMPT_FILES = sorted(
    [p for p in CONSULTANTS_DIR.glob('*.json') if p.stat().st_size > 1000],
    key=lambda p: p.stat().st_size,
)

print(f"{'File':<45} {'Size (bytes)':>12}")
print('-' * 59)
for p in PROMPT_FILES:
    print(f"{p.name:<45} {p.stat().st_size:>12,}")

# ── Serializer helpers ────────────────────────────────────────────────────────

def yaml_standard(data, width=80):
    """PyYAML standard YAML. width=80 is the PyYAML default (enables line-wrapping)."""
    return yaml.dump(data, allow_unicode=True, default_flow_style=False,
                     sort_keys=False, width=width)


def yaml_literal(data):
    """PyYAML YAML with literal block scalars (`|`) for all string values."""
    class LiteralStr(str):
        pass

    class LiteralDumper(yaml.Dumper):
        pass

    LiteralDumper.add_representer(
        LiteralStr,
        lambda dumper, val: dumper.represent_scalar('tag:yaml.org,2002:str', val, style='|'),
    )

    def apply_literal(obj):
        if isinstance(obj, str) and '\n' in obj:
            return LiteralStr(obj)
        if isinstance(obj, str):
            return LiteralStr(obj + '\n')
        if isinstance(obj, dict):
            return {k: apply_literal(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [apply_literal(i) for i in obj]
        return obj

    return yaml.dump(apply_literal(data), Dumper=LiteralDumper,
                     allow_unicode=True, default_flow_style=False, sort_keys=False)


def yq_standard(path):
    return subprocess.run(['yq', '-oy', str(path)], capture_output=True, text=True).stdout


def yq_literal(path):
    return subprocess.run(['yq', '-oy', '..style="literal"', str(path)],
                          capture_output=True, text=True).stdout


def tokens(text):
    return len(enc.encode(text))


JSON_VARIANTS = [
    'Original JSON',
    'Minified JSON',
    'Pretty JSON',
]

YAML_VARIANTS = [
    'Standard YAML (PyYAML)',
    'Standard YAML (yq)',
    'YAML Literal (PyYAML)',
    'YAML Literal (yq)',
]

VARIANT_NAMES = JSON_VARIANTS + YAML_VARIANTS


def variants_for(path, data):
    return {
        'Original JSON':          path.read_text(),
        'Minified JSON':          json.dumps(data, ensure_ascii=False, separators=(',', ':')),
        'Pretty JSON':            json.dumps(data, ensure_ascii=False, indent=2),
        'Standard YAML (PyYAML)': yaml_standard(data, width=80),
        'Standard YAML (yq)':     yq_standard(path),
        'YAML Literal (PyYAML)':  yaml_literal(data),
        'YAML Literal (yq)':      yq_literal(path),
    }

# Build long-format DataFrame
records = []
for path in PROMPT_FILES:
    data = json.loads(path.read_text())
    vmap = variants_for(path, data)
    for variant, text in vmap.items():
        records.append({
            'file':    path.name,
            'group':   'JSON' if variant in JSON_VARIANTS else 'YAML',
            'variant': variant,
            'tokens':  tokens(text),
        })

df = pd.DataFrame(records)
df['variant'] = pd.Categorical(df['variant'], categories=VARIANT_NAMES, ordered=True)

# Wide format: files × variants — used by all downstream cells
df_wide = df.pivot(index='file', columns='variant', values='tokens')[VARIANT_NAMES]
df_wide.columns.name = None
df_wide.index.name   = 'File'

# ── Display: one table, formats as columns, files as rows ─────────────────────
# MultiIndex columns group JSON and YAML variants visually
col_idx = pd.MultiIndex.from_tuples(
    [('JSON' if v in JSON_VARIANTS else 'YAML', v) for v in VARIANT_NAMES],
    names=['Group', 'Variant'],
)

# Token counts — with median summary row
counts_display = df_wide.copy()
counts_display.loc['median'] = df_wide.median().astype(int)
counts_display.columns = col_idx
display(counts_display.style
        .format('{:,.0f}')
        .bar(color='#d4e6f1', vmin=0, axis=None)
        .set_caption('Token counts — all files × all formats'))

# Δ % vs Minified JSON — with median summary row
baseline      = df_wide['Minified JSON']
delta_display = df_wide.subtract(baseline, axis=0).div(baseline, axis=0).mul(100).round(1)
delta_display.loc['median'] = delta_display.median().round(1)
delta_display.columns = col_idx
display(delta_display.style
        .format(lambda x: 'baseline' if abs(x) < 0.01 else f'+{x:.1f}%')
        .background_gradient(cmap='YlOrRd', vmin=0, axis=None)
        .set_caption('Δ % vs Minified JSON'))

# Median delta (%) across all files — grouped by format family
baseline_col = df_wide['Minified JSON']
delta_pct = df_wide.subtract(baseline_col, axis=0).div(baseline_col, axis=0).mul(100)

summary = pd.DataFrame({
    'Group':    ['JSON' if v in JSON_VARIANTS else 'YAML' for v in VARIANT_NAMES],
    'Variant':  VARIANT_NAMES,
    'Median Δ %': delta_pct.median().round(1).values,
    'Min Δ %':    delta_pct.min().round(1).values,
    'Max Δ %':    delta_pct.max().round(1).values,
}).set_index(['Group', 'Variant'])

# Expose as Series for the bar chart below
median_deltas = delta_pct.median()

display(summary)

# Bar chart: JSON group (blue) | YAML group (orange/pink); Minified JSON excluded (baseline = 0%)
plot_json = [v for v in JSON_VARIANTS if v != 'Minified JSON']
plot_yaml = YAML_VARIANTS
all_plot  = plot_json + plot_yaml
values    = [median_deltas[v] for v in all_plot]

colors  = (['#8da0cb'] * len(plot_json) +
           ['#fc8d62', '#fc8d62', '#e78ac3', '#e78ac3'])
hatches = ([''] * len(plot_json) +
           ['', '///', '', '///'])

x   = np.arange(len(all_plot))
fig, ax = plt.subplots(figsize=(13, 5))
bars = ax.bar(x, values, color=colors, edgecolor='grey', linewidth=0.8)
for bar, hatch, val in zip(bars, hatches, values):
    bar.set_hatch(hatch)
    ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.2,
            f'+{val:.1f}%', ha='center', va='bottom', fontsize=9)

split = len(plot_json) - 0.5
ax.axvline(split, color='grey', linestyle='--', linewidth=0.8, alpha=0.6)
ax.text(split - 0.05, max(values) * 1.1, 'JSON', ha='right', fontsize=9, color='#555')
ax.text(split + 0.05, max(values) * 1.1, 'YAML', ha='left',  fontsize=9, color='#555')

ax.set_xticks(x)
ax.set_xticklabels(all_plot, rotation=20, ha='right')
ax.set_ylabel('Median Δ % vs Minified JSON')
ax.set_title('Format Token Cost by Serializer — Median across prompt files\n'
             '(striped = yq  |  blue = JSON, orange = Std YAML, pink = Literal YAML)')
ax.set_ylim(0, max(values) * 1.3)
ax.grid(axis='y', linestyle='--', alpha=0.3)
plt.tight_layout()
plt.show()

rep_path = PROMPT_FILES[-1]
rep_data = json.loads(rep_path.read_text())
print(f"Representative file: {rep_path.name} ({rep_path.stat().st_size:,} bytes)")

t_min  = tokens(json.dumps(rep_data, ensure_ascii=False, separators=(',', ':')))
t_pretty = tokens(json.dumps(rep_data, ensure_ascii=False, indent=2))
w80    = yaml_standard(rep_data, width=80)
w_inf  = yaml_standard(rep_data, width=float('inf'))
t80, t_inf = tokens(w80), tokens(w_inf)

print(f"Minified JSON:\t\t\t{t_min:>6} tokens  (baseline)")
print(f"Standard YAML width=80:\t\t{t80:>6} tokens  ({len(w80.splitlines())} lines)  "
      f"[{(t80-t_min)/t_min*100:+.1f}% vs min JSON]")
print(f"Standard YAML width=∞:\t\t{t_inf:>6} tokens  ({len(w_inf.splitlines())} lines)  "
      f"[{(t_inf-t_min)/t_min*100:+.1f}% vs min JSON]")
print(f"Pretty JSON:\t\t\t{t_pretty:>6} tokens  [{(t_pretty-t_min)/t_min*100:+.1f}% vs min JSON]")
print(f"\nWrapping overhead:\t\t{t80 - t_inf:>+6} tokens  ({(t80 - t_inf) / t_inf * 100:.1f}%)")
print(f"YAML w=∞ vs Pretty JSON:\t{t_inf - t_pretty:>+6} tokens  "
      f"({'cheaper' if t_inf < t_pretty else 'more expensive'})")

default_style = yaml_standard(rep_data, width=float('inf'))
double_quoted = yaml.dump(rep_data, allow_unicode=True, default_flow_style=False,
                          sort_keys=False, width=float('inf'), default_style='"')

t_def, t_dq = tokens(default_style), tokens(double_quoted)
print(f"Default style (w=∞):       {t_def:>6} tokens")
print(f"Double-quote style (w=∞):  {t_dq:>6} tokens")
print(f"Delta:                     {t_dq - t_def:>+6} tokens")

date_default = [l for l in default_style.splitlines() if re.search(r'202\d-', l)][:3]
date_dquoted = [l for l in double_quoted.splitlines() if re.search(r'202\d-', l)][:3]
print("\nDate quoting — default:")
for l in date_default:
    print(f"  {l}")
print("Date quoting — double-quote:")
for l in date_dquoted:
    print(f"  {l}")

t_pretty = tokens(json.dumps(rep_data, ensure_ascii=False, indent=2))
py_lit   = yaml_literal(rep_data)
yq_lit   = yq_literal(rep_path)
t_py_lit, t_yq_lit = tokens(py_lit), tokens(yq_lit)

print(f"Pretty JSON:           {t_pretty:>6} tokens  (reference)")
print(f"PyYAML literal (`|`):  {t_py_lit:>6} tokens  [{t_py_lit - t_pretty:+d} vs Pretty JSON]"
      f"  {'← more expensive' if t_py_lit > t_pretty else '← cheaper'}")
print(f"yq literal (`|-`):     {t_yq_lit:>6} tokens  [{t_yq_lit - t_pretty:+d} vs Pretty JSON]"
      f"  {'← more expensive' if t_yq_lit > t_pretty else '← cheaper'}")
print(f"\nPyYAML vs yq literal delta: {t_py_lit - t_yq_lit:+d} tokens  (negative = PyYAML cheaper)")

# --- Decompose the delta by progressively transforming py_lit toward yq_lit ---

# Step 0: baseline
t0 = tokens(py_lit)

# Step 1: change all `|` indicators to `|-` (string scalars only — same content, different indicator)
py_step1 = re.sub(r'(: \|)(\n)', r'\1-\2', py_lit)
t1 = tokens(py_step1)

# Step 2: the remaining gap is everything else — boolean coercion, extra lines, quoting, etc.
# We can't transform py_lit further without reproducing yq's full output,
# so we report the remainder as "structural differences beyond indicators"

py_pipe_count = len(re.findall(r': \|$',  py_lit, re.MULTILINE))
yq_pipe_count = len(re.findall(r': \|-$', yq_lit, re.MULTILINE))
coerced       = re.findall(r'\|-\n(\s+)(true|false|\d+\.?\d*)\n', yq_lit)
n_coerced     = len(coerced)

actual_delta    = t_yq_lit - t0
indicator_cost  = t1 - t0                          # Source 1: `|` → `|-` on string scalars
structural_cost = t_yq_lit - t1                    # Source 2: everything else

print(f"\nDelta breakdown (yq literal − PyYAML literal = {actual_delta:+d} tokens):")
print(f"\n  Source 1 — `|-` vs `|` indicator on {py_pipe_count} string scalars:")
print(f"    Measured on actual document (all `|` → `|-`): {indicator_cost:+d} tokens")
print(f"    Per-indicator cost: {indicator_cost / py_pipe_count:+.1f} tokens")
print(f"\n  Source 2 — structural differences beyond indicators:")
print(f"    yq `|-` blocks (total): {yq_pipe_count}  (PyYAML `|`: {py_pipe_count})")
print(f"    Non-string scalars coerced to `|-`: {n_coerced}  {[v for _, v in coerced[:6]]}")
print(f"    Extra lines in yq: {len(yq_lit.splitlines()) - len(py_lit.splitlines())}")
print(f"    Measured cost (actual yq − step-1 document): {structural_cost:+d} tokens")
print(f"\n  Verification: {indicator_cost:+d} + {structural_cost:+d} = "
      f"{indicator_cost + structural_cost:+d} tokens  (actual: {actual_delta:+d})")

def show_diff(label_a, text_a, label_b, text_b, max_diffs=8):
    lines_a, lines_b = text_a.splitlines(), text_b.splitlines()
    print(f"Line counts: {label_a}={len(lines_a)}, {label_b}={len(lines_b)}")
    diffs = 0
    for i, (a, b) in enumerate(zip(lines_a, lines_b)):
        if a != b:
            print(f"  [{i}] {label_a}:\t{a!r}")
            print(f"       {label_b}:\t{b!r}")
            diffs += 1
            if diffs >= max_diffs:
                print(f"  ... (stopped at {max_diffs})")
                break
    if diffs == 0:
        print("  No differences in overlapping lines.")

print("=== Standard YAML: PyYAML vs yq ===")
show_diff('PyYAML', yaml_standard(rep_data), 'yq', yq_standard(rep_path))
print("\n=== YAML Literal: PyYAML vs yq ===")
show_diff('PyYAML', yaml_literal(rep_data), 'yq', yq_literal(rep_path))

def find_diff(a, b, path='root'):
    if isinstance(a, dict) and isinstance(b, dict):
        for k in set(list(a.keys()) + list(b.keys())):
            find_diff(a.get(k), b.get(k), f"{path}.{k}")
    elif isinstance(a, list) and isinstance(b, list):
        for i, (x, y) in enumerate(zip(a, b)):
            find_diff(x, y, f"{path}[{i}]")
    elif a != b:
        print(f"  DIFF at {path}: {a!r} → {b!r}")

print(f"{'File':<45} {'PyYAML':>8} {'yq':>8}")
print('-' * 63)
for path in PROMPT_FILES:
    data = json.loads(path.read_text())
    # yaml.safe_load is safe here: yaml.dump quotes YAML 1.1 implicit types
    # (yes/no/on/off/true/false, date strings) so safe_load recovers them as strings.
    py_rt = yaml.safe_load(yaml_standard(data, width=float('inf')))
    py_ok = py_rt == data
    yq_yaml = yq_standard(path)
    yq_json = subprocess.run(['yq', '-oj', '-'], input=yq_yaml,
                             capture_output=True, text=True).stdout
    yq_rt = json.loads(yq_json)
    yq_ok = yq_rt == data
    print(f"{path.name:<45} {'✓' if py_ok else '✗':>8} {'✓' if yq_ok else '✗':>8}")
    if not py_ok: find_diff(data, py_rt)
    if not yq_ok: find_diff(data, yq_rt)

# Edge case scalars.
# YAML 1.1 implicit types: 'yes'/'no'/'on'/'off' are booleans in YAML 1.1 but not in
# YAML 1.2 or JSON. PyYAML (YAML 1.1) quotes these strings on dump, so safe_load
# recovers them as strings. yq (YAML 1.2 / Go) does not coerce them — both paths
# should round-trip correctly, but we test explicitly to confirm.
edge_cases = {
    'booleans':       {'flag_true': True, 'flag_false': False},
    'numbers':        {'integer': 42, 'float_val': 3.14},
    'null':           {'empty': None},
    'date string':    {'date': '2026-01-07'},
    'bool string':    {'truthy': 'true', 'falsy': 'false'},
    'yaml11 strings': {'yes_str': 'yes', 'no_str': 'no', 'on_str': 'on', 'off_str': 'off'},
}

print(f"{'Scalar type':<15} {'PyYAML w=∞':>12} {'yq':>8}")
print('-' * 37)
for label, test_data in edge_cases.items():
    py_rt = yaml.safe_load(yaml_standard(test_data, width=float('inf')))
    py_ok = py_rt == test_data
    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
        json.dump(test_data, f)
        tmp = f.name
    try:
        yq_yaml = yq_standard(tmp)
        yq_json = subprocess.run(['yq', '-oj', '-'], input=yq_yaml,
                                 capture_output=True, text=True).stdout
        yq_rt = json.loads(yq_json)
        yq_ok = yq_rt == test_data
    finally:
        os.unlink(tmp)
    print(f"{label:<15} {'✓' if py_ok else f'✗ {py_rt}':>12} {'✓' if yq_ok else f'✗ {yq_rt}':>8}")

# HuggingFace authentication — required once per environment for gated models.
# Steps:
#   1. Accept licence at https://huggingface.co/google/gemma-2-2b
#   2. Accept licence at https://huggingface.co/meta-llama/Llama-3.2-1B
#   3. Create a read token at https://huggingface.co/settings/tokens
#   4. Uncomment the line below and run this cell (token is cached to ~/.cache/huggingface/):
from huggingface_hub import login; login()

import transformers
transformers.logging.set_verbosity_error()
from transformers import AutoTokenizer

# Claude is excluded: no downloadable tokenizer for Claude 3+.
# Gemma 2 is the recommended Gemini proxy (see note above).
_candidates = [
    ('GPT-2 (50k)',       'openai-community/gpt2'),
    ('Llama 3 (128k)',    'meta-llama/Llama-3.2-1B'),
    ('Qwen3 (151k)',      'Qwen/Qwen3-Embedding-0.6B'),
    ('Gemma 2 (256k)',    'google/gemma-2-2b'),
]

hf_tokenizers = {}
for label, model_id in _candidates:
    try:
        hf_tokenizers[label] = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        print(f'  ✓ {label}')
    except Exception as e:
        print(f'  ✗ {label} — skipped ({e})')

tok_records = []
for tok_label, tokenizer in hf_tokenizers.items():
    for path in PROMPT_FILES:
        data = json.loads(path.read_text())
        vmap = variants_for(path, data)
        for variant, text in vmap.items():
            tok_records.append({
                'tokenizer': tok_label,
                'file':      path.name,
                'group':     'JSON' if variant in JSON_VARIANTS else 'YAML',
                'variant':   variant,
                'tokens':    len(tokenizer.encode(text)),
            })

df_tok = pd.DataFrame(tok_records)
df_tok['variant'] = pd.Categorical(df_tok['variant'], categories=VARIANT_NAMES, ordered=True)

# Median across files — variants as rows, tokenizers as columns
summary_tok = (df_tok
    .groupby(['group', 'variant', 'tokenizer'], observed=True)['tokens']
    .median()
    .astype(int)
    .unstack('tokenizer'))
summary_tok.columns.name = None

# Sort columns left-to-right by vocabulary size
_vocab_k = lambda col: int(re.search(r'(\d+)k', col).group(1)) if re.search(r'\d+k', col) else 0
summary_tok = summary_tok[sorted(summary_tok.columns, key=_vocab_k)]

display(summary_tok.style
        .highlight_min(color='#c7f2a4', axis=0)
        .highlight_max(color='#f2c4c4', axis=0)
        .set_caption('Median token counts by tokenizer — green = min per tokenizer, red = max per tokenizer'))

# Per-file rankings — computed from df_wide
PYYAML_VARIANTS = ['Original JSON', 'Minified JSON', 'Pretty JSON', 'Standard YAML (PyYAML)', 'YAML Literal (PyYAML)']
YQ_VARIANTS     = ['Original JSON', 'Minified JSON', 'Pretty JSON', 'Standard YAML (yq)',     'YAML Literal (yq)']

SHORT = {
    'Original JSON':          'Orig JSON',
    'Minified JSON':          'Min JSON',
    'Pretty JSON':            'Pretty JSON',
    'Standard YAML (PyYAML)': 'Std YAML (py)',
    'Standard YAML (yq)':     'Std YAML (yq)',
    'YAML Literal (PyYAML)':  'Lit YAML (py)',
    'YAML Literal (yq)':      'Lit YAML (yq)',
}

flip_records = []
for path in PROMPT_FILES:
    fname = path.name
    row   = df_wide.loc[fname]
    py_sorted = row[PYYAML_VARIANTS].sort_values()
    yq_sorted = row[YQ_VARIANTS].sort_values()
    flip_records.append({
        'file':           fname,
        'PyYAML ranking': ' < '.join(SHORT[v] for v in py_sorted.index),
        'yq ranking':     ' < '.join(SHORT[v] for v in yq_sorted.index),
        'yq flip':        row['YAML Literal (yq)'] > row['Pretty JSON'],
        'py flip':        row['YAML Literal (PyYAML)'] > row['Pretty JSON'],
    })

df_flips = pd.DataFrame(flip_records).set_index('file')
display(df_flips)

print(f"\nFiles where Lit YAML (yq)     > Pretty JSON: {df_flips['yq flip'].sum()}/{len(PROMPT_FILES)}")
print(f"Files where Lit YAML (PyYAML) > Pretty JSON: {df_flips['py flip'].sum()}/{len(PROMPT_FILES)}")

f_records = []
for path in PROMPT_FILES:
    data      = json.loads(path.read_text())
    t_optimal = tokens(yaml_standard(data, width=float('inf')))
    f_records.append({
        'file':               path.name,
        'Orig JSON':          int(df_wide.loc[path.name, 'Original JSON']),
        'Min JSON':           int(df_wide.loc[path.name, 'Minified JSON']),
        'YAML w=80':          int(df_wide.loc[path.name, 'Standard YAML (PyYAML)']),
        'YAML w=∞':           t_optimal,
        # Wrap overhead = tokens eliminated by disabling line-wrapping (w=80 → w=∞)
        'Wrap overhead':      int(df_wide.loc[path.name, 'Standard YAML (PyYAML)']) - t_optimal,
        'Pretty JSON':        int(df_wide.loc[path.name, 'Pretty JSON']),
        'YAML w=∞ vs Pretty': t_optimal - int(df_wide.loc[path.name, 'Pretty JSON']),
    })

df_f = pd.DataFrame(f_records).set_index('file')
df_f.loc['median'] = df_f.median().round(0).astype(int)

display(df_f.style
        .format({'Wrap overhead':      lambda x: f'+{int(x)}',
                 'YAML w=∞ vs Pretty': lambda x: f'{int(x):+d}'})
        .bar(subset=['Wrap overhead'], color='#d5f5e3', vmin=0)
        .map(lambda v: 'color: green' if isinstance(v, (int, float)) and v < 0 else '',
             subset=['YAML w=∞ vs Pretty']))

print("Round-trip fidelity with width=∞ — all prompt files:\n")
all_pass = True
for path in PROMPT_FILES:
    data = json.loads(path.read_text())
    rt = yaml.safe_load(yaml_standard(data, width=float('inf')))
    ok = rt == data
    if not ok:
        all_pass = False
    print(f"  {path.name:<45} {'✓' if ok else '✗ FAILED'}")

print(f"\nAll files pass: {all_pass}")
if all_pass:
    print("""
Recommended configuration for prepare_prompt.py:

    yaml.dump(data, allow_unicode=True, default_flow_style=False,
              sort_keys=False, width=float('inf'))
""")
else:
    failed = [
        path.name for path in PROMPT_FILES
        if yaml.safe_load(yaml_standard(json.loads(path.read_text()), width=float('inf')))
           != json.loads(path.read_text())
    ]
    print(f"\n⚠ Round-trip failures detected in: {failed}")
    print("Investigate with find_diff() from Section C before adopting width=∞.")
    print("Likely cause: a value type that yaml.safe_load coerces on parse")
    print("(e.g. a string that looks like a date or boolean).")