#!/usr/bin/env python

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

# fmt: off

import json, re
from pathlib import Path
from upyog.cli import Param as P
from upyog.cli import call_parse

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

@call_parse
def notebook2script(fname: P("Path to the notebook to be converted")):
    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'''
        # Autogenerated from {fname}
        # 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}")

# fmt: on