import gradio as gr
import torch
import librosa
import numpy as np
import spaces
from transformers import ASTFeatureExtractor, ASTForAudioClassification

# ============================================================
# LOAD MODEL
# ============================================================
MODEL_PATH = "."

print("Loading AST model...")
feature_extractor = ASTFeatureExtractor.from_pretrained(MODEL_PATH)
model = ASTForAudioClassification.from_pretrained(MODEL_PATH)
model.eval()
print("Model loaded successfully!")
print(model.config.id2label)

# ============================================================
# KONSTANTA
# ============================================================
SAMPLE_RATE = 16000
CLIP_DURATION = 3
MAX_SAMPLES = SAMPLE_RATE * CLIP_DURATION


# ============================================================
# PREPROCESSING
# ============================================================
def preprocess_audio(audio_path):
    audio, sr = librosa.load(audio_path, sr=SAMPLE_RATE, mono=True)

    # Normalisasi
    max_val = np.max(np.abs(audio))
    if max_val > 1e-8:
        audio = audio / max_val * 0.95

    # Potong 3 detik
    audio = audio[:MAX_SAMPLES]

    # Padding jika kurang
    if len(audio) < MAX_SAMPLES:
        audio = np.pad(audio, (0, MAX_SAMPLES - len(audio)))

    return audio


# ============================================================
# PREDIKSI
# ============================================================
@spaces.GPU
def prediksi(audio_path):
    if audio_path is None:
        return {"Tidak ada audio": 1.0}

    try:
        audio = preprocess_audio(audio_path)

        inputs = feature_extractor(
            audio,
            sampling_rate=SAMPLE_RATE,
            return_tensors="pt"
        )

        device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(device)
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1)[0]

        hasil = {
            model.config.id2label[i]: float(probs[i])
            for i in range(len(probs))
        }
        return hasil

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


# ============================================================
# UI
# ============================================================
with gr.Blocks() as demo:
    gr.Markdown("# ⚙️ Sistem Deteksi Kerusakan Motor Matik")
    gr.Markdown(
        """
        Upload atau rekam suara mesin.
        Audio akan otomatis diproses: mono, 16 kHz, 3 detik, lalu diklasifikasi menggunakan model AST.
        """
    )

    with gr.Row():
        input_audio = gr.Audio(
            sources=["upload", "microphone"],
            type="filepath",
            label="Audio Mesin"
        )
        output = gr.Label(
            label="Hasil Prediksi",
            num_top_classes=3
        )

    btn = gr.Button("Analisis", variant="primary")
    btn.click(fn=prediksi, inputs=input_audio, outputs=output)

demo.queue()
demo.launch()