import gradio as gr
import numpy as np
import torch
import librosa
import soundfile as sf
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
import matplotlib.pyplot as plt
import tempfile
import os

# Constants
SAMPLING_RATE = 16000
MODEL_NAME = "MIT/ast-finetuned-audioset-10-10-0.4593"
DEFAULT_THRESHOLD = 0.7

# Load model components
try:
    feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
    model = AutoModelForAudioClassification.from_pretrained(MODEL_NAME)
except Exception as e:
    print(f"Error loading model: {str(e)}")

# Equipment knowledge base
EQUIPMENT_RECOMMENDATIONS = {
    "bearing": {
        "high_frequency": "• Replace bearings immediately\n• Check lubrication system\n• Monitor vibration levels",
        "low_frequency": "• Inspect bearing installation\n• Check for contamination\n• Verify lubrication",
        "irregular": "• Perform vibration analysis\n• Schedule bearing replacement\n• Check alignment"
    },
    "pump": {
        "cavitation": "• Check NPSH available\n• Inspect suction strainer\n• Adjust operating speed",
        "impeller": "• Inspect impeller for damage\n• Perform dynamic balancing\n• Check wear rings",
        "misalignment": "• Perform laser alignment\n• Check coupling condition\n• Verify baseplate level"
    },
    "motor": {
        "electrical": "• Megger test windings\n• Check connections\n• Inspect starter contacts",
        "mechanical": "• Perform dynamic balancing\n• Check alignment\n• Inspect cooling fins",
        "bearing": "• Replace motor bearings\n• Check lubrication\n• Monitor temperature"
    }
}

def analyze_frequency_patterns(audio, sr):
    """Analyze frequency patterns to identify potential issues"""
    patterns = []
    features = {}
    
    # Spectral analysis
    spectral_centroid = librosa.feature.spectral_centroid(y=audio, sr=sr)[0]
    spectral_rolloff = librosa.feature.spectral_rolloff(y=audio, sr=sr)[0]
    
    features['centroid_mean'] = np.mean(spectral_centroid)
    features['rolloff_mean'] = np.mean(spectral_rolloff)
    
    if features['centroid_mean'] > 3000:
        patterns.append("high_frequency")
    elif features['centroid_mean'] < 1000:
        patterns.append("low_frequency")
        
    if features['rolloff_mean'] > 8000:
        patterns.append("harmonic_rich")
    
    return patterns, features

def generate_recommendation(prediction, confidence, audio, sr):
    """Generate maintenance recommendations based on analysis"""
    if prediction == "Normal":
        return "✅ No immediate action required. Equipment operating within normal parameters."
    
    patterns, features = analyze_frequency_patterns(audio, sr)
    
    # Equipment classification
    spectral_flatness = librosa.feature.spectral_flatness(y=audio)[0]
    mean_flatness = np.mean(spectral_flatness)
    
    if mean_flatness < 0.2:
        equipment_type = "bearing"
    elif 0.2 <= mean_flatness < 0.6:
        equipment_type = "pump"
    else:
        equipment_type = "motor"
    
    # Generate recommendations
    recommendations = [
        "🔧 MAINTENANCE RECOMMENDATIONS",
        f"Equipment Type: {equipment_type.upper()}",
        f"Confidence: {confidence:.1%}",
        ""
    ]
    
    for pattern in patterns:
        if pattern in EQUIPMENT_RECOMMENDATIONS.get(equipment_type, {}):
            recommendations.append(EQUIPMENT_RECOMMENDATIONS[equipment_type][pattern])
    
    if prediction == "Anomaly":
        recommendations.extend([
            "",
            "🛠️ GENERAL ACTIONS:",
            "1. Isolate equipment if possible",
            "2. Perform visual inspection",
            "3. Schedule detailed diagnostics",
        ])
        
        if confidence > 0.8:
            recommendations.append("\n🚨 URGENT: High-confidence abnormality detected!")

    return "\n".join(recommendations)

def process_audio(file_path):
    """Handle audio file processing"""
    try:
        audio, sr = librosa.load(file_path, sr=SAMPLING_RATE, mono=True)
        return audio, sr
    except Exception as e:
        raise RuntimeError(f"Audio processing error: {str(e)}")

def analyze_audio(audio_input, threshold=DEFAULT_THRESHOLD):
    """Main analysis function"""
    try:
        # Handle file upload
        if isinstance(audio_input, str):
            audio, sr = process_audio(audio_input)
        else:  # Handle file object
            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp:
                tmp.write(audio_input.read())
                tmp_path = tmp.name
            audio, sr = process_audio(tmp_path)
            os.unlink(tmp_path)
        
        # Model prediction
        inputs = feature_extractor(audio, sampling_rate=SAMPLING_RATE, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1)
        
        predicted_class = "Normal" if probs[0][0] > threshold else "Anomaly"
        confidence = probs[0][0].item() if predicted_class == "Normal" else 1 - probs[0][0].item()
        
        # Generate visualization
        plt.figure(figsize=(10, 4))
        S = librosa.feature.melspectrogram(y=audio, sr=SAMPLING_RATE, n_mels=64)
        S_db = librosa.power_to_db(S, ref=np.max)
        librosa.display.specshow(S_db, x_axis='time', y_axis='mel', sr=SAMPLING_RATE, fmax=8000)
        plt.colorbar(format='%+2.0f dB')
        plt.title('Mel Spectrogram')
        
        spec_path = os.path.join(tempfile.gettempdir(), 'spec.png')
        plt.savefig(spec_path, bbox_inches='tight')
        plt.close()
        
        # Generate recommendations
        recommendations = generate_recommendation(predicted_class, confidence, audio, SAMPLING_RATE)
        
        return (
            predicted_class,
            f"{confidence:.1%}",
            spec_path,
            recommendations
        )
    
    except Exception as e:
        return f"Error: {str(e)}", "", None, ""

# Gradio Interface
with gr.Blocks(title="Industrial Audio Analyzer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏭 Industrial Equipment Sound Analyzer
    ### Acoustic Anomaly Detection & Maintenance Recommendation System
    """)
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(
                label="Upload Equipment Audio (.wav)",
                type="filepath"
            )
            threshold = gr.Slider(
                minimum=0.5, maximum=0.95, step=0.05, value=DEFAULT_THRESHOLD,
                label="Detection Sensitivity"
            )
            analyze_btn = gr.Button("🔍 Analyze & Diagnose", variant="primary")
            
        with gr.Column():
            result_label = gr.Label(label="Diagnosis Result")
            confidence = gr.Textbox(label="Confidence Score")
            spectrogram = gr.Image(label="Spectrogram Analysis")
            recommendations = gr.Textbox(
                label="Maintenance Recommendations",
                lines=10,
                interactive=False
            )
    
    analyze_btn.click(
        fn=analyze_audio,
        inputs=[audio_input, threshold],
        outputs=[result_label, confidence, spectrogram, recommendations]
    )
    
    gr.Markdown("""
    **Instructions:**
    - Upload 5-10 second .wav recordings
    - Results include:
      ✓ Anomaly detection
      ✓ Equipment classification
      ✓ Maintenance recommendations
      ✓ Spectrogram visualization
    """)

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