test.py

import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import random

class ECGADCGenerators:
    def __init__(self, adc_bits=10, baseline_adc=510, fs=250):
        """
        ECG ADC value generators with various patterns including ST depression
        
        Parameters:
        - adc_bits: ADC resolution (10-bit = 0-1023)
        - baseline_adc: Baseline ADC value (from your data: ~510)
        - fs: Sampling frequency in Hz
        """
        self.adc_bits = adc_bits
        self.adc_max = (2 ** adc_bits) - 1  # 1023 for 10-bit
        self.baseline_adc = baseline_adc
        self.fs = fs
        
    def add_noise(self, signal, noise_level=2):
        """Add realistic ADC noise to signal"""
        noise = np.random.normal(0, noise_level, len(signal))
        return np.clip(signal + noise, 0, self.adc_max).astype(int)
    
    def normal_ecg_adc(self, duration=10, hr=70):
        """Generate normal ECG in ADC values"""
        samples = int(duration * self.fs)
        t = np.linspace(0, duration, samples)
        
        # Heart rate parameters
        rr_interval = 60 / hr  # seconds between beats
        
        ecg_signal = np.zeros(samples)
        
        # Generate each heartbeat
        beat_times = np.arange(0, duration, rr_interval)
        
        for beat_time in beat_times:
            beat_start = int(beat_time * self.fs)
            if beat_start + int(0.8 * self.fs) < samples:
                beat = self._generate_normal_beat()
                end_idx = min(beat_start + len(beat), samples)
                ecg_signal[beat_start:end_idx] += beat[:end_idx-beat_start]
        
        # Convert to ADC values
        ecg_signal = ecg_signal * 80 + self.baseline_adc  # Scale and offset
        ecg_signal = self.add_noise(ecg_signal)
        
        return ecg_signal
    
    def st_depression_ecg_adc(self, duration=10, hr=70, depression_mv=0.5):
        """Generate ECG with ST depression in ADC values"""
        samples = int(duration * self.fs)
        t = np.linspace(0, duration, samples)
        
        # Heart rate parameters
        rr_interval = 60 / hr
        
        ecg_signal = np.zeros(samples)
        
        # Generate each heartbeat with ST depression
        beat_times = np.arange(0, duration, rr_interval)
        
        for beat_time in beat_times:
            beat_start = int(beat_time * self.fs)
            if beat_start + int(0.8 * self.fs) < samples:
                beat = self._generate_st_depression_beat(depression_mv)
                end_idx = min(beat_start + len(beat), samples)
                ecg_signal[beat_start:end_idx] += beat[:end_idx-beat_start]
        
        # Convert to ADC values
        ecg_signal = ecg_signal * 80 + self.baseline_adc
        ecg_signal = self.add_noise(ecg_signal)
        
        return ecg_signal
    
    def intermittent_st_depression_adc(self, duration=10, hr=70, depression_probability=0.3):
        """Generate ECG with intermittent ST depression"""
        samples = int(duration * self.fs)
        rr_interval = 60 / hr
        
        ecg_signal = np.zeros(samples)
        beat_times = np.arange(0, duration, rr_interval)
        
        for beat_time in beat_times:
            beat_start = int(beat_time * self.fs)
            if beat_start + int(0.8 * self.fs) < samples:
                # Randomly decide if this beat has ST depression
                if random.random() < depression_probability:
                    beat = self._generate_st_depression_beat(0.3 + random.random() * 0.4)
                else:
                    beat = self._generate_normal_beat()
                
                end_idx = min(beat_start + len(beat), samples)
                ecg_signal[beat_start:end_idx] += beat[:end_idx-beat_start]
        
        ecg_signal = ecg_signal * 80 + self.baseline_adc
        ecg_signal = self.add_noise(ecg_signal)
        
        return ecg_signal
    
    def progressive_st_depression_adc(self, duration=10, hr=70):
        """Generate ECG with progressively worsening ST depression"""
        samples = int(duration * self.fs)
        rr_interval = 60 / hr
        
        ecg_signal = np.zeros(samples)
        beat_times = np.arange(0, duration, rr_interval)
        
        for i, beat_time in enumerate(beat_times):
            beat_start = int(beat_time * self.fs)
            if beat_start + int(0.8 * self.fs) < samples:
                # Progressive worsening of ST depression
                depression_level = min(0.8, i * 0.1)  # Increase depression over time
                
                if depression_level > 0:
                    beat = self._generate_st_depression_beat(depression_level)
                else:
                    beat = self._generate_normal_beat()
                
                end_idx = min(beat_start + len(beat), samples)
                ecg_signal[beat_start:end_idx] += beat[:end_idx-beat_start]
        
        ecg_signal = ecg_signal * 80 + self.baseline_adc
        ecg_signal = self.add_noise(ecg_signal)
        
        return ecg_signal
    
    def st_elevation_ecg_adc(self, duration=10, hr=70, elevation_mv=0.5):
        """Generate ECG with ST elevation in ADC values"""
        samples = int(duration * self.fs)
        t = np.linspace(0, duration, samples)
        
        # Heart rate parameters
        rr_interval = 60 / hr
        
        ecg_signal = np.zeros(samples)
        
        # Generate each heartbeat with ST elevation
        beat_times = np.arange(0, duration, rr_interval)
        
        for beat_time in beat_times:
            beat_start = int(beat_time * self.fs)
            if beat_start + int(0.8 * self.fs) < samples:
                beat = self._generate_st_elevation_beat(elevation_mv)
                end_idx = min(beat_start + len(beat), samples)
                ecg_signal[beat_start:end_idx] += beat[:end_idx-beat_start]
        
        # Convert to ADC values
        ecg_signal = ecg_signal * 80 + self.baseline_adc
        ecg_signal = self.add_noise(ecg_signal)
        
        return ecg_signal
    
    def arrhythmia_with_st_depression_adc(self, duration=10, base_hr=70):
        """Generate ECG with irregular rhythm and ST depression"""
        samples = int(duration * self.fs)
        ecg_signal = np.zeros(samples)
        
        current_time = 0
        while current_time < duration:
            # Irregular RR intervals
            rr_variation = random.uniform(0.7, 1.3)  # ±30% variation
            rr_interval = (60 / base_hr) * rr_variation
            
            beat_start = int(current_time * self.fs)
            if beat_start + int(0.8 * self.fs) < samples:
                # 40% chance of ST depression
                if random.random() < 0.4:
                    beat = self._generate_st_depression_beat(0.2 + random.random() * 0.5)
                else:
                    beat = self._generate_normal_beat()
                
                end_idx = min(beat_start + len(beat), samples)
                ecg_signal[beat_start:end_idx] += beat[:end_idx-beat_start]
            
            current_time += rr_interval
        
        ecg_signal = ecg_signal * 80 + self.baseline_adc
        ecg_signal = self.add_noise(ecg_signal, noise_level=3)  # More noise for arrhythmia
        
        return ecg_signal
    
    def atrial_fibrillation_adc(self, duration=10, base_hr=90):
        """Generate atrial fibrillation ECG - no P waves, irregular rhythm"""
        samples = int(duration * self.fs)
        ecg_signal = np.zeros(samples)
        
        current_time = 0
        while current_time < duration:
            # Highly irregular RR intervals (0.3-1.2s)
            rr_interval = random.uniform(0.3, 1.2)
            
            beat_start = int(current_time * self.fs)
            if beat_start + int(0.6 * self.fs) < samples:
                # No P wave beats with varying QRS morphology
                beat = self._generate_afib_beat()
                end_idx = min(beat_start + len(beat), samples)
                ecg_signal[beat_start:end_idx] += beat[:end_idx-beat_start]
            
            current_time += rr_interval
        
        # Add baseline fibrillation waves
        ecg_signal += self._generate_fibrillation_baseline(samples)
        
        ecg_signal = ecg_signal * 80 + self.baseline_adc
        ecg_signal = self.add_noise(ecg_signal, noise_level=4)  # More noise for AFib
        
        return ecg_signal
    
    def ventricular_tachycardia_adc(self, duration=10, hr=180):
        """Generate ventricular tachycardia ECG - wide QRS, no P waves"""
        samples = int(duration * self.fs)
        t = np.linspace(0, duration, samples)
        
        # Fast, regular rhythm
        rr_interval = 60 / hr
        
        ecg_signal = np.zeros(samples)
        
        # Generate each VT beat
        beat_times = np.arange(0, duration, rr_interval)
        
        for beat_time in beat_times:
            beat_start = int(beat_time * self.fs)
            if beat_start + int(0.4 * self.fs) < samples:
                beat = self._generate_vt_beat()
                end_idx = min(beat_start + len(beat), samples)
                ecg_signal[beat_start:end_idx] += beat[:end_idx-beat_start]
        
        ecg_signal = ecg_signal * 80 + self.baseline_adc
        ecg_signal = self.add_noise(ecg_signal)
        
        return ecg_signal
    
    def junctional_rhythm_adc(self, duration=10, hr=50):
        """Generate junctional rhythm ECG - absent/inverted P waves"""
        samples = int(duration * self.fs)
        t = np.linspace(0, duration, samples)
        
        # Slow, regular rhythm
        rr_interval = 60 / hr
        
        ecg_signal = np.zeros(samples)
        
        # Generate each junctional beat
        beat_times = np.arange(0, duration, rr_interval)
        
        for beat_time in beat_times:
            beat_start = int(beat_time * self.fs)
            if beat_start + int(0.8 * self.fs) < samples:
                beat = self._generate_junctional_beat()
                end_idx = min(beat_start + len(beat), samples)
                ecg_signal[beat_start:end_idx] += beat[:end_idx-beat_start]
        
        ecg_signal = ecg_signal * 80 + self.baseline_adc
        ecg_signal = self.add_noise(ecg_signal)
        
        return ecg_signal
    
    def _generate_normal_beat(self):
        """Generate a single normal heartbeat"""
        beat_samples = int(0.8 * self.fs)  # 0.8 seconds per beat
        t = np.linspace(0, 0.8, beat_samples)
        
        # P wave (0.08-0.12s)
        p_wave = 0.2 * np.exp(-((t - 0.1) / 0.03)**2)
        
        # QRS complex (0.06-0.10s)
        q_wave = -0.1 * np.exp(-((t - 0.25) / 0.01)**2)
        r_wave = 1.0 * np.exp(-((t - 0.27) / 0.015)**2)
        s_wave = -0.2 * np.exp(-((t - 0.29) / 0.01)**2)
        
        # T wave (0.15-0.25s)
        t_wave = 0.3 * np.exp(-((t - 0.45) / 0.08)**2)
        
        beat = p_wave + q_wave + r_wave + s_wave + t_wave
        return beat
    
    def _generate_st_depression_beat(self, depression_mv=0.5):
        """Generate a heartbeat with ST depression"""
        beat_samples = int(0.8 * self.fs)
        t = np.linspace(0, 0.8, beat_samples)
        
        # Normal components
        p_wave = 0.2 * np.exp(-((t - 0.1) / 0.03)**2)
        q_wave = -0.1 * np.exp(-((t - 0.25) / 0.01)**2)
        r_wave = 1.0 * np.exp(-((t - 0.27) / 0.015)**2)
        s_wave = -0.2 * np.exp(-((t - 0.29) / 0.01)**2)
        
        # ST segment depression (0.08s after QRS)
        st_start = 0.32
        st_end = 0.42
        st_depression = np.zeros_like(t)
        st_mask = (t >= st_start) & (t <= st_end)
        st_depression[st_mask] = -depression_mv
        
        # T wave (may be inverted with severe depression)
        t_wave_amplitude = 0.3 if depression_mv < 0.4 else -0.2
        t_wave = t_wave_amplitude * np.exp(-((t - 0.45) / 0.08)**2)
        
        beat = p_wave + q_wave + r_wave + s_wave + st_depression + t_wave
        return beat
    
    def _generate_st_elevation_beat(self, elevation_mv=0.5):
        """Generate a heartbeat with ST elevation"""
        beat_samples = int(0.8 * self.fs)
        t = np.linspace(0, 0.8, beat_samples)
        
        # Normal components
        p_wave = 0.2 * np.exp(-((t - 0.1) / 0.03)**2)
        q_wave = -0.1 * np.exp(-((t - 0.25) / 0.01)**2)
        r_wave = 1.0 * np.exp(-((t - 0.27) / 0.015)**2)
        s_wave = -0.2 * np.exp(-((t - 0.29) / 0.01)**2)
        
        # ST segment elevation (0.08s after QRS)
        st_start = 0.32
        st_end = 0.42
        st_elevation = np.zeros_like(t)
        st_mask = (t >= st_start) & (t <= st_end)
        st_elevation[st_mask] = elevation_mv
        
        # T wave (often tall and peaked with ST elevation)
        t_wave_amplitude = 0.3 + (elevation_mv * 0.5)  # Taller T wave
        t_wave = t_wave_amplitude * np.exp(-((t - 0.45) / 0.08)**2)
        
        beat = p_wave + q_wave + r_wave + s_wave + st_elevation + t_wave
        return beat
    
    def _generate_afib_beat(self):
        """Generate atrial fibrillation beat - no P wave, variable QRS"""
        beat_samples = int(0.6 * self.fs)  # Shorter beats
        t = np.linspace(0, 0.6, beat_samples)
        
        # Variable QRS morphology
        qrs_variation = random.uniform(0.8, 1.2)
        q_wave = -0.1 * np.exp(-((t - 0.15) / 0.01)**2) * qrs_variation
        r_wave = 1.0 * np.exp(-((t - 0.17) / 0.015)**2) * qrs_variation
        s_wave = -0.2 * np.exp(-((t - 0.19) / 0.01)**2) * qrs_variation
        
        # T wave
        t_wave = 0.3 * np.exp(-((t - 0.35) / 0.08)**2)
        
        beat = q_wave + r_wave + s_wave + t_wave
        return beat
    
    def _generate_vt_beat(self):
        """Generate ventricular tachycardia beat - wide QRS, no P wave"""
        beat_samples = int(0.4 * self.fs)  # Shorter cycle
        t = np.linspace(0, 0.4, beat_samples)
        
        # Wide, bizarre QRS complex
        qrs_complex = np.zeros_like(t)
        
        # Multi-phasic wide QRS
        qrs_complex += -0.3 * np.exp(-((t - 0.08) / 0.02)**2)  # Wide Q
        qrs_complex += 1.5 * np.exp(-((t - 0.12) / 0.025)**2)  # Tall R
        qrs_complex += -0.4 * np.exp(-((t - 0.16) / 0.02)**2)  # Deep S
        qrs_complex += 0.2 * np.exp(-((t - 0.20) / 0.02)**2)   # Late R'
        
        # Abnormal T wave
        t_wave = -0.4 * np.exp(-((t - 0.28) / 0.05)**2)  # Inverted T
        
        beat = qrs_complex + t_wave
        return beat
    
    def _generate_junctional_beat(self):
        """Generate junctional rhythm beat - absent/inverted P wave"""
        beat_samples = int(0.8 * self.fs)
        t = np.linspace(0, 0.8, beat_samples)
        
        # Absent or inverted P wave (20% chance of inverted)
        if random.random() < 0.2:
            p_wave = -0.1 * np.exp(-((t - 0.1) / 0.03)**2)  # Inverted P
        else:
            p_wave = np.zeros_like(t)  # No P wave
        
        # Normal narrow QRS
        q_wave = -0.1 * np.exp(-((t - 0.25) / 0.01)**2)
        r_wave = 0.8 * np.exp(-((t - 0.27) / 0.015)**2)  # Slightly smaller R
        s_wave = -0.15 * np.exp(-((t - 0.29) / 0.01)**2)
        
        # T wave
        t_wave = 0.25 * np.exp(-((t - 0.45) / 0.08)**2)
        
        beat = p_wave + q_wave + r_wave + s_wave + t_wave
        return beat
    
    def _generate_fibrillation_baseline(self, samples):
        """Generate fibrillation baseline waves for AFib"""
        # Random fibrillation waves (350-600 per minute)
        fib_freq = random.uniform(350, 600) / 60  # Convert to Hz
        t = np.linspace(0, samples/self.fs, samples)
        
        # Multiple frequency components for chaotic appearance
        fib_waves = np.zeros(samples)
        for i in range(5):
            freq = fib_freq * (0.8 + 0.4 * random.random())
            amplitude = 0.05 * random.random()
            phase = 2 * np.pi * random.random()
            fib_waves += amplitude * np.sin(2 * np.pi * freq * t + phase)
        
        return fib_waves
    
    def plot_comparison(self, duration=5):
        """Plot comparison of different ECG patterns"""
        fig, axes = plt.subplots(3, 3, figsize=(18, 12))
        
        # Generate different patterns
        normal = self.normal_ecg_adc(duration)
        st_dep = self.st_depression_ecg_adc(duration, depression_mv=0.4)
        st_elev = self.st_elevation_ecg_adc(duration, elevation_mv=0.3)
        intermittent = self.intermittent_st_depression_adc(duration)
        progressive = self.progressive_st_depression_adc(duration)
        afib = self.atrial_fibrillation_adc(duration)
        vt = self.ventricular_tachycardia_adc(duration)
        junctional = self.junctional_rhythm_adc(duration)
        arrhythmia = self.arrhythmia_with_st_depression_adc(duration)
        
        t = np.linspace(0, duration, len(normal))
        
        patterns = [
            (normal, 'Normal ECG'),
            (st_dep, 'ST Depression'),
            (st_elev, 'ST Elevation'),
            (intermittent, 'Intermittent ST Depression'),
            (progressive, 'Progressive ST Depression'),
            (afib, 'Atrial Fibrillation (No P waves)'),
            (vt, 'Ventricular Tachycardia'),
            (junctional, 'Junctional Rhythm'),
            (arrhythmia, 'Arrhythmia + ST Depression')
        ]
        
        for i, (pattern, title) in enumerate(patterns):
            row = i // 3
            col = i % 3
            
            t_pattern = np.linspace(0, duration, len(pattern))
            axes[row, col].plot(t_pattern, pattern)
            axes[row, col].set_title(title)
            axes[row, col].set_ylabel('ADC Value')
            axes[row, col].grid(True)
            
            if row == 2:  # Bottom row
                axes[row, col].set_xlabel('Time (s)')
        
        plt.tight_layout()
        plt.show()

