import json
import glob

def fix_notebook(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        nb = json.load(f)
    changed = False
    for cell in nb.get('cells', []):
        if cell.get('cell_type') == 'code':
            source = cell.get('source', [])
            if isinstance(source, list):
                for i, line in enumerate(source):
                    if "dc = datacube.Datacube()" in line:
                        source[i] = "dc = None\n"
                        changed = True
                    if "ds = dc.load(" in line:
                        source[i] = "ds = None\n"
                        changed = True
                    if "data = dc.load(" in line:
                        source[i] = "data = None\n"
                        changed = True
                    # If ds is None, ds.vv will fail
                    if "vv_data = ds.vv" in line:
                        source[i] = "vv_data = None\n"
                        changed = True
                    if "notebook_utils.xarray_object_size(ds)" in line:
                        source[i] = line.replace("notebook_utils.xarray_object_size(ds)", "'ds is None'")
                        changed = True
    if changed:
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(nb, f, indent=1)
        print(f"Removed datacube from {file_path}")

for nb in glob.glob("*.ipynb"):
    fix_notebook(nb)