import gradio as gr
import numpy as np
import librosa
import onnxruntime as rt
from transformers import Pipeline

class MujawwadPipeline:
    def __init__(self, model_path='model/mujawwad_classifier.onnx'):
        self.session = rt.InferenceSession(model_path)
        self.labels = ['bayati','hijaz','jiharkah','nihawand','rast','shoba','sikah']
        self.descriptions = {
            'bayati': 'Bayati (بياتي)',
            'hijaz': 'Hijaz (حجاز)',
            'jiharkah': 'Jiharkah (جهاركاه)',
            'nihawand': 'Nihawand (نهاوند)',
            'rast': 'Rast (راست)',
            'shoba': 'Shoba (صبا)',
            'sikah': 'Sikah (سيكاه)'
        }

    def predict(self, audio_path):
        try:
            input_data = self.preprocess(audio_path)
            input_name = self.session.get_inputs()[0].name
            output_name = self.session.get_outputs()[0].name
            predictions = self.session.run([output_name], {input_name: input_data.astype(np.float32)})[0]
            
            # Format results with descriptions
            results = {self.descriptions[self.labels[i]]: float(predictions[0][i]) 
                      for i in range(len(self.labels))}
            return results
        except Exception as e:
            return {"Error": str(e)}

def classify_audio(audio_path):
    classifier = MujawwadPipeline()
    result = classifier.predict(audio_path)
    return result

# Create Gradio interface
iface = gr.Interface(
    fn=classify_audio,
    inputs=gr.Audio(type="filepath", label="Upload Quranic Recitation"),
    outputs=gr.Label(num_top_classes=7),
    title="Maqamat Classification Model",
    description="Classify Quranic recitation maqamat into 7 traditional Arabic melodic modes",
    examples=[["examples/bayati_example.mp3"], ["examples/hijaz_example.mp3"]]
)

if __name__ == "__main__":
    iface.launch()