"""Generate article figures and a separate, explicitly lossy media experiment.

Usage: python visualize.py BENCHMARK_OUTPUT PUBLIC_ASSET_DIRECTORY
Requires matplotlib and Pillow, plus FFmpeg on PATH.
Pass --figures-only to preserve the existing media measurements.
"""
import io
import json
import math
import re
import statistics
import subprocess
import sys
import time
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from PIL import Image

root, assets = map(Path, sys.argv[1:3])
assets.mkdir(parents=True, exist_ok=True)
results = json.loads((root / 'results.json').read_text())
plt.rcParams.update({'font.family': 'DejaVu Sans', 'font.size': 11,
                     'axes.spines.top': False, 'axes.spines.right': False,
                     'svg.fonttype': 'none', 'figure.facecolor': 'white',
                     'axes.facecolor': 'white', 'text.color': '#202020'})


def save(fig, name):
    svg = assets / f'{name}.svg'
    fig.savefig(svg, bbox_inches='tight')
    svg.write_text('\n'.join(line.rstrip() for line in svg.read_text(encoding='utf-8').splitlines()) + '\n',
                   encoding='utf-8', newline='\n')
    fig.savefig(assets / f'{name}.png', bbox_inches='tight', dpi=160)
    plt.close(fig)


files = ['guide.txt', 'events.jsonl', 'gradient.png', 'gradient.jpg', 'motion.mp4',
         'noise.png', 'random.bin', 'report-compressed.pdf', 'report-plain.pdf']
names = ['Original prose', 'Synthetic logs', 'Gradient PNG', 'Gradient JPEG',
         'H.264 video', 'Noise PNG', 'Random bytes', 'Compressed PDF', 'Plain PDF']
codecs = ['gzip-6', 'xz-9', 'brotli-11', 'zstd-3']
colors = ['#a5a5a5', '#444444', '#b6632a', '#ef792b']
fig, ax = plt.subplots(figsize=(10, 7))
for j, (codec, color) in enumerate(zip(codecs, colors)):
    vals = [next(r['percent'] for r in results if r['file'] == f and r['codec'] == codec) for f in files]
    bars = ax.barh([i + j * 0.19 for i in range(len(files))], vals, height=0.18, label=codec, color=color)
    for bar, value in zip(bars, vals):
        ax.text(value + 1, bar.get_y() + bar.get_height() / 2, f'{value:.1f}', va='center', fontsize=8)
ax.set_yticks([i + .285 for i in range(len(files))], names)
ax.invert_yaxis()
ax.axvline(100, color='#222222', linestyle=':', linewidth=1)
ax.set_xlim(0, 118)
ax.set_xlabel('Stored size as % of original bytes; lower is smaller')
ax.set_title('Some synthetic files compress well. Others grow.', loc='left', fontsize=16, pad=20)
ax.legend(ncol=4, loc='upper center', bbox_to_anchor=(.5, -.12), frameon=False)
fig.tight_layout()
save(fig, 'compression-ratios')

fig, axes = plt.subplots(1, 2, figsize=(11, 5))
for i, name in enumerate(codecs + ['zstd-19', 'zstd-22', 'lz4-0']):
    r = next(r for r in results if r['file'] == 'events.jsonl' and r['codec'] == name)
    color = '#ef792b' if name.startswith('zstd') else '#444444'
    axes[0].scatter(r['encode_ms'], r['percent'], color=color, s=50)
    shifts = {'gzip-6': (5, -18), 'xz-9': (-32, -24), 'zstd-19': (6, 24), 'zstd-22': (-35, 7)}
    axes[0].annotate(name, (r['encode_ms'], r['percent']), xytext=shifts.get(name, (5, 5)), textcoords='offset points', fontsize=9)
    peak = r['process_peak_working_set_bytes']
    axes[1].barh(name, (peak or 0) / 2**20, color=color)
    if peak is None:
        axes[1].text(0, name, 'Unavailable on this platform', va='center', fontsize=8)
