#!/usr/bin/env python3
"""Recomputes every table number of https://soba-app.com/accuracy/models/ from the scans CSV.

Standard library only. Keep this file next to soba-model-benchmark-2026-10-scans.csv.

    python3 reproduce.py           # prints all tables
    python3 reproduce.py --json    # the same numbers as JSON
    python3 reproduce.py --check   # compares with expected_output.txt and expected_results.json
"""
import argparse
import contextlib
import csv
import io
import json
import math
import random
import statistics as st
import sys
from operator import mul
from pathlib import Path

HERE = Path(__file__).resolve().parent
SCANS = HERE / 'soba-model-benchmark-2026-10-scans.csv'
EXPECTED_TEXT = HERE / 'expected_output.txt'
EXPECTED_JSON = HERE / 'expected_results.json'
SEED, REPS = 20261007, 10_000
THRESHOLD_G, GROSS_G = 10.0, 30.0

G36, G37, G38 = 'google/gemini-3.6-flash', 'google/gemini-3.7-flash', 'google/gemini-3.8-flash'
LITE, LUNA, PRO, SOL = 'google/gemini-3.5-flash-lite', 'openai/gpt-6-luna', 'openai/gpt-6-luna-pro', 'openai/gpt-6.1-sol'
CLAUDE, PREVIEW = 'anthropic/claude-sonnet-5.5', 'google/gemini-3.1-pro-preview'
EIGHT = (G37, G38, LITE, G36, LUNA, PRO, SOL, CLAUDE)
NAMES = {G36: 'Gemini 3.6 Flash', G37: 'Gemini 3.7 Flash', G38: 'Gemini 3.8 Flash', LITE: 'Gemini 3.5 Flash-Lite',
         LUNA: 'GPT-6 Luna', PRO: 'GPT-6 Luna Pro', SOL: 'GPT-6.1 Sol', CLAUDE: 'Claude Sonnet 5.5',
         PREVIEW: 'Gemini 3.1 Pro Preview'}
# A config is 'model|reasoning'. Sorted by the share within 10 g.
POOL = (f'{G36}|medium', f'{G37}|medium', f'{G37}|low', f'{G38}|low', f'{G36}|low', f'{LITE}|low', f'{LUNA}|low')
HOLDOUT = (f'{G36}|medium', f'{G36}|low', f'{G37}|low', f'{G37}|medium', f'{LITE}|low', f'{G38}|low', f'{LUNA}|low')
HOLDOUT_LOW = tuple(c for c in HOLDOUT if c.endswith('|low'))
BASE = f'{G36}|low'
CARB_BANDS = ((0, 15), (15, 30), (30, 60), (60, math.inf))
WEIGHT_BANDS = ((0, 100), (100, 200), (200, 400), (400, 600), (600, math.inf))
# First-set meals quoted in the text. Holdout meals stay out of dish-level discussion.
EXAMPLES = {'coffee': 'b325a1fe', 'sweet_tea': 'c4975569', 'potatoes': '9a8f2076', 'curry_rice': '0e206269',
            'breakfast_plate': 'dish_1563551220', 'cherry_tomatoes': 'dish_1558641287', 'spinach': 'dish_1565810969',
            'five_cookies': '2d4c9ab3', 'pancake': '5a72e137', 'sweet_potato': 'b398db11', 'mooncake': '818a6db2',
            'brownies_label': 'dish_1566501575'}


def load():
    rows = []
    with SCANS.open(newline='') as f:
        for r in csv.DictReader(f):
            r['ref'], r['ref_mass'] = float(r['reference_carbs_g']), float(r['reference_mass_g'])
            r['est'] = float(r['estimate_g']) if r['estimate_g'] else None
            r['mass'] = float(r['estimate_mass_g']) if r['estimate_mass_g'] else None
            r['seconds'], r['cost'], r['requests'] = float(r['seconds']), float(r['cost_usd']), int(r['requests'])
            r['key'] = r['dataset'] + '/' + r['dish_id']
            rows.append(r)
    return rows


def within(est, ref, limit=THRESHOLD_G):
    return est is not None and abs(est - ref) <= limit


def quantile(ordered, q):
    pos = q * (len(ordered) - 1)
    low, high = int(pos), min(int(pos) + 1, len(ordered) - 1)
    return ordered[low] + (ordered[high] - ordered[low]) * (pos - low)


