import gradio as gr
import torch
from model import MachineSoundCNN 
from config import NUM_CLASSES

# Import YOUR exact pipeline just like infer.py does!
from preprocessing import preprocess_audio
from features import audio_to_tensor

# 1. Define your exact classes
CLASSES = [
    "Machine 1 Normal", "Machine 1 Abnormal",
    "Machine 2 Normal", "Machine 2 Abnormal",
    "Machine 3 Normal", "Machine 3 Abnormal"
]

# 2. Load the Model 
device = torch.device('cpu')
model = MachineSoundCNN(num_classes=NUM_CLASSES)
model.load_state_dict(torch.load('best_model.pth', map_location=device, weights_only=True))
model.eval() 

# 3. Preprocessing & Inference Function
def predict_machine_sound(audio_path):
    if audio_path is None:
        return "Please upload an audio file."

    try:
        # Step A: Your exact Preprocessing (Resample -> Noise Reduce -> Silence -> Normalize)
        audio, sr = preprocess_audio(audio_path)

        # Step B: Your exact Feature Extraction (Mel Spec -> Padding -> Tensor)
        tensor = audio_to_tensor(audio, sr, augment=False)
        
        # Step C: Add batch dimension and move to CPU
        tensor = tensor.unsqueeze(0).to(device)

        # Step D: Forward Pass
        with torch.no_grad():
            output = model(tensor)
            probabilities = torch.nn.functional.softmax(output[0], dim=0)

        # Format output for Gradio
        result = {CLASSES[i]: float(probabilities[i]) for i in range(len(CLASSES))}
        return result
        
    except Exception as e:
        return {f"Error processing audio: {str(e)}": 1.0}

# 4. Build the Web Interface
interface = gr.Interface(
    fn=predict_machine_sound,
    inputs=gr.Audio(type="filepath", label="Upload Machine Audio (.wav)"),
    outputs=gr.Label(num_top_classes=6, label="CNN Prediction Confidence"),
    title="Industrial Machine Sound Anomaly Detector",
    description="Upload an audio clip of an industrial machine. The Custom CNN will analyze the audio using the exact training pipeline.",
    flagging_mode="never"
)

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