import os
import librosa
import librosa.display
import matplotlib.pyplot as plt
import warnings
import numpy as np
warnings.filterwarnings("ignore", category=RuntimeWarning)

def load_audio(file_path, duration=None, offset=0):
    audio_data, sample_rate = librosa.load(file_path, duration=duration, offset=offset)
    return audio_data, sample_rate

def plot_waveform(audio_data, sample_rate, output_path):
    plt.figure(figsize=(10, 4))
    librosa.display.waveshow(audio_data, sr=sample_rate)
    plt.title("Waveform")
    plt.xlabel("Time (seconds)")
    plt.ylabel("Amplitude")
    plt.savefig(os.path.join(output_path, "waveform.png"))
    plt.show()
    plt.close()

def plot_spectrogram(audio_data, sample_rate, output_path):
    plt.figure(figsize=(10, 4))
    spectrogram = librosa.feature.melspectrogram(y=audio_data, sr=sample_rate)
    librosa.display.specshow(librosa.power_to_db(spectrogram, ref=np.max), y_axis='mel', x_axis='time')
    plt.colorbar(format='%+2.0f dB')
    plt.title("Spectrogram")
    plt.savefig(os.path.join(output_path, "spectrogram.png"))
    plt.show()
    plt.close()

if __name__ == "__main__":
    input_file = "recorded_voice_sounddevice.wav"
    output_path = "."  # You can change this to the desired output directory

    audio_data, sample_rate = load_audio(input_file)

    # Print some information about the audio file
    print(f"Audio data shape: {audio_data.shape}")
    print(f"Sample rate: {sample_rate} Hz")

    # Plot the audio waveform and spectrogram
    plot_waveform(audio_data, sample_rate, output_path)
    plot_spectrogram(audio_data, sample_rate, output_path)

    print("Plots saved as 'waveform.png' and 'spectrogram.png' in the current directory.")