"""
Authors: Peter Mawhorter
Date: 2022-10-6
Purpose: A script for clearing outputs in a notebook, to help recover
  from infinite loops that make the notebook unloadable.
"""

import os
import sys
import json

__version__ = "1.0"
"""
Current version number.
"""

print(f"Running clearNotebook version {__version__}")

# Find all .ipynb files in the current directory, and ask which one to
# clean up.
notebooks = []
for f in sorted(os.listdir()):
    if f.endswith(".ipynb"):
        notebooks.append((f, os.path.abspath(f)))


def affirmative(answer):
    "Returns true for affirmative answers which start with 'y'."
    return answer.strip().lower() in ('y', 'yes', 'yeah', 'yup')


if len(notebooks) == 0:
    print("There are no notebook files in the current directory.")
    searchSubs = input(
        "Would you like to search for notebook files in subdirectories"
        " (y/n)? "
    )
    if affirmative(searchSubs):
        print("Okay, search for any notebook files in subdirectories...")
        for dirpath, dirnames, filenames in os.walk('.'):
            for f in filenames:
                if f.endswith('.ipynb'):
                    path = os.path.join(dirpath, f)
                    notebooks.append((path, os.path.abspath(path)))

if len(notebooks) == 0:
    print("We couldn't find any notebooks to clear; aborting.")
    sys.exit(0)

if len(notebooks) == 1:
    print(f"We found one notebook: '{notebooks[0][0]}'.")
    doClear = input("Do you want to clear its outputs (y/n)? ")
    if affirmative(doClear):
        target = notebooks[0]
else:
    print(
        f"We found {len(notebooks)} notebooks... Select one to clear,"
        f" or use 0 to cancel."
    )
    print("  [0] <cancel operation>")
    for i, (name, fullpath) in enumerate(notebooks):
        print(f"  [{i + 1}] '{name}'")
    whichOne = input("Which notebook would you like to clear? ")
    try:
        selected = int(whichOne)
    except ValueError:
        print(
            f"Your response '{whichOne}' could not be interpreted as an"
            f" integer. Aborting."
        )
        sys.exit(0)

    if (selected - 1) not in range(len(notebooks)):
        print(
            f"Your selection ({selected}) did not correspond to any"
            f" of the listed notebooks. Aborting."
        )
        sys.exit(0)

    if selected == 0:
        print("Cancelling operation.")
        sys.exit(0)

    target = notebooks[selected - 1]

print(f"Okay, we'll clear all outputs in the notebook '{target[0]}'.")

print("Loading notebook...")
with open(target[1], 'r') as fileInput:
    raw = fileInput.read()
    parsed = json.loads(raw)

print("Clearing outputs...")
for cell in parsed["cells"]:
    if cell["cell_type"] == "code":
        cell["outputs"] = []

print("Re-encoding notebook...")
encoded = json.dumps(parsed)

if len(encoded) > 10000000:
    print(
        "After clearing outputs, the notebook is still larger than 10"
        " megabytes... You may need to get additional help to clean"
        " things up."
    )

# Find a good name for a backup
backup = target[1] + '.backup'
while os.path.exists(backup):
    if backup.endswith(')'):
        bits = backup.split('(')
        backupn = int(bits[-1][:-1])
        backup = '('.join(bits[:-1]) + f'({backupn + 1})'
    else:
        backup = backup + '(1)'

print(f"Backing up original notebook as '{backup}'...")
with open(backup, 'w') as fileOutput:
    fileOutput.write(raw)

print("Replacing original notebook with cleared version...")
with open(target[1], 'w') as fileOutput:
    fileOutput.write(encoded)

print("...done clearing notebook. You should try to open it again now.")
print(
    f"(If that works, you should probably delete the backup file"
    f" '{backup}' since it may be taking up a lot of space.)"
)