import pyaudio
import numpy as np
import scipy.signal as signal
import matplotlib.pyplot as plt
# 파라미터 설정
RATE = 44100  # 샘플링 주파수
CHUNK = 1024  # 읽을 샘플의 수
THRESHOLD = 256  # 피크를 검출하기 위한 threshold 값
WIN_SIZE = 1024  # STFT를 적용할 윈도우 사이즈
HOP_SIZE = 512  # STFT에서 윈도우 사이의 거리 (오버랩 사이즈)

# PyAudio 객체 생성
p = pyaudio.PyAudio()

# 콜백 함수 정의
def process_audio(in_data, frame_count, time_info, status):
    # 오디오 데이터 변환
    data = np.frombuffer(in_data, dtype=np.int16)

    # STFT 수행
    f, t, Zxx = signal.stft(data, RATE, nperseg=WIN_SIZE, noverlap=HOP_SIZE)

    # 피크 검출
    peaks, _ = signal.find_peaks(np.abs(np.mean(Zxx, axis=1)), height=THRESHOLD, distance=WIN_SIZE)
    # 파라미터 추정
    if len(peaks) > 0:
        peak_idx = peaks[0]  # 첫 번째 피크 선택
        height = np.abs(Zxx[peak_idx, 0])  # 피크의 높이 추정
        freq = f[peak_idx]  # 피크의 주파수 추정
        amp = np.max(np.abs(data))  # 신호의 진폭 추정
        progress = (peak_idx + HOP_SIZE) / RATE  # 충돌음의 진행 길이 추정

        # 결과 출력
        print("Height: {:.2f}, Frequency: {:.2f}, Amplitude: {:.2f}, Progress: {:.2f}".format(height, freq, amp, progress))

    # 반환할 데이터 없음
    return (in_data, pyaudio.paContinue)

# 입력 스트림 열기
stream = p.open(format=p.get_format_from_width(2),
                channels=1,
                rate=RATE,
                input=True,
                output=True,
                frames_per_buffer=CHUNK,
                stream_callback=process_audio
                )

# 스트림 시작
stream.start_stream()

# 프로그램 실행 중지 전까지 무한 대기
while stream.is_active():
    pass

# 스트림과 PyAudio 객체 종료
stream.stop_stream()
stream.close()
p.terminate()