eeg_preprocessing_functions.py

# -*- coding: utf-8 -*-
"""EEG Preprocessing Functions

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1waeYVfxOhj9iqUa0PAfdlw5uh5sjMJTV
"""

import numpy as np
import mne
from mne.preprocessing import ICA
from scipy.stats import zscore
import pywt # <-- Added PyWavelets

# --- Configuration Constants (Modify as needed) ---
# Assuming your ANT Neuro system configuration:
SFREQ = 256.0           # Sampling rate in Hz (common for BCI2020/ANT systems)
N_CHANNELS = 64         # Number of EEG channels
LINE_FREQ = 50          # Line noise frequency (50 Hz for India/Europe, 60 Hz for US)
LOW_FREQ = 0.5          # High-pass filter cutoff (removes DC drift)
HIGH_FREQ = 100         # Low-pass filter cutoff (retains high gamma, removes excessive noise)
T_MIN = -0.5            # Start time of epoch relative to event (baseline period in seconds)
T_MAX = 3.5             # End time of epoch (for a total 4.0s epoch length)

# --- PHASE 3: SLIDING WINDOW CONFIGURATION ---
WINDOW_DURATION = 1.0   # Duration of the analysis window in seconds (1.0s recommended)
STRIDE_DURATION = 0.25  # How much the window shifts in seconds (0.25s recommended for 75% overlap)
# -----------------------------------------------

# --- PHASE 4: DWT CONFIGURATION ---
WAVELET_NAME = 'db4'    # Daubechies 4 wavelet recommended for EEG
DWT_LEVEL = 5           # Level 5 decomposition to separate Theta, Alpha, Beta, Gamma
# The relevant bands for features will be A5, D5, D4, D3, D2 (5 bands, 4 features each)
# ----------------------------------


def create_mne_info(sfreq, n_channels):
    """Creates the MNE Info structure required to handle raw EEG data."""
    # Create channel names (assuming standard 10-20 layout for 64 channels)
    # Note: MNE automatically loads standard channel locations for 10-20 system if
    # using standard names (e.g., Fz, Cz, Pz). We use generic names here for safety.
    ch_names = [f'EEG {i:02d}' for i in range(1, n_channels + 1)]

    # Create channel types array (all are EEG)
    ch_types = ['eeg'] * n_channels

    # Create the info structure
    info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
    return info

def apply_temporal_filtering(raw):
    """Applies Bandpass and Notch filtering to the MNE Raw object."""
    print("--- Applying Temporal Filtering ---")

    # 1. Bandpass Filtering (0.5 Hz to 100 Hz)
    raw.filter(l_freq=LOW_FREQ, h_freq=HIGH_FREQ, fir_design='firwin', verbose=False)
    print(f"   -> Bandpass Filter: {LOW_FREQ} Hz to {HIGH_FREQ} Hz applied.")

    # 2. Notch Filtering (Line Noise Removal)
    raw.notch_filter(freqs=LINE_FREQ, filter_length='auto', phase='zero', verbose=False)
    print(f"   -> Notch Filter: {LINE_FREQ} Hz line noise removed.")

    # 3. Reference (Crucial step often performed before ICA/epoching)
    # Assumes common average reference (CAR) unless specific reference (e.g., Mastoids) is known
    raw.set_eeg_reference('average', verbose=False)
    print("   -> Common Average Reference (CAR) applied.")

    return raw

def correct_artifacts_ica(raw):
    """Corrects EOG and EMG artifacts using Independent Component Analysis (ICA)."""
    print("\n--- Artifact Correction (ICA) ---")

    # 1. Fit ICA model
    ica = ICA(n_components=N_CHANNELS - 1, random_state=42, method='fastica')
    ica.fit(raw)
    print(f"   -> ICA fitted with {ica.n_components} components.")

    # --- Automated Artifact Detection (Needs manual verification in a real project) ---
    try:
        eog_indices, eog_scores = ica.find_bads_eog(raw, ch_name='EEG 01', threshold=3.0)
    except ValueError:
        eog_indices = []

    emg_indices, emg_scores = ica.find_bads_muscle(raw, threshold=0.1)

    bad_components = set(eog_indices) | set(emg_indices)

    if bad_components:
        ica.exclude = list(bad_components)
        print(f"   -> Detected and excluded {len(ica.exclude)} components (EOG/EMG).")

        # 2. Apply inverse ICA transformation to remove artifacts
        raw_clean = ica.apply(raw.copy())
        print("   -> ICA applied. Data cleaned.")
    else:
        raw_clean = raw.copy()
        print("   -> No clear artifact components detected. Returning original filtered data.")

    return raw_clean

def segment_and_epoch(raw_clean, marker_array, unique_words):
    """Segments the continuous cleaned data into fixed-length trials (epochs)."""
    print("\n--- Segmentation and Epoching ---")

    # 1. Create MNE Events Array
    events = []
    event_indices = np.where(marker_array != 0)[0]

    for idx in event_indices:
        event_id = int(marker_array[idx])
        events.append([idx, 0, event_id])

    events = np.array(events)

    # 2. Define Event IDs and Epochs
    event_id_dict = {word: i+1 for i, word in enumerate(unique_words)}

    epochs = mne.Epochs(
        raw_clean,
        events=events,
        event_id=event_id_dict,
        tmin=T_MIN,
        tmax=T_MAX,
        baseline=(-0.5, 0), # Baseline correct the data
        preload=True,        # Load data into memory immediately
        verbose=False
    )

    # Final data shape: (N_events, N_channels, N_time_points)
    eeg_data_3d = epochs.get_data()
    eeg_labels_1d = epochs.events[:, 2] # The third column contains the Event ID (label)

    print(f"   -> Identified {len(events)} total events for {len(unique_words)} words.")
    print(f"   -> Epoching data from {T_MIN}s to {T_MAX}s ({T_MAX - T_MIN}s trials).")

    return eeg_data_3d, eeg_labels_1d, event_id_dict

