import json
import glob
import os
import sys

# Đảm bảo import được new_import_ODC
sys.path.insert(0, os.getcwd())
import new_import_ODC
from new_import_ODC import load_data, load_sen1

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

def extract_params(nb_file):
    params = {}
    try:
        with open(nb_file, 'r', encoding='utf-8') as f:
            nb = json.load(f)
            
        for cell in nb.get('cells', []):
            if cell.get('cell_type') == 'code':
                source = cell.get('source', [])
                if isinstance(source, list):
                    source_code = "".join(source)
                else:
                    source_code = source
                    
                # Phân tích các dòng
                for line in source_code.split('\n'):
                    line = line.strip()
                    if line.startswith('date_range = '):
                        # Lấy giá trị của date_range
                        try:
                            val = eval(line.split('=', 1)[1].strip())
                            params['date_range'] = val
                        except: pass
                    elif line.startswith('longtitude_range = '):
                        try:
                            val = eval(line.split('=', 1)[1].strip())
                            params['longtitude_range'] = val
                        except: pass
                    elif line.startswith('latitude_range = '):
                        try:
                            val = eval(line.split('=', 1)[1].strip())
                            params['latitude_range'] = val
                        except: pass
                    elif line.startswith('time_range = '):
                        try:
                            val = eval(line.split('=', 1)[1].strip())
                            params['time_range'] = val
                        except: pass
    except Exception as e:
        print(f"Error reading {nb_file}: {e}")
        
    return params

print("Starting to cache data for all notebooks...")
for nb_file in NOTEBOOKS_TO_RUN:
    if os.path.exists(nb_file):
        params = extract_params(nb_file)
        if 'date_range' in params and 'longtitude_range' in params and 'latitude_range' in params:
            date_range = params['date_range']
            lon_range = params['longtitude_range']
            lat_range = params['latitude_range']
            
            print(f"\n--- Caching for {nb_file} ---")
            print(f"Date: {date_range}, Lon: {lon_range}, Lat: {lat_range}")
            
            # Caching Sentinel-2
            print("Loading Sentinel-2 (load_data)...")
            try:
                load_data(None, date_range, lon_range, lat_range)
            except Exception as e:
                print(f"Failed Sentinel-2: {e}")
                
            # Caching Sentinel-1
            time_range = f"{date_range[0]}/{date_range[1]}"
            bbox = [lon_range[0], lat_range[0], lon_range[1], lat_range[1]]
            print("Loading Sentinel-1 (load_sen1)...")
            try:
                load_sen1(bbox, time_range)
            except Exception as e:
                print(f"Failed Sentinel-1: {e}")
        else:
            print(f"\nSkipped {nb_file}: Could not find all parameters.")

print("\nDone caching all data!")