# make_charts.py - render benchmark charts for DNA-DiskChat-2B.
import json, numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

cpu = json.load(open('/root/dna/cpu_bench.json'))
ssd = json.load(open('/root/dna/ssd_bench.json'))
BLUE, GREEN, ORANGE, GREY = '#2563eb', '#059669', '#ea580c', '#64748b'

def bars(data, color, title, fname, mean_label):
    v = data['tok_s']; m, s = data['mean'], data['std']
    fig, ax = plt.subplots(figsize=(7, 4))
    x = np.arange(1, len(v) + 1)
    ax.bar(x, v, color=color, alpha=.85, width=.6)
    ax.axhline(m, color='black', ls='--', lw=1.2, label=f'{mean_label}: {m:.1f} tok/s (±{s:.1f})')
    ax.set_xlabel('run'); ax.set_ylabel('tokens / second'); ax.set_title(title)
    ax.set_xticks(x); ax.set_ylim(0, max(v) * 1.25); ax.legend(loc='lower right')
    ax.grid(axis='y', alpha=.3)
    fig.tight_layout(); fig.savefig(fname, dpi=130); plt.close(fig)
    print('wrote', fname)

bars(cpu, BLUE, 'DNA-DiskChat-2B — CPU inference (table warm in RAM)',
     '/root/dna/cpu_throughput.png', 'mean')
bars(ssd, GREEN, 'DNA-DiskChat-2B — CPU + SSD inference (256 B/token from disk)',
     '/root/dna/ssd_throughput.png', 'mean')

# combined comparison: throughput + per-token I/O
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.2))
labels = ['GPU\ntrain', 'CPU\n(RAM)', 'CPU+SSD\n(disk)']
vals = [45574, cpu['mean'], ssd['mean']]
cols = [GREY, BLUE, GREEN]
b = ax1.bar(labels, vals, color=cols, alpha=.88)
ax1.set_ylabel('tokens / second (log)'); ax1.set_yscale('log')
ax1.set_title('Throughput: GPU training vs CPU/SSD inference')
for r, val in zip(b, vals):
    ax1.text(r.get_x() + r.get_width() / 2, val * 1.05, f'{val:,.0f}', ha='center', va='bottom', fontsize=9)
ax1.grid(axis='y', alpha=.3)

io_labels = ['DNA-2B\n(this model)', 'DiskChat-0.5B-v2']
io_vals = [ssd['io_bytes_per_token'], 8192]
b2 = ax2.bar(io_labels, io_vals, color=[GREEN, ORANGE], alpha=.88)
ax2.set_ylabel('bytes read / token'); ax2.set_title('Per-token SSD I/O (2 routed rows)')
for r, val in zip(b2, io_vals):
    ax2.text(r.get_x() + r.get_width() / 2, val + 120, f'{val:,} B', ha='center', va='bottom', fontsize=9)
ax2.grid(axis='y', alpha=.3)
fig.tight_layout(); fig.savefig('/root/dna/throughput_compare.png', dpi=130); plt.close(fig)
print('wrote /root/dna/throughput_compare.png')
print('CHARTS_DONE')