class Resampler:
    """Cluster bootstrap: clusters are resampled within each dataset, and the datasets weigh the same."""

    def __init__(self, dishes):
        rng = random.Random(SEED)
        self.groups = []
        for dataset in sorted({d['dataset'] for d in dishes}):
            index = [i for i, d in enumerate(dishes) if d['dataset'] == dataset]
            clusters = sorted({dishes[i]['cluster'] for i in index})
            members = [[i for i in index if dishes[i]['cluster'] == c] for c in clusters]
            sizes = [len(m) for m in members]
            draws = []
            for _ in range(REPS):
                counts = [0] * len(clusters)
                for _ in clusters:
                    counts[int(rng.random() * len(clusters))] += 1
                draws.append((counts, sum(map(mul, counts, sizes))))
            self.groups.append((members, draws))

    def interval(self, values):
        totals = [0.0] * REPS
        for members, draws in self.groups:
            sums = [sum(values[i] for i in m) for m in members]
            for rep, (counts, bottom) in enumerate(draws):
                totals[rep] += sum(map(mul, counts, sums)) / bottom / len(self.groups)
        totals.sort()
        return quantile(totals, .025) * 100, quantile(totals, .975) * 100


def row(dishes, scans, resampler):
    """Replicate scans are averaged per dish. A failed scan counts as a miss and has no error."""
    hit, abs_err, err, gross = [], [], [], []
    for d in dishes:
        group = scans[d['key']]
        hit.append(st.fmean(within(s['est'], s['ref']) for s in group))
        answered = [s['est'] - s['ref'] for s in group if s['est'] is not None]
        if answered:
            abs_err.append(st.fmean(abs(e) for e in answered))
            err.append(st.fmean(answered))
            gross.append(st.fmean(abs(e) > GROSS_G for e in answered))
    every = [s for d in dishes for s in scans[d['key']]]
    low, high = resampler.interval(hit)
    return {'n': len(dishes), 'hits': sum(hit), 'share': st.fmean(hit) * 100, 'low': low, 'high': high,
            'answered': len(abs_err), 'mae': st.fmean(abs_err), 'median_abs': st.median(abs_err), 'bias': st.fmean(err),
            'over30': sum(gross), 'over30_share': sum(gross) / len(gross) * 100,
            'failed': sum(s['est'] is None for s in every), 'cost': st.fmean(s['cost'] for s in every),
            'seconds': st.median(s['seconds'] for s in every)}, hit


def difference(a, b, resampler):
    delta = [x - y for x, y in zip(a, b)]
    low, high = resampler.interval(delta)
    return {'diff': st.fmean(delta) * 100, 'low': low, 'high': high}


def table(dishes, select, configs, resampler, bases=()):
    """One row per config, and paired differences from each base config."""
    rows, hits = {}, {}
    for config in configs:
        rows[config], hits[config] = row(dishes, select(config), resampler)
    pairs = {b: {c: difference(hits[c], hits[b], resampler) for c in configs if c != b} for b in bases}
    return {'rows': rows, 'vs': pairs}, hits


def band_table(dishes, select, configs):
    """By reference carbs. Failed scans count as misses, errors use answered scans."""
    out = {}
    for config in configs:
        scans, cells = select(config), []
        for lo, hi in CARB_BANDS:
            part = [d for d in dishes if lo <= d['ref'] < hi]
            errors = [s['est'] - s['ref'] for d in part for s in scans[d['key']] if s['est'] is not None]
            cells.append({'n': len(part), 'hits': sum(within(scans[d['key']][0]['est'], d['ref']) for d in part),
                          'answered': len(errors), 'mean_error': st.fmean(errors), 'median_error': st.median(errors),
                          'too_low': sum(e < 0 for e in errors), 'lowest': min(errors), 'highest': max(errors)})
        out[config] = cells
    return out


def weight_table(dishes, select, configs):
    """Estimated weight divided by reference weight. Failed scans and placeholder answers stay out."""
    out = {}
    for config in configs:
        scans, cells = select(config), []
        for lo, hi in WEIGHT_BANDS:
            part = [d for d in dishes if lo <= d['ref_mass'] < hi]
            ratios = [s['mass'] / d['ref_mass'] for d in part for s in scans[d['key']]
                      if s['mass'] is not None and s['note'] != 'placeholder_item']
            cells.append({'n': len(part), 'scans': len(ratios), 'too_high': sum(r > 1 for r in ratios),
                          'median_ratio': st.median(ratios) if ratios else None})
        out[config] = cells
    return out


