import json
import glob
import subprocess
import time
import os

NOTEBOOKS_TO_RUN = [
    "01.train_ODC.ipynb",
    "01.train_ODC_XGBoost.ipynb",
    "02.predict_ODC.ipynb",
    "new_train.ipynb"
]

def limit_time_range(file_path):
    try:
        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):
                        # Replace 2023-12-31 with 2023-04-01
                        if '"2023-12-31"' in line:
                            source[i] = line.replace('"2023-12-31"', '"2023-04-01"')
                            changed = True
                        if "'2023-10-01'" in line:
                            source[i] = line.replace("'2023-10-01'", "'2022-10-01'")
                            changed = True
                        if '"2023-10-01"' in line:
                            source[i] = line.replace('"2023-10-01"', '"2022-10-01"')
                            changed = True
                        # For time_range="2022-09-01/2023-10-01"
                        if "2022-09-01/2023-10-01" in line:
                            source[i] = line.replace("2022-09-01/2023-10-01", "2022-09-01/2022-10-01")
                            changed = True
                            
                elif isinstance(source, str):
                    new_source = source.replace('"2023-12-31"', '"2023-04-01"')
                    new_source = new_source.replace("'2023-10-01'", "'2022-10-01'")
                    new_source = new_source.replace('"2023-10-01"', '"2022-10-01"')
                    new_source = new_source.replace("2022-09-01/2023-10-01", "2022-09-01/2022-10-01")
                    if new_source != source:
                        cell['source'] = new_source
                        changed = True
        
        if changed:
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(nb, f, indent=1)
            print(f"Limited time_range to 1 month in {file_path}")
    except Exception as e:
        print(f"Error on {file_path}: {e}")

# 1. Modify the time ranges
for nb_file in glob.glob("*.ipynb"):
    limit_time_range(nb_file)

# 2. Run them in parallel
print("\nStarting parallel execution of notebooks...")
processes = []
for nb_file in NOTEBOOKS_TO_RUN:
    if os.path.exists(nb_file):
        print(f"Launching {nb_file}...")
        cmd = f"source /home/x79/miniconda3/etc/profile.d/conda.sh && conda activate env_01 && jupyter nbconvert --execute --ExecutePreprocessor.timeout=-1 --inplace {nb_file}"
        p = subprocess.Popen(["bash", "-c", cmd], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        processes.append((nb_file, p))

# 3. Wait and print output
for nb_file, p in processes:
    p.wait()
    output = p.stdout.read().decode('utf-8')
    if p.returncode == 0:
        print(f"[{nb_file}] SUCCESS")
    else:
        print(f"[{nb_file}] FAILED (code {p.returncode})")
        print(f"--- OUTPUT START ({nb_file}) ---")
        print(output)
        print(f"--- OUTPUT END ({nb_file}) ---")

print("\nAll tasks finished.")