import numpy as np
import matplotlib.pyplot as plt
import sounddevice as sd
from scipy.fftpack import fft
import time
import os

# Configuration
SAMPLE_RATE = 44100  # Hz
DURATION = 5  # seconds
CHUNK_SIZE = 1024  # samples
PLOT_REFRESH_RATE = 0.1  # seconds

# Create a figure with two subplots (waveform and spectrum)
plt.ion()  # Interactive mode on
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))
fig.suptitle('Live Microphone Waveform & Spectrum')

# Initialize empty plots
line_wave, = ax1.plot(np.arange(CHUNK_SIZE), np.zeros(CHUNK_SIZE))
line_spectrum, = ax2.plot(np.arange(CHUNK_SIZE//2), np.zeros(CHUNK_SIZE//2))

# Axis settings
ax1.set_ylim(-1, 1)
ax1.set_xlim(0, CHUNK_SIZE)
ax1.set_title('Waveform')
ax1.set_ylabel('Amplitude')

ax2.set_xlim(0, CHUNK_SIZE//2)
ax2.set_ylim(0, 1)
ax2.set_title('Spectrum')
ax2.set_ylabel('Magnitude')
ax2.set_xlabel('Frequency Bin')

# Audio callback function
def audio_callback(indata, frames, time, status):
    if status:
        print(status, flush=True)
    
    # Update waveform plot
    line_wave.set_ydata(indata[:, 0])
    
    # Compute and update spectrum
    N = len(indata[:, 0])
    yf = fft(indata[:, 0])
    xf = np.linspace(0, SAMPLE_RATE//2, N//2)
    line_spectrum.set_ydata(2/N * np.abs(yf[:N//2]))
    
    # Redraw the plots
    fig.canvas.draw()
    fig.canvas.flush_events()

# Start streaming
stream = sd.InputStream(
    callback=audio_callback,
    samplerate=SAMPLE_RATE,
    channels=1,
    blocksize=CHUNK_SIZE
)

try:
    print("Starting audio capture... Press Ctrl+C to stop.")
    with stream:
        while True:
            time.sleep(PLOT_REFRESH_RATE)
except KeyboardInterrupt:
    print("\nStopping audio capture...")
    
    # Save the final plot as an image
    timestamp = time.strftime("%Y%m%d-%H%M%S")
    filename = f"waveform_{timestamp}.png"
    fig.savefig(filename)
    print(f"Saved waveform image as {filename}")
    
    plt.close()
except Exception as e:
    print(f"Error: {e}")
    plt.close()