def compute(scans):
    by_config = {}
    for s in scans:
        by_config.setdefault((s['set'], f"{s['model']}|{s['reasoning']}"), {}).setdefault(s['key'], []).append(s)

    def dish_list(name):
        seen = {}
        for s in scans:
            if s['set'] == name:
                seen.setdefault(s['key'], {k: s[k] for k in ('key', 'dish_id', 'dataset', 'cluster', 'ref', 'ref_mass')})
        return list(seen.values())

    first, holdout = dish_list('food200'), dish_list('holdout')
    res = {}

    # Tables 1 and 2: eight models on the first set, one run each.
    runs = {m: {k: [s for s in v if s['replicate'] == '1'] for k, v in by_config[('food200', f'{m}|low')].items()}
            for m in EIGHT}
    resample_first = Resampler(first)
    res['first'], hits = table(first, runs.get, EIGHT, resample_first, (G37, G38, LITE, G36))

    # Table 3: the same run by dataset.
    for dataset in ('nutrition5k', 'snapme'):
        part = [d for d in first if d['dataset'] == dataset]
        res['first_' + dataset], _ = table(part, runs.get, EIGHT, Resampler(part), (G36,))

    # Table 4: Claude, every view of the same run.
    claude = {k: v[0] for k, v in runs[CLAUDE].items()}
    placeholders = {k for k, s in claude.items() if s['note'] == 'placeholder_item'}
    answered = [d for d in first if claude[d['key']]['est'] is not None]
    real = [d for d in answered if d['key'] not in placeholders]
    one = [s for s in claude.values() if s['requests'] == 1]
    two = [s for s in claude.values() if s['requests'] == 2]
    res['claude'] = {
        'answered': table(answered, runs.get, EIGHT, Resampler(answered), (CLAUDE,))[0],
        'real': table(real, runs.get, EIGHT, Resampler(real), (CLAUDE,))[0],
        'hits_without_placeholders': sum(within(claude[d['key']]['est'], d['ref']) for d in real),
        'placeholders': len(placeholders), 'providers': sorted({s['provider'] for s in claude.values()}),
        'requests': sum(s['requests'] for s in claude.values()),
        'failed_by_dataset': {ds: sum(s['est'] is None and s['dataset'] == ds for s in claude.values())
                              for ds in ('nutrition5k', 'snapme')},
        'one_request': {'n': len(one), 'hits': sum(within(s['est'], s['ref']) for s in one),
                        'cost': st.fmean(s['cost'] for s in one), 'seconds': st.median(s['seconds'] for s in one)},
        'two_requests': {'n': len(two), 'failed': sum(s['est'] is None for s in two),
                         'placeholders': sum(s['note'] == 'placeholder_item' for s in two),
                         'real': sum(s['est'] is not None and s['note'] != 'placeholder_item' for s in two),
                         'cost': st.fmean(s['cost'] for s in two), 'seconds': st.median(s['seconds'] for s in two)}}

    # Table 5: every config that ran on all 400 meals. Replicate runs are averaged per dish.
    both = first + holdout

    def pool_scans(config):
        return {**by_config[('food200', config)], **by_config[('holdout', config)]}

    resample_both = Resampler(both)
    res['pool'], hits_pool = table(both, pool_scans, POOL, resample_both, (BASE,))
    res['pool']['medium_37_vs_low_37'] = difference(hits_pool[f'{G37}|medium'], hits_pool[f'{G37}|low'], resample_both)
    res['pool']['thresholds'] = {
        c: {str(g): st.fmean(st.fmean(within(s['est'], s['ref'], g) for s in pool_scans(c)[d['key']]) for d in both) * 100
            for g in (5, 10, 15, 20)} for c in POOL}

    # Table 6: the holdout, one run per config.
    holdout_scans = lambda c: by_config[('holdout', c)]
    res['holdout'], _ = table(holdout, holdout_scans, HOLDOUT, Resampler(holdout), (BASE,))

    # Tables 7 and 8: plate size and portion weight.
    n5k_first = [d for d in first if d['dataset'] == 'nutrition5k']
    n5k_holdout = [d for d in holdout if d['dataset'] == 'nutrition5k']
    res['carb_bands'] = {'first': band_table(first, runs.get, EIGHT),
                         'first_nutrition5k': band_table(n5k_first, runs.get, EIGHT),
                         'holdout': band_table(holdout, holdout_scans, HOLDOUT_LOW)}
    res['weight_bands'] = {'first': weight_table(first, runs.get, EIGHT),
                           'first_nutrition5k': weight_table(n5k_first, runs.get, EIGHT),
                           'holdout': weight_table(holdout, holdout_scans, HOLDOUT_LOW),
                           'holdout_nutrition5k': weight_table(n5k_holdout, holdout_scans, HOLDOUT_LOW)}

    # Repeat runs of three models on the first set.
    res['repeat'] = {}
    for m in (G36, G37, G38):
        pairs = by_config[('food200', f'{m}|low')]
        same, second, gaps = 0, 0, []
        for d in first:
            a, b = pairs[d['key']]
            same += within(a['est'], d['ref']) == within(b['est'], d['ref'])
            second += within(b['est'], d['ref'])
            gaps.append(abs(a['est'] - b['est']))
        gaps.sort()
        res['repeat'][m] = {'same_outcome': same / len(first) * 100, 'second_run': second / len(first) * 100,
                            'median_gap': st.median(gaps), 'p90_gap': quantile(gaps, .9), 'max_gap': gaps[-1]}

    # A fixed answer in whole grams as a floor for the share within 10 g.
    shares = {g: sum(within(float(g), d['ref']) for d in first) / len(first) * 100 for g in range(101)}
    best = max(sorted(shares), key=shares.get)
    res['fixed_answer'] = {'grams': best, 'first': shares[best],
                           'holdout': sum(within(float(best), d['ref']) for d in holdout) / len(holdout) * 100}

    # Agreement of the eight models on the first set.
    ests = {d['key']: [runs[m][d['key']][0]['est'] for m in EIGHT] for d in first}
    missed = [d['dataset'] for d in first if not any(within(e, d['ref']) for e in ests[d['key']])]
    res['agreement'] = {'hit_by_all': sum(all(within(e, d['ref']) for e in ests[d['key']]) for d in first),
                        'missed_by_all': len(missed), 'missed_snapme': missed.count('snapme'),
                        'missed_nutrition5k': missed.count('nutrition5k')}

    # Composition of the two sets.
    res['sets'] = {}
    for name, dishes in (('first', first), ('holdout', holdout)):
        refs = [d['ref'] for d in dishes]
        res['sets'][name] = {'meals': len(dishes), 'median_ref': st.median(refs),
                             'under_15_share': sum(r < 15 for r in refs) / len(refs) * 100}
        for ds in ('nutrition5k', 'snapme'):
            part = [d for d in dishes if d['dataset'] == ds]
            res['sets'][name][ds] = {'meals': len(part), 'clusters': len({d['cluster'] for d in part}),
                                     'median_ref': st.median(d['ref'] for d in part),
                                     'mean_ref': st.fmean(d['ref'] for d in part)}

    # Meals quoted in the text: estimates of all eight models.
    res['examples'] = {}
    for label, prefix in EXAMPLES.items():
        d = next(d for d in first if d['dish_id'].startswith(prefix))
        picked = [runs[m][d['key']][0] for m in EIGHT]
        res['examples'][label] = {'dish_id': d['dish_id'], 'ref_carbs': d['ref'], 'ref_mass': d['ref_mass'],
                                  'carbs': {m: s['est'] for m, s in zip(EIGHT, picked)},
                                  'mass': {m: s['mass'] for m, s in zip(EIGHT, picked)}}

    # MAE and bias without the coffee, the largest single error of six models.
    coffee = next(d for d in first if d['dish_id'].startswith(EXAMPLES['coffee']))
    res['without_coffee'] = {}
    for m in EIGHT:
        errors = [runs[m][d['key']][0]['est'] - d['ref'] for d in first
                  if d is not coffee and runs[m][d['key']][0]['est'] is not None]
        res['without_coffee'][m] = {'mae': st.fmean(abs(e) for e in errors), 'bias': st.fmean(errors)}

    # A ninth model on later request code, reported apart from the eight.
    preview, preview_hits = row(first, by_config[('food200', f'{PREVIEW}|low')], resample_first)
    res['other'] = {PREVIEW: preview, 'vs_36': difference(preview_hits, hits[G36], resample_first)}
    return rounded(res)


