from machine import Pin, I2C
from ssd1306 import SSD1306_I2C
import time
import math
import random

# Setup I2C and OLED
i2c = I2C(0, scl=Pin(1), sda=Pin(0), freq=400000)
oled = SSD1306_I2C(128, 64, i2c)

# Variables
ecg_mV = 0.0
heart_rate = 75
last_beat = 0
waveform = [2048] * 128
idx = 0

print("=== ECG SIMULATION ===")
print("mV -> Amplifier(x1000) -> ADC -> HR")

while True:
    t = time.ticks_ms() % 800
    
    # Generate ECG in mV (0-3mV)
    if t < 200:
        x = (t - 100) / 40.0
        ecg_mV = 2.5 * math.exp(-(x * x))
    else:
        ecg_mV = 0.1 * math.sin(t / 50.0) + 0.05
    
    ecg_mV += random.uniform(-0.01, 0.01)
    
    # Amplify (Gain = 1000) and convert to ADC
    adc = int((ecg_mV * 1000 / 3300.0) * 4095)
    adc = max(0, min(4095, adc))
    
    # Beat detection
    now = time.ticks_ms()
    if adc > 2480 and time.ticks_diff(now, last_beat) > 300:
        if last_beat != 0:
            interval = time.ticks_diff(now, last_beat)
            heart_rate = int(60000 / interval)
        last_beat = now
        print("BEAT!")
    
    # Store waveform
    waveform[idx] = adc
    idx = (idx + 1) % 128
    
    # Draw
    oled.fill(0)
    
    for i in range(1, 128):
        y1 = int(55 - (waveform[(idx + i - 1) % 128] / 4095) * 40)
        y2 = int(55 - (waveform[(idx + i) % 128] / 4095) * 40)
        y1 = max(10, min(55, y1))
        y2 = max(10, min(55, y2))
        oled.line(i-1, y1, i, y2, 1)
    
    oled.text(f"HR:{heart_rate} BPM", 0, 56)
    oled.text(f"{ecg_mV:.2f}mV", 75, 56)
    oled.show()
    
    print(f"ECG:{ecg_mV:.2f}mV | Amp:{ecg_mV*1000:.0f}mV | ADC:{adc} | HR:{heart_rate}")
    
    time.sleep(0.02)