#!/usr/bin/env python3
import numpy as np
import sounddevice as sd
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

# ========== USER SETTINGS ==========
samplerate = 48000       # sample rate of your system
blocksize = 2048         # larger block -> better frequency resolution
fft_size = 2048          # number of samples per FFT
n_lines = 300            # number of lines in the waterfall (vertical scrolling)
freq_max = 10000          # get the full 20khz audio range
device_name = None       # set to your monitor source if you know it# "None" if you dont (without quotes)
delay = 0
# Example: device_name = "alsa_output.pci-0000_00_1b.0.analog-stereo.monitor"
# Leave as None to be prompted

# ===================================

print("\n=== Available audio devices ===")
print(sd.query_devices())

if device_name is None:
    device_name = input("\nEnter the name or number of your Monitor device: ").strip()

window = np.hanning(fft_size)
freqs = np.fft.rfftfreq(fft_size, d=1/samplerate)
freq_bins = freqs <= freq_max

waterfall = np.zeros((n_lines, np.sum(freq_bins)))
fig, ax = plt.subplots(figsize=(10, 6))
img = ax.imshow(waterfall, origin='lower', aspect='auto', cmap='inferno',
                extent=[0, freq_max, 0, n_lines],
                norm=LogNorm(vmin=1e-6, vmax=1))
ax.set_xlabel("Frequency (Hz)")
ax.set_ylabel("Time (scrolling up)")
ax.set_title("Hydrophone Waterfall Spectrogram (0–1 kHz)")
plt.tight_layout()

def audio_callback(indata, frames, time, status):
    global delay
    global waterfall
    if status:
        print(status)
    data = indata[:, 0] * window[:frames]
    spectrum = np.abs(np.fft.rfft(data))**2
    spectrum = spectrum[freq_bins]
    waterfall = np.roll(waterfall, -1, axis=0)
    waterfall[-1, :] = spectrum
    if delay > 40:
    	img.set_data(waterfall)
    	fig.canvas.draw_idle()
    	delay = 0
    else:
    	delay = delay + 1
    #plt.pause(0.001)

with sd.InputStream(device=device_name, channels=1,
                    samplerate=samplerate, blocksize=blocksize,
                    callback=audio_callback):
    print("\n✅ Listening to output via:", device_name)
    print("Close the window or Ctrl+C to stop.\n")
    plt.show()