#!/usr/bin/env python

import sys
import json

def renumber_nb(nb_file_name):
    assert nb_file_name.endswith('.ipynb')
    base = '.'.join(nb_file_name.split('.')[:-1])
    i = 1
    j = json.load(open(nb_file_name, 'r'))
    for cell in j['cells']:
        if cell['cell_type'] == 'code':
            cell['execution_count'] = i
            for o in cell['outputs']:
                if 'data' in o and o['output_type'] == 'execute_result':
                    o['execution_count'] = i
            i += 1
    outfile = '{}-numbered.ipynb'.format(base)
    json.dump(j, open(outfile, 'w'))
    return outfile

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print('Usage: {} <jupyter_notebook.ipynb>'.format(sys.argv[0]))
        sys.exit(1)
    outfile = renumber_nb(sys.argv[1])
    print('Wrote {}'.format(outfile))