main.py

import numpy as np
from preprocess import butter_bandpass_filter, apply_zscore_normalization
from entropy import extract_2d_entropy_map
from train import train_model
from tqdm import tqdm

def main():
    print("=== MECNN Framework Pipeline ===")
    
    # 1. Configuration (e.g., BCI Comp IV 2a settings)
    n_trials = 100
    n_channels = 22
    n_times = 500  # 2.0s epoch at 250 Hz
    tau_max = 24
    num_classes = 4
    
    print("1. Generating simulated EEG data...")
    # Synthetic data generation: (Trials, Channels, Timepoints)
    raw_eeg = np.random.randn(n_trials, n_channels, n_times)
    labels = np.random.randint(0, num_classes, n_trials)
    
    # 2. Preprocessing
    print("2. Applying 4-40Hz Butterworth Bandpass Filter...")
    filtered_eeg = butter_bandpass_filter(raw_eeg, fs=250.0)
    
    # 3. 2D Spatial-Scale Entropy Extraction
    print("3. Extracting 2D RCMDE Spatial-Scale Maps (Multicore)...")
    feature_maps = np.zeros((n_trials, n_channels, tau_max))
    for i in tqdm(range(n_trials), desc="Extracting"):
        feature_maps[i] = extract_2d_entropy_map(filtered_eeg[i], tau_max=tau_max)
        
    # Reshape for CNN input: (Batch, 1, Channels, Scales)
    feature_maps = np.expand_dims(feature_maps, axis=1)
    
    # 4. Train/Val Split (80/20)
    split_idx = int(n_trials * 0.8)
    X_train, y_train = feature_maps[:split_idx], labels[:split_idx]
    X_val, y_val = feature_maps[split_idx:], labels[split_idx:]
    
    # 5. Normalization
    print("4. Applying Z-score Normalization...")
    X_train, X_val = apply_zscore_normalization(X_train, X_val)
    
    # 6. Model Training
    print("5. Initiating MECNN Training...")
    trained_model = train_model(X_train, y_train, X_val, y_val, 
                                num_channels=n_channels, num_scales=tau_max, num_classes=num_classes)
                                
    print("=== Pipeline Completed Successfully! ===")

if __name__ == "__main__":
    main()