def apply_sliding_window(eeg_data_3d, eeg_labels_1d):
    """
    Applies the sliding window technique to augment the segmented EEG data.
    """
    print("\n--- PHASE 3: Data Augmentation (Sliding Window) ---")

    # Calculate sizes in samples
    window_samples = int(WINDOW_DURATION * SFREQ)  # 1.0s * 256 Hz = 256 samples
    stride_samples = int(STRIDE_DURATION * SFREQ)  # 0.25s * 256 Hz = 64 samples

    if window_samples > eeg_data_3d.shape[2]:
        print(f"ERROR: Window size ({window_samples}) is larger than epoch length ({eeg_data_3d.shape[2]}). Skipping augmentation.")
        return eeg_data_3d, eeg_labels_1d

    augmented_data = []
    augmented_labels = []

    total_trials = eeg_data_3d.shape[0]

    for trial_index in range(total_trials):
        trial_data = eeg_data_3d[trial_index]
        original_label = eeg_labels_1d[trial_index]

        for start_time in range(0,
                                 trial_data.shape[1] - window_samples + 1,
                                 stride_samples):
            end_time = start_time + window_samples
            segment = trial_data[:, start_time:end_time]

            augmented_data.append(segment)
            augmented_labels.append(original_label)

    augmented_data = np.array(augmented_data)
    augmented_labels = np.array(augmented_labels)

    print(f"   -> Original Samples: {total_trials}")
    print(f"   -> Augmented Samples: {augmented_data.shape[0]}")
    print(f"   -> Data Multiplication Factor: {augmented_data.shape[0] / total_trials:.1f}x")
    print(f"   -> New Sample Shape: (64 Channels, {window_samples} Samples)")

    return augmented_data, augmented_labels

def extract_dwt_features(augmented_data):
    """
    Performs DWT decomposition and extracts statistical features (Energy, Power, Std Dev, Entropy)
    from the relevant EEG frequency bands for every channel and segment.

    Args:
        augmented_data (np.ndarray): Augmented EEG data (Augmented Samples, Channels, TimePoints).

    Returns:
        np.ndarray: The final feature matrix (Augmented Samples, 1024 Features).
    """
    print("\n--- PHASE 4: DWT Feature Extraction and Assembly ---")

    def shannon_entropy(signal):
        """Calculate Shannon entropy of a signal."""
        # Convert to probability distribution
        signal_normalized = signal - np.min(signal)
        if np.sum(signal_normalized) > 0:
            prob_dist = signal_normalized / np.sum(signal_normalized)
            # Remove zeros to avoid log(0)
            prob_dist = prob_dist[prob_dist > 0]
            # Calculate entropy
            entropy = -np.sum(prob_dist * np.log2(prob_dist))
            return entropy
        else:
            return 0.0

    # Relevant bands based on Level 5 decomposition and 256 Hz SFREQ:
    # A5 (Delta: 0-4 Hz), D5 (Theta: 4-8 Hz), D4 (Alpha: 8-16 Hz),
    # D3 (Beta: 16-32 Hz), D2 (Gamma: 32-64 Hz)

    feature_vector_list = []
    n_samples, n_channels, _ = augmented_data.shape

    # Estimate the final feature dimension based on configuration:
    n_bands = 5  # A5, D5, D4, D3, D2
    n_features_per_band = 4 # Energy, Power, Std Dev, Entropy
    expected_dim = n_channels * n_bands * n_features_per_band

    print(f"   -> Target feature dimension: {expected_dim} per sample.")

    for i, segment in enumerate(augmented_data):
        # segment shape: (64 Channels, 256 Samples)
        feature_vector = []

        for channel_data in segment: # Iterate through all 64 channels

            # 1. Perform DWT Decomposition
            # Returns coefficients list: [A5, D5, D4, D3, D2, D1]
            coeffs = pywt.wavedec(channel_data, WAVELET_NAME, level=DWT_LEVEL)

            # Select relevant bands (A5, D5, D4, D3, D2) and exclude D1 (high-frequency noise)
            relevant_bands = [coeffs[0], coeffs[1], coeffs[2], coeffs[3], coeffs[4]]

            for band_coeff in relevant_bands:

                # --- Calculate the 4 Core Statistical Features (The "Report Card") ---

                # Ensure coefficients are not empty
                if band_coeff.size == 0:
                    features = [0.0] * n_features_per_band
                else:
                    # 1. Energy (E) - Sum of the squared coefficients
                    energy = np.sum(band_coeff**2)

                    # 2. Mean Power (P) - Mean of the squared coefficients
                    power = np.mean(band_coeff**2)

                    # 3. Standard Deviation (σ)
                    std_dev = np.std(band_coeff)

                    # 4. Shannon Entropy (H) - Measures complexity/disorder
                    entropy = shannon_entropy(band_coeff)

                    features = [energy, power, std_dev, entropy]

                # Append the 4 features for this band/channel to the vector
                feature_vector.extend(features)

        # Add the final 1024-dimension vector for this segment to the list
        feature_vector_list.append(feature_vector)

    final_feature_matrix = np.array(feature_vector_list)
    print(f"   -> Final Feature Matrix Shape: {final_feature_matrix.shape}")

    return final_feature_matrix