import gradio as gr
import numpy as np
import onnxruntime
import librosa
import os

# Load the ONNX model
sess = onnxruntime.InferenceSession('weightsgru.onnx')

# Define a function to extract MFCC features from an audio file
utterance_length = 384
def extract_mfcc(file_path, utterance_length):
    """
    Extracts MFCC features from an audio file and pads/truncates them
    to a fixed utterance length.
    """
    try:
        # Get raw .wav data and sampling rate from librosa's load function
        raw_w, sampling_rate = librosa.load(file_path, sr=None)

        # Obtain MFCC Features from raw data
        mfcc_features = librosa.feature.mfcc(y=raw_w, sr=sampling_rate, n_mfcc=40, hop_length=512)

        # Adjust the utterance length
        if mfcc_features.shape[1] > utterance_length:
            mfcc_features_new = mfcc_features[:, :utterance_length]
        else:
            mfcc_features_new = np.pad(
                mfcc_features,
                ((0, 0), (0, utterance_length - mfcc_features.shape[1])),
                mode='constant',
                constant_values=0
            )

        # Transpose the matrix for the desired shape
        mfcc_features_new = mfcc_features_new.T
        return mfcc_features_new
    except Exception as e:
        print(f"Error extracting MFCC: {e}")
        return None

# Define a function to classify audio using the ONNX model
def classify_audio(audio_file):
    """
    Classifies an audio file of a bird call and returns the two most likely species.
    """
    if audio_file is None:
        return "Please upload an audio file.", "No second label."
        
    utterance_length = 384
    mfcc_features = extract_mfcc(audio_file, utterance_length)

    if mfcc_features is None:
        return "Error processing audio file.", "No second label."

    # Split the features into 6 segments for the GRU model
    features = np.split(mfcc_features, 6)
    inputs = [np.expand_dims(segment, axis=0) for segment in features]
    classes = np.zeros((1, 10))

    # Run each segment through the ONNX model and sum the outputs
    for input_data in inputs:
        output = sess.run(None, {'gru_input': input_data.astype(np.float32)})[0]
        classes += output

    # Get the indices of the top two predictions
    pred_final1 = np.argmax(classes)
    sorted_indices = np.argsort(classes.flatten())
    
    # The second highest prediction is the second to last element in the sorted array
    pred_final2 = sorted_indices[-2]

    # Map the indices to human-readable labels
    bird_labels = {
        0: 'crow', 1: 'duck', 2: 'owl', 3: 'goaway', 4: 'peafowl',
        5: 'sparrow', 6: 'bluejay', 7: 'Asiankoel', 8: 'lapwing', 9: 'woodpeewe'
    }
    highest_bird_label = bird_labels[pred_final1]
    second_highest_bird_label = bird_labels[pred_final2]

    return highest_bird_label, second_highest_bird_label

# Gradio interface definition
inputs = gr.Audio(type="filepath", label="Upload an audio file (.wav)")
# outputs = [gr.Label(label="Bird Label 1"), gr.Label(label="Bird Label 2")]
outputs = [
    gr.Textbox(label="Bird Label 1"),
    gr.Textbox(label="Bird Label 2")
]

interface = gr.Interface(
    fn=classify_audio,
    inputs=inputs,
    outputs=outputs,
    title="Multi-label Bird Call Identification from Field Recordings",
    description="The multi-label bird call identification system is designed to classify 10 different bird species: House Crow, Mallard Duck, Eurasian Owl, Grey Go-away, Indian Peafowl, House Sparrow, Blue Jay, Asian Koel, Red Lapwing, and Western Wood Pewee. This system utilizes the MFCC-GRU model and is trained on data obtained from the Xeno-Canto bird sound database.",
    css="""
    body {
        background-color: #f0f8ff; /* Light sky blue background */
    }
    .gradio-container {
        background: linear-gradient(135deg, #e0f7fa, #b2ebf2); /* A nice gradient for the container */
        border-radius: 15px;
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
        padding: 20px;
    }
    h1 {
        color: #00796b; /* Dark cyan for the title */
        text-align: center;
    }
    .markdown-heading {
        color: #004d40; /* Even darker shade for markdown headings */
    }
    .gr-button {
        background-color: #009688; /* Teal color for buttons */
        color: white;
        border: none;
        border-radius: 10px;
        font-weight: bold;
        transition: background-color 0.3s ease;
    }
    .gr-button:hover {
        background-color: #00796b; /* Darker teal on hover */
    }
    .gr-textbox {
        background-color: #e0f2f1; /* Light cyan background for textboxes */
        border-radius: 10px;
        border: 1px solid #b2dfdb;
    }
"""
)

# Launch the app
interface.launch()