import machine
import ustruct
import ulab as np

# FFT Configuration
SAMPLE_RATE = 44100  # 44.1 kHz
BUFFER_SIZE = 1024   # Number of samples per FFT
FREQ_RES = SAMPLE_RATE / BUFFER_SIZE  # Frequency resolution of the FFT

# I2S Configuration
i2s_input = machine.I2S(
    id=0,
    sck=machine.Pin(14),  # BCLK
    ws=machine.Pin(15),   # WS
    sd=machine.Pin(32),   # Data
    mode=machine.I2S.RX,  # Receive
    bits=16,              # 16-bit audio
    format=machine.I2S.MONO,  # Mono
    rate=SAMPLE_RATE,      # 44.1 kHz
    ibuf=1024              # Buffer
)

# Note Frequency Mapping
NOTES = {
    440: 'A4', 494: 'B4', 523: 'C5', 587: 'D5', 
    659: 'E5', 698: 'F5', 784: 'G5', 880: 'A5'
}

def freq_to_note(freq):
    """Convert frequency to the closest musical note."""
    closest_note = min(NOTES, key=lambda x: abs(x - freq))
    return NOTES[closest_note]

# Buffer to read I2S audio data
audio_buffer = bytearray(BUFFER_SIZE * 2)  # 16 bits = 2 bytes per sample

try:
    while True:
        # 1. Read from I2S (read raw audio samples)
        num_bytes_read = i2s_input.readinto(audio_buffer)
        
        # 2. Convert audio to NumPy array (16-bit samples)
        samples = np.array(ustruct.unpack('<' + 'h' * (num_bytes_read // 2), audio_buffer))
        
        # 3. Apply FFT to the audio data
        fft_result = np.fft.fft(samples)
        magnitudes = np.abs(fft_result)

        # 4. Find the dominant frequency
        max_index = np.argmax(magnitudes)
        dominant_frequency = max_index * FREQ_RES

        # 5. Convert frequency to note
        note = freq_to_note(dominant_frequency)
        
        print(f"Dominant Frequency: {dominant_frequency:.2f} Hz - Note: {note}")

except KeyboardInterrupt:
    i2s_input.deinit()