def rounded(value):
    """Six decimals are enough, and they hide float differences between Python versions."""
    if isinstance(value, dict):
        return {k: rounded(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [rounded(v) for v in value]
    return round(value, 6) if isinstance(value, float) else value


def label(config):
    model, _, reasoning = config.partition('|')
    return f'{NAMES[model]}, {reasoning}' if reasoning else NAMES[model]


def signed(x):
    return f'{x:+.1f}'


def print_rows(title, rows):
    print(f'\n== {title}')
    print(f'{"model":34}{"n":>5}{"within 10 g":>13}{"95% interval":>16}{"MAE":>7}{"bias":>7}{">30 g":>7}{"share":>8}'
          f'{"failed":>8}{"$/scan":>9}{"median s":>10}')
    for config, r in rows.items():
        print(f'{label(config):34}{r["n"]:>5}{r["share"]:>12.1f}%{r["low"]:>9.1f} {r["high"]:>6.1f}{r["mae"]:>7.1f}'
              f'{r["bias"]:>+7.1f}{r["over30"]:>7.1f}{r["over30_share"]:>7.1f}%{r["failed"]:>8}{r["cost"]:>9.4f}{r["seconds"]:>10.1f}')


def print_pairs(title, pairs):
    print(f'\n== {title}')
    for base, others in pairs.items():
        print(f'  minus {label(base)}, points [95% interval]')
        for config, d in others.items():
            print(f'    {label(config):32}{signed(d["diff"]):>7} [{signed(d["low"])}; {signed(d["high"])}]')


def report(res):
    print('Soba model benchmark, October 2026. All numbers are recomputed from the scans CSV.')
    print(f'Bootstrap: {REPS} resamples, seed {SEED}, clusters: SNAPMe participant or Nutrition5k ten-minute interval.')
    print_rows('Table 1. Eight models, first set of 200 meals, one run each', res['first']['rows'])
    print_pairs('Table 2. Paired differences in the share within 10 g, first set', res['first']['vs'])
    print_rows('Table 3a. Same run, 100 Nutrition5k plates (weighed)', res['first_nutrition5k']['rows'])
    print_pairs('Table 3a, differences', res['first_nutrition5k']['vs'])
    print_rows('Table 3b. Same run, 100 SNAPMe photos (food diary)', res['first_snapme']['rows'])
    print_pairs('Table 3b, differences', res['first_snapme']['vs'])
    print('\n== Median absolute error by dataset, first set, grams')
    for m in EIGHT:
        print(f'  {NAMES[m]:32} nutrition5k {res["first_nutrition5k"]["rows"][m]["median_abs"]:>5.1f}  '
              f'snapme {res["first_snapme"]["rows"][m]["median_abs"]:>5.1f}')
    c = res['claude']
    print('\n== Claude Sonnet 5.5, how the run went')
    print(f'  providers {c["providers"]}, requests {c["requests"]}, failed by dataset {c["failed_by_dataset"]}, placeholders {c["placeholders"]}')
    print(f'  one request: {c["one_request"]}')
    print(f'  two requests: {c["two_requests"]}')
    print(f'  hits without placeholders: {c["hits_without_placeholders"]} of 200')
    print_rows('Table 4a. The 159 meals where Claude returned an answer', c['answered']['rows'])
    print_pairs('Table 4a, differences', c['answered']['vs'])
    print_rows('Table 4b. The 154 meals left without placeholder answers', c['real']['rows'])
    print_pairs('Table 4b, differences', c['real']['vs'])
    print_rows('Table 5. Every config on all 400 meals, replicate runs averaged per meal', res['pool']['rows'])
    print_pairs('Table 5, differences', res['pool']['vs'])
    d = res['pool']['medium_37_vs_low_37']
    print(f'  Gemini 3.7 Flash, medium minus low: {signed(d["diff"])} [{signed(d["low"])}; {signed(d["high"])}]')
    print('  share within 5, 10, 15 and 20 g:')
    for config, t in res['pool']['thresholds'].items():
        print(f'    {label(config):32}' + ''.join(f'{t[g]:>7.1f}' for g in ('5', '10', '15', '20')))
    print_rows('Table 6. Holdout of 200 meals, one run per config', res['holdout']['rows'])
    print_pairs('Table 6, differences', res['holdout']['vs'])
    titles = ['<15', '15-30', '30-60', '>=60']
    for scope, title in (('first', 'first set'), ('first_nutrition5k', 'first set, Nutrition5k only'), ('holdout', 'holdout')):
        print(f'\n== Table 7. By reference carbs, {title}: hits/meals, then the errors at 60 g or more')
        for config, cells in res['carb_bands'][scope].items():
            top = cells[-1]
            print(f'  {label(config):34}' + ''.join(f'{t} {x["hits"]}/{x["n"]}  ' for t, x in zip(titles, cells))
                  + f'mean {top["mean_error"]:+.1f}, median {top["median_error"]:+.1f}, too low {top["too_low"]}/{top["answered"]},'
                  + f' range {top["lowest"]:+.1f} to {top["highest"]:+.1f}')
    titles = ['<100 g', '100-200', '200-400', '400-600', '>=600']
    for scope, title in (('first', 'first set'), ('first_nutrition5k', 'first set, Nutrition5k only'),
                         ('holdout', 'holdout'), ('holdout_nutrition5k', 'holdout, Nutrition5k only')):
        print(f'\n== Table 8. Estimated weight / reference weight, median by reference weight, {title}')
        cells = next(iter(res['weight_bands'][scope].values()))
        print('  meals: ' + ', '.join(f'{t} {x["n"]}' for t, x in zip(titles, cells)))
        for config, cells in res['weight_bands'][scope].items():
            print(f'  {label(config):34}' + ''.join(f'{x["median_ratio"]:>7.2f}' if x['median_ratio'] else f'{"-":>7}' for x in cells))
    print('\n== Repeat runs on the first set')
    for m, r in res['repeat'].items():
        print(f'  {NAMES[m]:32} same outcome {r["same_outcome"]:.1f}%, median gap {r["median_gap"]:.1f} g, 90th percentile '
              f'{r["p90_gap"]:.1f} g, largest {r["max_gap"]:.1f} g, second run {r["second_run"]:.1f}%')
    f, a = res['fixed_answer'], res['agreement']
    print('\n== A fixed answer')
    print(f'  {f["grams"]} g for every meal: {f["first"]:.1f}% within 10 g on the first set, {f["holdout"]:.1f}% on the holdout')
    print('\n== Agreement of the eight models on the first set')
    print(f'  all eight within 10 g: {a["hit_by_all"]}; none within 10 g: {a["missed_by_all"]} ({a["missed_snapme"]} SNAPMe, {a["missed_nutrition5k"]} Nutrition5k)')
    print('\n== Sets')
    for s, parts in res['sets'].items():
        print(f'  {s}: {json.dumps(parts, sort_keys=True)}')
    print('\n== Meals quoted in the text, first set. Estimates in the order of Table 1, x means a failed scan')
    for name, e in res['examples'].items():
        carbs = ' '.join('x' if v is None else f'{v:.0f}' for v in e['carbs'].values())
        mass = ' '.join('x' if v is None else f'{v:.0f}' for v in e['mass'].values())
        print(f'  {name} ({e["dish_id"]}): reference {e["ref_carbs"]:.1f} g carbs, {e["ref_mass"]:.0f} g weight; carbs {carbs}; weight {mass}')
    print('\n== MAE and bias without the coffee')
    for m, r in res['without_coffee'].items():
        print(f'  {NAMES[m]:32} MAE {r["mae"]:.1f}, bias {r["bias"]:+.1f}')
    print_rows('Other run, later request code, outside the eight', {PREVIEW: res['other'][PREVIEW]})
    d = res['other']['vs_36']
    print(f'  minus Gemini 3.6 Flash: {signed(d["diff"])} [{signed(d["low"])}; {signed(d["high"])}]')


def report_text(res):
    out = io.StringIO()
    with contextlib.redirect_stdout(out):
        report(res)
    return out.getvalue()


def results_text(res):
    return json.dumps(res, indent=1, sort_keys=True) + '\n'


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('--json', action='store_true', help='print the results as JSON')
    ap.add_argument('--check', action='store_true', help='compare with the expected files')
    args = ap.parse_args()
    res = compute(load())
    if args.check:
        if not (EXPECTED_TEXT.exists() and EXPECTED_JSON.exists()):
            sys.exit('Put expected_output.txt and expected_results.json next to this script.')
        if report_text(res) != EXPECTED_TEXT.read_text() or results_text(res) != EXPECTED_JSON.read_text():
            sys.exit('The output differs from the expected files.')
        print('The output matches the expected files.')
    elif args.json:
        print(results_text(res), end='')
    else:
        report(res)


if __name__ == '__main__':
    main()