import argparse
from traitlets.config import Config
import nbformat as nbf
from nbconvert.preprocessors import TagRemovePreprocessor, ClearOutputPreprocessor
from nbconvert.exporters import NotebookExporter


def get_arg_parser():
    parser = argparse.ArgumentParser()

    parser.add_argument('input_file')
    parser.add_argument('output_file')

    return parser


def convert(input_file, output_file):
    c = Config()
    c.TagRemovePreprocessor.remove_cell_tags = ("solution",)
    c.TagRemovePreprocessor.enabled = True
    c.ClearOutputPreprocesser.enabled = True
    c.NotebookExporter.preprocessors = [
        "nbconvert.preprocessors.TagRemovePreprocessor",
        "nbconvert.preprocessors.ClearOutputPreprocessor"
    ]

    exporter = NotebookExporter(config=c)
    exporter.register_preprocessor(TagRemovePreprocessor(config=c), True)
    exporter.register_preprocessor(ClearOutputPreprocessor(), True)

    output = NotebookExporter(config=c).from_filename(input_file)
    with open(output_file, 'w') as f:
        f.write(output[0])


if __name__ == "__main__":
    parser = get_arg_parser()
    args = parser.parse_args()

    convert(args.input_file, args.output_file)
    print(f'Converted {args.input_file} to {args.output_file}')