axes[0].set_xscale('log')
axes[0].set_xlabel('Median encoding time, ms (log scale)')
axes[0].set_ylabel('Stored bytes, % of original logs')
axes[0].set_title('Smaller files can take much longer')
axes[0].margins(x=.25, y=.25)
axes[1].set_xlabel('Peak process working set, MiB')
axes[1].set_title('Memory includes Python, input and verification')
fig.suptitle('One 846,281-byte synthetic log file, one Windows machine', x=.05, ha='left', fontsize=15)
fig.tight_layout()
save(fig, 'time-and-memory')

fig, axes = plt.subplots(1, 2, figsize=(10, 4.8))
chunks = [4096, 65536, 1048576]
for f, label, color in [('events.jsonl', 'Logs', '#ef792b'), ('random.bin', 'Random bytes', '#444444')]:
    vals = [next(r['percent'] for r in results if r['file'] == f and r['codec'] == f'blocks-{c}') for c in chunks]
    axes[0].plot(range(3), vals, marker='o', label=label, color=color)
axes[0].set_xticks(range(3), ['4 KiB', '64 KiB', '1 MiB'])
axes[0].set_xlabel('Independent block size')
axes[0].set_ylabel('Stored bytes including index, %')
axes[0].legend(frameon=False)
axes[0].set_title('Measured storage in the toy container')
axes[1].bar(['4 KiB', '64 KiB', '1 MiB'], [1, 16, 256], color=['#aaa', '#666', '#ef792b'])
axes[1].set_ylabel('Decoded bytes / requested bytes')
axes[1].set_title('Model: one aligned 4 KiB read, full block')
for i, value in enumerate([1, 16, 256]):
    axes[1].text(i, value + 4, f'{value}x', ha='center')
axes[1].set_ylim(0, 290)
fig.suptitle('Block size changes the storage and read cost', x=.05, ha='left', fontsize=16)
fig.tight_layout()
save(fig, 'block-tradeoff')

fig, ax = plt.subplots(figsize=(11, 4.6))
ax.axis('off')
labels = [
    ('Input bytes', 'Keep originals and hashes'),
    ('Predict patterns', 'Dictionary, context or neural state'),
    ('Encode surprises', 'Entropy coder writes\nthe bitstream'),
    ('Restore bytes', 'Same state, model and rules needed'),
]
for i, (title, detail) in enumerate(labels):
    x = .015 + i * .25
    ax.text(x, .78, title, transform=ax.transAxes, fontsize=14, weight='bold',
            bbox=dict(boxstyle='round,pad=.5', facecolor='#fff1e6', edgecolor='#ef792b'))
    ax.text(x, .54, detail.replace(', ', ',\n').replace(' and ', '\nand '),
            transform=ax.transAxes, fontsize=11, va='top', wrap=True)
    if i < 3:
        ax.annotate('', xy=(x + .24, .81), xytext=(x + .195, .81), xycoords='axes fraction',
                    arrowprops=dict(arrowstyle='->', color='#444'))
ax.text(.015, .25, 'Where work grows: match searches, context tables, model weights, KV cache and sequential prediction.',
        transform=ax.transAxes, fontsize=12)
ax.text(.015, .12, 'Research direction: skip weak predictors and repeated context work; measure decoding cost too.',
        transform=ax.transAxes, fontsize=12)
ax.set_title('A smaller bitstream can move the cost into decoder state', loc='left', fontsize=16, pad=12)
save(fig, 'predictor-state')

fig, ax = plt.subplots(figsize=(10, 3.4))
ax.axis('off')
stages = [(.01, 'Original', 'AAAAAAAAAAAA', '12 ASCII bytes'),
          (.37, 'Write a count and a letter', '12:A', '4 ASCII bytes'),
          (.73, 'Read the instruction', 'AAAAAAAAAAAA', '12 ASCII bytes restored')]
for x, title, value, detail in stages:
    ax.text(x, .72, title, transform=ax.transAxes, fontsize=12)
    ax.text(x, .48, value, transform=ax.transAxes, fontsize=15,
            bbox=dict(boxstyle='round,pad=.5', facecolor='#fff1e6', edgecolor='#ef792b'))
    ax.text(x, .24, detail, transform=ax.transAxes, fontsize=11)
for start, end in [(.29, .35), (.64, .71)]:
    ax.annotate('', xy=(end, .51), xytext=(start, .51), xycoords='axes fraction',
                arrowprops=dict(arrowstyle='->', color='#444'))