# Example usage
if __name__ == "__main__":
    # Initialize generator with your parameters
    ecg_gen = ECGADCGenerators(baseline_adc=510, fs=250)
    
    # Generate 10 seconds of different ECG patterns
    print("Generating different ECG patterns...")
    
    # Normal ECG
    normal_ecg = ecg_gen.normal_ecg_adc(duration=10, hr=70)
    print(f"Normal ECG: {len(normal_ecg)} samples, range: {normal_ecg.min()}-{normal_ecg.max()}")
    
    # ST Depression ECG
    st_dep_ecg = ecg_gen.st_depression_ecg_adc(duration=10, hr=70, depression_mv=0.5)
    print(f"ST Depression ECG: {len(st_dep_ecg)} samples, range: {st_dep_ecg.min()}-{st_dep_ecg.max()}")
    
    # Intermittent ST Depression
    intermittent_ecg = ecg_gen.intermittent_st_depression_adc(duration=10, hr=70)
    print(f"Intermittent ST Depression: {len(intermittent_ecg)} samples")
    
    # Progressive ST Depression
    progressive_ecg = ecg_gen.progressive_st_depression_adc(duration=10, hr=70)
    print(f"Progressive ST Depression: {len(progressive_ecg)} samples")
    
    # Arrhythmia with ST Depression
    arrhythmia_ecg = ecg_gen.arrhythmia_with_st_depression_adc(duration=10, base_hr=70)
    print(f"Arrhythmia with ST Depression: {len(arrhythmia_ecg)} samples")
    
    # Plot comparison
    ecg_gen.plot_comparison(duration=5)
    
    # Example: Get a specific pattern for testing
    test_data = ecg_gen.st_depression_ecg_adc(duration=30, hr=75, depression_mv=0.3)
    print(f"\nTest data generated: {len(test_data)} ADC values")
    print(f"Sample values: {test_data[:10]}")