#!/usr/bin/env python

## SOURCE ==> https://github.com/fastai/fastai_dev/blob/master/dev_nb/notebook2script.py
## with minor modifications

# use tools/sync-nb-exports to autogenerate all relevant files automatically

import json, re, os.path
from pathlib import Path

def is_export(cell):
    if cell['cell_type'] != 'code': return False
    src = cell['source']
    if len(src) == 0 or len(src[0]) < 7: return False
    #import pdb; pdb.set_trace()
    return re.match(r'^\s*#\s*export\s*$', src[0], re.IGNORECASE) is not None

def notebook2script(fname):
    fname = Path(fname)
    fname_out = fname.parent/f"{fname.stem}.py"
    print(fname_out)
    main_dic = json.load(open(fname,'r'))
    cells = main_dic['cells']
    code_cells = [c for c in cells if is_export(c)]
    module = f'''
        #################################################
        ### THIS FILE WAS AUTOGENERATED! DO NOT EDIT! ###
        #################################################
        # file to edit: {fname.name}\n\n'''
    for cell in code_cells: module += ''.join(cell['source'][1:]) + '\n\n'
    # remove trailing spaces
    module = re.sub(r' +$', '', module, flags=re.MULTILINE)
    with open(fname_out,'w') as f: f.write(module[:-2])
    print(f"Converted {fname} to {fname_out}")

import argparse
desc = '''
This program converts a Jupyter notebook with '#export' at the beginning
of each cell into a script with same name as the notebook
'''
parser = argparse.ArgumentParser(usage=desc, 
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('fname', type=str, metavar='NB_FNAME',
                    help='Path to the notebook to be converted')
args = parser.parse_args()

notebook2script(args.fname)