# Imports
import streamlit as st
import os
import numpy as np
import librosa
import joblib
from pydub import AudioSegment

# Function to predict audio anomalies using an Isolation Forest model
def predict_anomalies(sample_audio_file, model_filename='isolation_forest_model_advanced.pkl'):
    """
    Loads an Isolation Forest model and predicts anomalies in an audio file.
    
    Parameters:
    - sample_audio_file: Path to the audio file for anomaly detection.
    - model_filename: Filename of the saved Isolation Forest model.
    
    Returns:
    - A prediction label (0 for normal, 1 for anomaly).
    """
    clf = joblib.load(model_filename)  # Load the Isolation Forest model
    
    # Extract features from the audio file
    audio, sample_rate = librosa.load(sample_audio_file, sr=None)
    mfccs = librosa.feature.mfcc(y=audio, sr=sample_rate, n_mfcc=13)
    spectral_contrast = librosa.feature.spectral_contrast(y=audio, sr=sample_rate)
    feature = np.hstack((mfccs.mean(axis=1), spectral_contrast.mean(axis=1)))
    
    feature = feature.reshape(1, -1)  # Reshape for prediction
    prediction = clf.predict(feature)  # Predict anomalies
    
    return prediction[0]

# Define sample audio files for testing
sample_files = {
    "Sample 1": "Audio Samples/good28.wav",
    "Sample 2": "Audio Samples/fault35.wav",
    "Sample 3": "Audio Samples/fault31.wav",
    "Sample 4": "Audio Samples/good30.wav",
    "Sample 5": "Audio Samples/fault30.wav",
    "Sample 6": "Audio Samples/fault32.wav",
    "Sample 7": "Audio Samples/good29.wav",
}

# Streamlit UI setup
st.title("AI powered Audio Anomaly Detector")
st.write("Utilizing Isolation Forest for efficient anomaly detection in audio data. Suitable for unsupervised ML applications.")

# Audio file selection in sidebar
selected_sample = st.sidebar.selectbox("Choose a Sample .wav File", list(sample_files.keys()))
sample_audio_file = sample_files[selected_sample]

# Anomaly detection button
predict_button = st.button("Detect Anomalies")

if predict_button:
    result = predict_anomalies(sample_audio_file)  # Perform anomaly detection
    
    # Display results
    if result == 1:
        st.error("Anomaly Detected")
    else:
        st.success("Normal Audio")
    
    # Setup for audio playback
    audio = AudioSegment.from_wav(sample_audio_file)
    audio.export("temp.wav", format="wav")
    st.audio("temp.wav", format="audio/wav")

# Instruction on how to use the application
st.sidebar.markdown("**Instructions:** Select a file → Detect anomalies → Listen to the audio.")

# Credits and acknowledgements
st.markdown("### Credits")
st.caption('''
            Developed by @Aditya-Neural-Net-Ninja \n
            Leveraging Isolation Forest ML model for anomaly detection \n
            Gratitude to Streamlit and 🤗 Spaces for hosting and deployment!
            ''')

import sys

print(f"Python version: {sys.version}")