import os
import numpy as np
import joblib
import gradio as gr
import warnings

from ml_features import preprocess_audio, extract_ml_features

warnings.filterwarnings('ignore')

LABEL_MAP_INVERSE = {
    0: 'Machine 1_Normal', 1: 'Machine 1_Abnormal',
    2: 'Machine 2_Normal', 3: 'Machine 2_Abnormal',
    4: 'Machine 3_Normal', 5: 'Machine 3_Abnormal'
}


# ==========================================
# LOAD MODEL
# ==========================================
model = None
model_error = ""
try:
    model = joblib.load("ml_XGBoost_v1_immed.joblib")
    print("XGBoost model loaded successfully.")
except Exception as e:
    model_error = str(e)
    print(f"Warning: Could not load model. Error: {model_error}")


# ==========================================
# PREDICTION
# ==========================================
def predict(audio_filepath):
    if model is None:
        return f"Model not loaded properly. Error: {model_error}"
    if audio_filepath is None:
        return "Please upload an audio file."

    try:
        # 1. Preprocess
        y = preprocess_audio(audio_filepath)

        # 2. Extract features
        features = extract_ml_features(y)

        # 3. Convert to array in the same column order the model expects
        feature_names = sorted(features.keys())
        X = np.array([[features[name] for name in feature_names]])

        # 4. Predict
        predicted_class = model.predict(X)[0]
        predicted_label = LABEL_MAP_INVERSE.get(int(predicted_class), "Unknown")

        # Get probabilities if available
        if hasattr(model, 'predict_proba'):
            proba = model.predict_proba(X)[0]
            confidence = float(np.max(proba))
            return f"Prediction: {predicted_label} (Confidence: {confidence:.2f})"
        else:
            return f"Prediction: {predicted_label}"

    except Exception as e:
        return f"Error processing file: {str(e)}"


# ==========================================
# GRADIO UI
# ==========================================
iface = gr.Interface(
    fn=predict,
    inputs=gr.Audio(type="filepath", label="Upload Machine Audio"),
    outputs="text",
    title="Machine Listener Diagnosis",
    description="Upload a sound from a machine to predict whether it is Normal or Abnormal."
)

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