ax.text(.01, .02, 'Teaching notation only. Real formats need rules for counts, escapes and other symbols.',
        transform=ax.transAxes, fontsize=10)
ax.set_title('A small lossless example: write the repetition as an instruction', loc='left', fontsize=15, pad=10)
save(fig, 'repetition-example')

selected = [next(r for r in results if r['file'] == 'events.jsonl' and r['codec'] == name)
            for name in ['lz4-0', 'zstd-3', 'bzip2-9', 'brotli-11']]
fig, axes = plt.subplots(1, 3, figsize=(12, 4.8))
labels = [r['codec'] for r in selected]
for ax, field, title in zip(axes, ['compressed_bytes', 'encode_ms', 'decode_ms'],
                           ['Stored file bytes', 'Encode time (ms)', 'Decode time (ms)']):
    values = [r[field] for r in selected]
    panel_labels = labels[:]
    if field == 'compressed_bytes':
        values.insert(0, selected[0]['original_bytes'])
        panel_labels.insert(0, 'Original')
    bars = ax.barh(panel_labels, values, color=['#ef792b' if label == 'zstd-3' else '#555'
                                              for label in panel_labels])
    ax.invert_yaxis()
    ax.set_title(title, loc='left', pad=12)
    ax.set_xlim(0, max(values) * 1.4)
    ax.ticklabel_format(style='plain', axis='x')
    ax.tick_params(axis='x', labelsize=8)
    for bar, value in zip(bars, values):
        label = f'{value:,.0f}' if field == 'compressed_bytes' else f'{value:,.2f}'
        ax.text(value + max(values) * .025, bar.get_y() + bar.get_height() / 2,
                label, va='center', fontsize=9)
fig.suptitle('The same log file: saving bytes and saving time give different winners',
             x=.02, ha='left', fontsize=15)
fig.tight_layout()
save(fig, 'logs-size-and-time')

if '--figures-only' in sys.argv[3:]:
    sys.exit(0)

# Media conversion is a separate experiment. Exact pixels do not imply exact file bytes.
original = Image.open(root / 'samples/gradient.ppm').convert('RGB')
media = []
for fmt, options in [('PNG', {'compress_level': 9}), ('WEBP', {'lossless': True, 'method': 6}),
                     ('JPEG', {'quality': 60}), ('JPEG', {'quality': 90})]:
    durations = []
    for _ in range(3):
        buffer = io.BytesIO()
        start = time.perf_counter_ns()
        original.save(buffer, format=fmt, **options)
        durations.append((time.perf_counter_ns() - start) / 1e6)
    blob = buffer.getvalue()
    restored = Image.open(io.BytesIO(blob)).convert('RGB')
    a, b = original.tobytes(), restored.tobytes()
    mse = sum((x-y)**2 for x, y in zip(a, b)) / len(a)
    media.append(dict(format=fmt, options=options, bytes=len(blob),
                      encode_ms=statistics.median(durations), pixels_exact=a == b,
                      rgb_psnr_db=None if mse == 0 else 10 * math.log10(255**2 / mse)))
target = root / 'motion-crf35.mp4'
command = ['ffmpeg', '-hide_banner', '-y', '-i', str(root / 'samples/motion.mp4'),
           '-an', '-c:v', 'libx264', '-preset', 'medium', '-crf', '35', '-threads', '1',
           '-movflags', '+faststart', str(target)]
start = time.perf_counter()
subprocess.run(command, check=True, capture_output=True)
elapsed = time.perf_counter() - start
comparison = subprocess.run(['ffmpeg', '-hide_banner', '-i', str(root / 'samples/motion.mp4'),
    '-i', str(target), '-lavfi', 'psnr', '-f', 'null', '-'], check=True, capture_output=True, text=True)
psnr = float(re.search(r'average:([\d.]+)', comparison.stderr).group(1))
(root / 'media.json').write_text(json.dumps(dict(image=media, video=dict(
    original_bytes=(root / 'samples/motion.mp4').stat().st_size, converted_bytes=target.stat().st_size,
    encode_seconds=elapsed, repeats=1, yuv_psnr_db=psnr, command=command,
    note='Lossy transcode of an already lossy generated source, not exact file recovery.')), indent=2))
print(json.dumps(media, indent=2))
print('Video PSNR', psnr)