import sys
import numpy as np
import pyaudio
import matplotlib.pyplot as plt

RECORD_SECONDS = 5
CHUNK = 1024
RATE = 44100

p = pyaudio.PyAudio()
stream = p.open(format=p.get_format_from_width(2),
                channels=1 if sys.platform == 'darwin' else 2,
                rate=RATE,
                input=True,
                output=True,
                frames_per_buffer=CHUNK)

print('* recording')

# Initialize plot
fig, ax = plt.subplots()
x = np.arange(0, RECORD_SECONDS, CHUNK / RATE)
line, = ax.plot(x, np.zeros(len(x)))

def add_echo(data, output_stream):
    output_stream.write(data)

# Initialize data arrays
db_data = np.zeros(len(x))


for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
    byte_stream = stream.read(CHUNK)
    data = np.frombuffer(byte_stream, dtype=np.int16)
    fft_data = np.fft.fft(data)
    # 주파수 대역 설정
    freq_range = (200, 2000)  # 200Hz ~ 2kHz

    # 주파수 스펙트럼에서 주파수 대역 추출
    freq_spectrum = fft_data[(freq_range[0] // (RATE // len(data))) : (freq_range[1] // (RATE // len(data)))]

    # 에너지 계산
    energy = np.sum(np.abs(freq_spectrum)**2)
    # 소리의 세기 계산
    db = 10 * np.log10(energy)

    # 소리의 높낮이 계산
    max_freq = (np.argmax(np.abs(freq_spectrum)) * RATE) / len(data)
    print("freq : ", max_freq)
    print("db : ", db)

    # Add data to arrays
    db_data = np.roll(db_data, -1)
    db_data[-1] = db

    # Update plot
    line.set_ydata(db_data)
    ax.relim()
    ax.autoscale_view()

    plt.draw()
    plt.pause(0.001)

print('* done')

stream.close()
p.terminate()

# Show plot
plt.show()