import numpy as np
import matplotlib.pyplot as plt
import librosa
import librosa.display

# Load the audio file
audio_file = "./sounds/s4.mp3"
audio, sr = librosa.load(audio_file, sr=None)

# Echo parameters
delay = int(0.2 * sr)  # Echo delay in samples
decay1 = 0.8  # Echo decay factor 1
decay2 = 0.3  # Echo decay factor 2

# Apply echo effect with decay factor 1
echoed_audio1 = np.zeros_like(audio)
echoed_audio1[delay:] = audio[:-delay] + decay1 * audio[delay:]

# Apply echo effect with decay factor 2
echoed_audio2 = np.zeros_like(audio)
echoed_audio2[delay:] = audio[:-delay] + decay2 * audio[delay:]

# Compute the Mel spectrogram of the original audio
mel_spec_orig = librosa.feature.melspectrogram(y=audio, sr=sr)

# Compute the Mel spectrogram of the echoed audio with decay factor 1
mel_spec_echoed1 = librosa.feature.melspectrogram(y=echoed_audio1, sr=sr)

# Compute the Mel spectrogram of the echoed audio with decay factor 2
mel_spec_echoed2 = librosa.feature.melspectrogram(y=echoed_audio2, sr=sr)

# Convert to dB scale
mel_spec_orig_db = librosa.power_to_db(mel_spec_orig, ref=np.max)
mel_spec_echoed1_db = librosa.power_to_db(S=mel_spec_echoed1, ref=np.max)
mel_spec_echoed2_db = librosa.power_to_db(S=mel_spec_echoed2, ref=np.max)

# Display the Mel spectrograms
plt.figure(figsize=(10, 4))

plt.subplot(1, 2, 1)
librosa.display.specshow(mel_spec_orig_db, sr=sr, x_axis='time', y_axis='mel')
plt.colorbar(format='%+2.0f dB')
plt.title('Original Mel Spectrogram')

plt.subplot(1, 2, 2)
librosa.display.specshow(mel_spec_echoed1_db, sr=sr, x_axis='time', y_axis='mel')
plt.colorbar(format='%+2.0f dB')
plt.title('Echoed Mel Spectrogram (Decay 1)')

# plt.subplot(1, 3, 3)
# librosa.display.specshow(mel_spec_echoed2_db, sr=sr, x_axis='time', y_axis='mel')
# plt.colorbar(format='%+2.0f dB')
# plt.title('Echoed Mel Spectrogram (Decay 2)')

plt.tight_layout()
plt.show()