import json
import glob

def patch_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', [])
            
            # Check if this cell should be fully commented out
            full_source = ''.join(source)
            if 'dc.load(' in full_source or 'ds.vv' in full_source:
                for i in range(len(source)):
                    if not source[i].startswith('#'):
                        source[i] = '# ' + source[i]
                changed = True
                continue
            
            # Otherwise, do line-by-line replacements
            for i, line in enumerate(source):
                if 'ST_training data_updated_1130points.shp' in line:
                    source[i] = line.replace('ST_training data_updated_1130points.shp', 'ST_training_data_updated_1130points.shp')
                    changed = True
                if 'from new_import import *' in line:
                    source[i] = line.replace('from new_import import *', 'from new_import_ODC import *')
                    changed = True
                if 'dc = datacube.Datacube()' in line:
                    source[i] = line.replace('dc = datacube.Datacube()', 'dc = None')
                    changed = True
                if 'load_data(dc,' in line:
                    source[i] = line.replace('load_data(dc,', 'load_data(None,')
                    changed = True
                    
    if changed:
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(nb, f, indent=1)
        print(f"Patched {file_path}")

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