import wave
import numpy as np
import soundfile as sf
import noisereduce as nr
from scipy.signal import butter, lfilter

FIFO_PATH = "/tmp/esp32_audio"
RAW_FILE = "recorded_audio.wav"
BANDPASS_FILE = "recorded_audio_bandpass.wav"
NOISE_REDUCED_FILE = "recorded_audio_noise_reduced.wav"

SAMPLE_RATE = 16000
CHANNELS = 2
BYTES_PER_SAMPLE = 2

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

b, a = butter_bandpass(80, 7000, 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

# --- Step 1: Capture ~2s of noise (runtime only) ---
print("Sampling background noise...")
with open(FIFO_PATH, "rb") as f:
    noise_data = f.read(SAMPLE_RATE * BYTES_PER_SAMPLE * CHANNELS * 2)  # 2 seconds
noise_np = np.frombuffer(noise_data, dtype=np.int16).reshape(-1, CHANNELS)
noise_mono = noise_np.mean(axis=1)  # downmix to mono

# --- 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)

bandpass_wav = wave.open(BANDPASS_FILE, "wb")
bandpass_wav.setnchannels(CHANNELS)
bandpass_wav.setsampwidth(BYTES_PER_SAMPLE)
bandpass_wav.setframerate(SAMPLE_RATE)

# --- Step 2: Record raw + band-pass in parallel ---
with open(FIFO_PATH, "rb") as f:
    print(f"Recording from {FIFO_PATH} into {RAW_FILE} and {BANDPASS_FILE}...")
    try:
        while True:
            data = f.read(8192)  # aligned block size
            if not data:
                continue

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

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

            # --- Band-pass filter each channel ---
            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)

            # --- Write band-pass filtered ---
            bandpass_wav.writeframes(filtered_np.tobytes())

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

# --- Step 3: Post-process noise reduction ---
print("Running noise reduction on raw file...")
raw_audio, sr = sf.read(RAW_FILE, dtype="float32")  # ensure float32 in [-1,1]

# Downmix both to mono
raw_mono = raw_audio.mean(axis=1)
noise_mono = noise_np.astype(np.float32).mean(axis=1) / 32768.0  # scale int16 -> [-1,1]

print("Noise RMS:", np.sqrt(np.mean(noise_mono**2)))
print("Raw RMS:", np.sqrt(np.mean(raw_mono**2)))

# Apply noise reduction
filtered_mono = nr.reduce_noise(y=raw_mono,
                                sr=sr,
                                y_noise=noise_mono,
                                prop_decrease=0.7)

# Mirror back to stereo
filtered_stereo = np.column_stack((filtered_mono, filtered_mono))
sf.write(NOISE_REDUCED_FILE, filtered_stereo, sr)
print(f"Saved {NOISE_REDUCED_FILE}")
