import wave
import numpy as np
from scipy.signal import butter, lfilter

FIFO_PATH = "/tmp/esp32_audio"
RAW_FILE = "recorded_audio.wav"
FILTERED_FILE = "recorded_audio_filtered.wav"

SAMPLE_RATE = 16000
CHANNELS = 2
BYTES_PER_SAMPLE = 2

# --- DSP helper: smoother band-pass filter ---
def butter_bandpass(lowcut, highcut, fs, order=4):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype="band")
    return b, a

# Keep filter state across chunks (avoids clicks)
b, a = butter_bandpass(120, 6000, SAMPLE_RATE, order=2)
zi = [np.zeros(max(len(a), len(b)) - 1) for _ in range(CHANNELS)]

def bandpass_filter(data, channel):
    global zi
    filtered, zi[channel] = lfilter(b, a, data, zi=zi[channel])
    return filtered

# --- Configure WAV writers ---
raw_wav = wave.open(RAW_FILE, "wb")
raw_wav.setnchannels(CHANNELS)
raw_wav.setsampwidth(BYTES_PER_SAMPLE)
raw_wav.setframerate(SAMPLE_RATE)

filtered_wav = wave.open(FILTERED_FILE, "wb")
filtered_wav.setnchannels(CHANNELS)
filtered_wav.setsampwidth(BYTES_PER_SAMPLE)
filtered_wav.setframerate(SAMPLE_RATE)

with open(FIFO_PATH, "rb") as f:
    print(f"Recording from {FIFO_PATH} into {RAW_FILE} and {FILTERED_FILE}...")
    try:
        while True:
            data = f.read(8000)  # ~0.125s stereo
            if not data:
                continue

            # --- Write raw ---
            raw_wav.writeframes(data)

            # --- Convert to numpy ---
            audio_np = np.frombuffer(data, dtype=np.int16).reshape(-1, CHANNELS)

            # --- Filter each channel with persistent state ---
            filtered_channels = []
            for ch in range(CHANNELS):
                filtered_channels.append(bandpass_filter(audio_np[:, ch], ch))
            filtered_np = np.stack(filtered_channels, axis=1).astype(np.int16)

            GAIN = 6.0  # 2x louder

            # --- Apply gain after filtering ---
            filtered_np = np.stack(filtered_channels, axis=1)
            
            # Scale and clip to int16 range
            filtered_np = np.clip(filtered_np * GAIN, -32768, 32767).astype(np.int16)
            # --- Write filtered ---
            filtered_wav.writeframes(filtered_np.tobytes())

    except KeyboardInterrupt:
        print("\nStopping recording.")
    finally:
        raw_wav.close()
        filtered_wav.close()
        print(f"Saved {RAW_FILE} and {FILTERED_FILE}")
