data_generator.py

import numpy as np
from scipy import signal
import random

class MockBrainwaveGenerator:
    """
    Generates realistic mock brainwave data for testing and demonstration.
    """
    
    def __init__(self):
        # Standard EEG channel names (10-20 system)
        self.channel_names = [
            'Fp1', 'Fp2', 'F3', 'F4', 'C3', 'C4', 'P3', 'P4',
            'O1', 'O2', 'F7', 'F8', 'T3', 'T4', 'T5', 'T6',
            'Fz', 'Cz', 'Pz', 'A1', 'A2', 'T1', 'T2', 'X1'
        ]
        
        # Frequency bands and their characteristics
        self.frequency_bands = {
            'delta': {'freq_range': (0.5, 4), 'amplitude': 50, 'prominence': 0.3},
            'theta': {'freq_range': (4, 8), 'amplitude': 30, 'prominence': 0.2},
            'alpha': {'freq_range': (8, 13), 'amplitude': 40, 'prominence': 0.4},
            'beta': {'freq_range': (13, 30), 'amplitude': 20, 'prominence': 0.3},
            'gamma': {'freq_range': (30, 100), 'amplitude': 10, 'prominence': 0.1}
        }
        
        # Brain state characteristics
        self.brain_states = {
            'relaxed': {
                'alpha_boost': 2.0,
                'beta_reduction': 0.5,
                'overall_amplitude': 0.8,
                'noise_level': 0.1
            },
            'focused': {
                'beta_boost': 1.8,
                'gamma_boost': 1.4,
                'alpha_reduction': 0.6,
                'overall_amplitude': 1.2,
                'noise_level': 0.15
            },
            'excited': {
                'beta_boost': 2.2,
                'gamma_boost': 1.8,
                'overall_amplitude': 1.5,
                'noise_level': 0.2
            },
            'drowsy': {
                'delta_boost': 1.8,
                'theta_boost': 1.5,
                'alpha_reduction': 0.4,
                'beta_reduction': 0.3,
                'overall_amplitude': 0.6,
                'noise_level': 0.08
            },
            'meditative': {
                'theta_boost': 2.0,
                'alpha_boost': 1.6,
                'beta_reduction': 0.4,
                'overall_amplitude': 0.9,
                'noise_level': 0.05
            }
        }
    
    def generate_mock_eeg(self, duration=10, sampling_rate=256, num_channels=8, 
                         noise_level=0.1, brain_states=['relaxed']):
        """
        Generate realistic mock EEG data.
        
        Args:
            duration: Duration in seconds
            sampling_rate: Sampling frequency in Hz
            num_channels: Number of EEG channels
            noise_level: Amount of noise to add
            brain_states: List of brain states to simulate
            
        Returns:
            EEG data array (samples x channels)
        """
        num_samples = int(duration * sampling_rate)
        eeg_data = np.zeros((num_samples, num_channels))
        
        # Time vector
        t = np.linspace(0, duration, num_samples)
        
        # Generate data for each channel
        for ch in range(num_channels):
            channel_data = self._generate_channel_data(
                t, sampling_rate, brain_states, ch, num_channels
            )
            
            # Add channel-specific characteristics
            channel_data = self._add_channel_characteristics(
                channel_data, ch, num_channels
            )
            
            # Add artifacts occasionally
            if random.random() < 0.1:  # 10% chance of artifacts
                channel_data = self._add_artifacts(channel_data, t)
            
            eeg_data[:, ch] = channel_data
        
        # Add cross-channel correlation
        eeg_data = self._add_spatial_correlation(eeg_data)
        
        # Add global noise
        noise = np.random.normal(0, noise_level * 10, eeg_data.shape)
        eeg_data += noise
        
        return eeg_data
    
    def _generate_channel_data(self, t, sampling_rate, brain_states, channel_idx, num_channels):
        """
        Generate data for a single channel.
        """
        signal_data = np.zeros_like(t)
        
        # Cycle through different brain states during the recording
        if len(brain_states) == 0:
            brain_states = ['relaxed']  # Default state if none selected
        
        state_duration = len(t) // len(brain_states)
        
        for i, state in enumerate(brain_states):
            start_idx = i * state_duration
            end_idx = min((i + 1) * state_duration, len(t))
            
            if start_idx >= len(t):
                break
            
            t_segment = t[start_idx:end_idx]
            state_signal = self._generate_state_signal(
                t_segment, state, channel_idx, num_channels
            )
            
            signal_data[start_idx:end_idx] = state_signal
        
        return signal_data
    
    def _generate_state_signal(self, t, brain_state, channel_idx, num_channels):
        """
        Generate signal for a specific brain state.
        """
        state_params = self.brain_states.get(brain_state, self.brain_states['relaxed'])
        signal_data = np.zeros_like(t)
        
        # Generate each frequency band
        for band_name, band_info in self.frequency_bands.items():
            freq_min, freq_max = band_info['freq_range']
            base_amplitude = band_info['amplitude']
            prominence = band_info['prominence']
            
            # Apply state-specific modifications
            amplitude_modifier = state_params.get(f'{band_name}_boost', 1.0)
            if f'{band_name}_reduction' in state_params:
                amplitude_modifier = state_params[f'{band_name}_reduction']
            
            amplitude = base_amplitude * amplitude_modifier * prominence
            
            # Generate multiple frequencies within the band
            num_freqs = random.randint(2, 5)
            for _ in range(num_freqs):
                freq = random.uniform(freq_min, freq_max)
                phase = random.uniform(0, 2 * np.pi)
                
                # Add frequency component
                component = amplitude * np.sin(2 * np.pi * freq * t + phase)
                
                # Add some amplitude modulation
                mod_freq = random.uniform(0.1, 2.0)
                mod_amplitude = random.uniform(0.1, 0.3)
                modulation = 1 + mod_amplitude * np.sin(2 * np.pi * mod_freq * t)
                
                signal_data += component * modulation
        
        # Apply overall amplitude scaling
        overall_amplitude = state_params.get('overall_amplitude', 1.0)
        signal_data *= overall_amplitude
        
        return signal_data
    
    def _add_channel_characteristics(self, data, channel_idx, num_channels):
        """
        Add channel-specific characteristics based on brain region.
        """
        # Simulate different brain regions
        if channel_idx < num_channels // 4:  # Frontal
            # Higher beta and gamma activity
            data *= 1.1
        elif channel_idx < num_channels // 2:  # Central
            # Balanced activity
            pass
        elif channel_idx < 3 * num_channels // 4:  # Parietal
            # Higher alpha activity
            alpha_boost = 0.2 * np.sin(2 * np.pi * 10 * np.linspace(0, len(data)/256, len(data)))
            data += alpha_boost * np.mean(np.abs(data))
        else:  # Occipital
            # Strong alpha rhythm
            alpha_rhythm = 0.3 * np.sin(2 * np.pi * 10 * np.linspace(0, len(data)/256, len(data)))
            data += alpha_rhythm * np.mean(np.abs(data))
        
        return data
    
    def _add_artifacts(self, data, t):
        """
        Add realistic EEG artifacts.
        """
        artifact_type = random.choice(['blink', 'muscle', 'movement'])
        
        if artifact_type == 'blink':
            # Eye blink artifacts (low frequency, high amplitude)
            blink_times = random.randint(1, 3)
            for _ in range(blink_times):
                blink_start = random.uniform(0, len(t) - 256)
                blink_idx = int(blink_start)
                blink_duration = random.randint(64, 128)
                
                # Create blink artifact
                blink_signal = 200 * np.exp(-np.linspace(0, 5, blink_duration))
                end_idx = min(blink_idx + blink_duration, len(data))
                actual_duration = end_idx - blink_idx
                
                data[blink_idx:end_idx] += blink_signal[:actual_duration]
        
        elif artifact_type == 'muscle':
            # Muscle artifacts (high frequency)
            muscle_start = random.uniform(0, len(t) - 512)
            muscle_idx = int(muscle_start)
            muscle_duration = random.randint(128, 512)
            
            muscle_signal = np.random.normal(0, 30, muscle_duration)
            # Create bandpass filter for muscle artifacts (50-100 Hz)
            sos = signal.butter(4, [50, 100], btype='band', fs=256, output='sos')
            muscle_signal = signal.sosfilt(sos, muscle_signal)
            
            end_idx = min(muscle_idx + muscle_duration, len(data))
            actual_duration = end_idx - muscle_idx
            
            data[muscle_idx:end_idx] += muscle_signal[:actual_duration]
        
        elif artifact_type == 'movement':
            # Movement artifacts (mixed frequencies)
            movement_start = random.uniform(0, len(t) - 1024)
            movement_idx = int(movement_start)
            movement_duration = random.randint(256, 1024)
            
            movement_signal = np.random.normal(0, 50, movement_duration)
            
            end_idx = min(movement_idx + movement_duration, len(data))
            actual_duration = end_idx - movement_idx
            
            data[movement_idx:end_idx] += movement_signal[:actual_duration]
        
        return data
    
    def _add_spatial_correlation(self, eeg_data):
        """
        Add realistic spatial correlation between channels.
        """
        num_channels = eeg_data.shape[1]
        
        # Create correlation matrix (nearby channels are more correlated)
        correlation_matrix = np.eye(num_channels)
        for i in range(num_channels):
            for j in range(num_channels):
                if i != j:
                    distance = abs(i - j)
                    correlation = max(0, 0.8 - 0.1 * distance)
                    correlation_matrix[i, j] = correlation
        
        # Apply correlation (simplified)
        for i in range(num_channels):
            for j in range(i + 1, num_channels):
                if correlation_matrix[i, j] > 0.3:
                    # Add correlated component
                    correlation_strength = correlation_matrix[i, j]
                    shared_component = correlation_strength * 0.2
                    
                    common_signal = (eeg_data[:, i] + eeg_data[:, j]) / 2
                    eeg_data[:, i] += shared_component * common_signal
                    eeg_data[:, j] += shared_component * common_signal
        
        return eeg_data
    
    def generate_event_related_potentials(self, num_trials=50, sampling_rate=256):
        """
        Generate event-related potentials (ERPs) for different stimulus types.
        """
        trial_duration = 1.0  # 1 second per trial
        num_samples = int(trial_duration * sampling_rate)
        
        # Different ERP components
        erp_components = {
            'P100': {'latency': 0.1, 'amplitude': 5, 'width': 0.02},
            'N170': {'latency': 0.17, 'amplitude': -8, 'width': 0.03},
            'P300': {'latency': 0.3, 'amplitude': 12, 'width': 0.05},
            'N400': {'latency': 0.4, 'amplitude': -6, 'width': 0.04}
        }
        
        trials = []
        
        for trial in range(num_trials):
            t = np.linspace(0, trial_duration, num_samples)
            erp_signal = np.zeros_like(t)
            
            # Add ERP components
            for comp_name, comp_params in erp_components.items():
                latency = comp_params['latency']
                amplitude = comp_params['amplitude']
                width = comp_params['width']
                
                # Gaussian-shaped component
                component = amplitude * np.exp(-0.5 * ((t - latency) / width) ** 2)
                erp_signal += component
            
            # Add noise
            noise = np.random.normal(0, 2, len(erp_signal))
            erp_signal += noise
            
            trials.append(erp_signal)
        
        return np.array(trials)
    
    def generate_sleep_stages(self, duration=3600, sampling_rate=256):
        """
        Generate EEG data with different sleep stages.
        """
        num_samples = int(duration * sampling_rate)
        t = np.linspace(0, duration, num_samples)
        
        # Define sleep stages and their characteristics
        sleep_stages = {
            'wake': {'duration': 0.1, 'alpha': 2.0, 'beta': 1.5},
            'n1': {'duration': 0.05, 'theta': 2.0, 'alpha': 0.5},
            'n2': {'duration': 0.4, 'theta': 1.5, 'sleep_spindles': True, 'k_complexes': True},
            'n3': {'duration': 0.2, 'delta': 3.0},
            'rem': {'duration': 0.25, 'theta': 2.0, 'beta': 1.2, 'saw_tooth': True}
        }
        
        eeg_data = np.zeros_like(t)
        
        # Cycle through sleep stages
        stage_sequence = ['wake', 'n1', 'n2', 'n3', 'n2', 'rem'] * (duration // 5400 + 1)
        
        current_time = 0
        for stage in stage_sequence:
            if current_time >= duration:
                break
            
            stage_duration = sleep_stages[stage]['duration'] * 5400  # Stage duration in seconds
            stage_end = min(current_time + stage_duration, duration)
            
            stage_start_idx = int(current_time * sampling_rate)
            stage_end_idx = int(stage_end * sampling_rate)
            
            if stage_start_idx >= len(t):
                break
            
            t_stage = t[stage_start_idx:stage_end_idx]
            stage_signal = self._generate_sleep_stage_signal(t_stage, stage, sleep_stages[stage])
            
            eeg_data[stage_start_idx:stage_end_idx] = stage_signal
            current_time = stage_end
        
        return eeg_data
    
    def _generate_sleep_stage_signal(self, t, stage_name, stage_params):
        """
        Generate EEG signal for a specific sleep stage.
        """
        signal_data = np.zeros_like(t)
        
        # Base frequencies for different bands
        if 'delta' in stage_params:
            delta_amp = stage_params['delta'] * 40
            for freq in np.arange(0.5, 4, 0.5):
                signal_data += delta_amp * np.sin(2 * np.pi * freq * t + random.uniform(0, 2*np.pi))
        
        if 'theta' in stage_params:
            theta_amp = stage_params['theta'] * 30
            for freq in np.arange(4, 8, 0.5):
                signal_data += theta_amp * np.sin(2 * np.pi * freq * t + random.uniform(0, 2*np.pi))
        
        if 'alpha' in stage_params:
            alpha_amp = stage_params['alpha'] * 35
            signal_data += alpha_amp * np.sin(2 * np.pi * 10 * t)
        
        if 'beta' in stage_params:
            beta_amp = stage_params['beta'] * 20
            for freq in np.arange(13, 30, 2):
                signal_data += beta_amp * np.sin(2 * np.pi * freq * t + random.uniform(0, 2*np.pi))
        
        # Add stage-specific features
        if stage_params.get('sleep_spindles'):
            # Add sleep spindles (12-14 Hz bursts)
            spindle_times = random.randint(1, 5)
            for _ in range(spindle_times):
                spindle_start = random.uniform(0, len(t) - 128)
                spindle_idx = int(spindle_start)
                spindle_duration = random.randint(64, 128)
                
                spindle_freq = random.uniform(12, 14)
                spindle_env = np.exp(-np.linspace(0, 3, spindle_duration))
                spindle_signal = 50 * spindle_env * np.sin(2 * np.pi * spindle_freq * 
                                                          np.linspace(0, spindle_duration/256, spindle_duration))
                
                end_idx = min(spindle_idx + spindle_duration, len(signal_data))
                signal_data[spindle_idx:end_idx] += spindle_signal[:end_idx-spindle_idx]
        
        return signal_data