"""Render the offline index-space JSONL using matplotlib."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

parser = argparse.ArgumentParser(__doc__)
parser.add_argument('results', type=Path)
parser.add_argument('output', type=Path)
args = parser.parse_args()
rows = [json.loads(line) for line in args.results.read_text().splitlines()]
colors = {'zkv': '#225ea8', 'rocks': '#bc5b13'}
names = {'zkv': 'ZKV', 'rocks': 'RocksDB'}
plt.rcParams.update({'font.size': 13, 'axes.spines.top': False,
                     'axes.spines.right': False, 'svg.fonttype': 'none'})
fig, axes = plt.subplots(3, 1, figsize=(7.5, 13), layout='constrained')
summary = {}
for engine in colors:
    data = [r for r in rows if r['engine'] == engine]
    by_phase = {r['phase']: r for r in data}
    assert len(data) == 30 and 'verify-final-reopen' in by_phase
    growth = [r for r in data if r['phase'].startswith('grow-')]
    axes[0].plot([r['count'] * 2 / 1e6 for r in growth],
                 [r['final']['logical_bytes'] / 2**20 for r in growth],
                 'o-', color=colors[engine], label=names[engine], linewidth=2)
    phases = ['settle-ingested', 'settle-duplicate', 'settle-refill-3']
    values = [by_phase[p]['final']['logical_bytes'] / 2**20 for p in phases]
    offset = -.18 if engine == 'zkv' else .18
    bars = axes[1].bar([i + offset for i in range(3)], values, .34,
                       color=colors[engine], label=names[engine])
    axes[1].bar_label(bars, fmt='%.1f', fontsize=11, padding=3)
    groups = [
        [r for r in data if r['phase'].startswith('grow-') or r['phase'] == 'verify-ingested'],
        [r for r in data if r['phase'] == 'duplicate'],
        [r for r in data if r['phase'].startswith(('delete-', 'refill-', 'verify-deleted-', 'verify-final-'))],
        [r for r in data if r['phase'].startswith('settle-')],
    ]
    writes = [sum(r['os_block_write_bytes'] for r in group) / 2**20 for group in groups]
    bars = axes[2].bar([i + offset for i in range(4)], writes, .34,
                       color=colors[engine], label=names[engine])
    axes[2].bar_label(bars, labels=['<0.1' if 0 < v < .1 else f'{v:.1f}' for v in writes], fontsize=11, padding=3)
    peak = max(r['sampled_peak']['logical_bytes'] for r in data)
    summary[engine] = dict(settled_bytes=[by_phase[p]['final']['logical_bytes'] for p in phases],
                           bytes_per_membership=by_phase['settle-ingested']['final']['logical_bytes'] / 2e6,
                           sampled_peak_bytes=peak, phase_write_bytes=[round(v * 2**20) for v in writes])
axes[0].set(title='Disk growth during ingestion', xlabel='Live memberships (millions)', ylabel='Database files (MiB)')
axes[1].set(title='Retained disk after maintenance', ylabel='Database files (MiB)',
            xticks=range(3), xticklabels=['Ingested', 'Duplicate pass', 'Three churn cycles'])
axes[2].set(title='Filesystem writes by workload phase', ylabel='OS-accounted writes (MiB)',
            xticks=range(4), xticklabels=['Ingest\nlifecycle', 'Duplicate\npass', 'Churn\nlifecycle', 'Explicit\nsettling'])
for axis in axes:
    axis.set_ylim(0, axis.get_ylim()[1] * 1.12)
    axis.grid(axis='y', alpha=.18)
    axis.set_axisbelow(True)
axes[0].legend(loc='upper left', frameon=False)
fig.suptitle('Zlay collection-index storage comparison\nIdentical offline input · synchronous writes', fontsize=17)
fig.savefig(args.output.with_suffix('.svg'), facecolor='white')
fig.savefig(args.output.with_suffix('.png'), dpi=150, facecolor='white')
args.output.with_suffix('.json').write_text(json.dumps(summary, indent=2) + '\n')
print(json.dumps(summary, indent=2))