import json
import glob
import re

def patch_python_script(filepath):
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
            
        original_content = content
        
        # Replace imports
        content = re.sub(r'from sklearn\.ensemble import RandomForestClassifier', 
                         'from xgboost import XGBClassifier', content)
        
        # Replace the model instantiations (for 01.train_ODC.py)
        rf_pattern = re.compile(r'model\s*=\s*RandomForestClassifier\([^)]+\)', re.DOTALL)
        xgb_replacement = """model = XGBClassifier(
                    n_estimators=200,
                    max_depth=30,
                    tree_method="hist",
                    device="cuda",
                    random_state=42,
                    n_jobs=-1,
                    verbosity=1
                )"""
        
        content = rf_pattern.sub(xgb_replacement, content)
        
        if content != original_content:
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(content)
            print(f"Patched {filepath}")
    except Exception as e:
        print(f"Error patching {filepath}: {e}")

def patch_notebook(filepath):
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            nb = json.load(f)
            
        changed = False
        import_pattern = re.compile(r'from\s+sklearn\.ensemble\s+import\s+RandomForestClassifier')
        inst_pattern = re.compile(r'RandomForestClassifier\([^)]*\)')
        
        for cell in nb.get('cells', []):
            if cell.get('cell_type') == 'code':
                source = cell.get('source', [])
                for i in range(len(source)):
                    if import_pattern.search(source[i]):
                        source[i] = import_pattern.sub('from xgboost import XGBClassifier', source[i])
                        changed = True
                        
                    if inst_pattern.search(source[i]):
                        source[i] = inst_pattern.sub("XGBClassifier(tree_method='hist', device='cuda', random_state=42, n_jobs=-1)", source[i])
                        changed = True
                        
                    if "'classifier__criterion': ['gini', 'entropy']" in source[i]:
                        source[i] = source[i].replace("'classifier__criterion': ['gini', 'entropy']",
                                                      "'classifier__learning_rate': [0.01, 0.1, 0.2]")
                        changed = True
                        
        if changed:
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(nb, f, indent=1)
            print(f"Patched {filepath}")
    except Exception as e:
        print(f"Error patching {filepath}: {e}")

if __name__ == "__main__":
    patch_python_script("01.train_ODC.py")
    patch_python_script("new_train.py")
    
    for nb in glob.glob("*.ipynb"):
        patch_notebook(nb)