import os
import numpy as np
import librosa
import scipy.signal
from scipy.stats import kurtosis
import gradio as gr
import tensorflow as tf
import warnings

warnings.filterwarnings('ignore', category=UserWarning)

MAX_TIME_FRAMES = 313
N_MELS = 128
N_1D_FEATURES = 22

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'
}

class MachineListenerPreprocessor:
    def __init__(self, target_sr=16000, n_fft=2048, hop_length=512, n_mels=N_MELS, n_mfcc=20):
        self.sr = target_sr
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.n_mels = n_mels
        self.n_mfcc = n_mfcc
        
    def _apply_highpass_filter(self, y, cutoff=60.0):
        nyquist = 0.5 * self.sr
        normal_cutoff = cutoff / nyquist
        if normal_cutoff >= 1.0:
            return y
        b, a = scipy.signal.butter(4, normal_cutoff, btype='high', analog=False)
        return scipy.signal.filtfilt(b, a, y)

    def _truncate_silence(self, y, top_db=25):
        y_trimmed, _ = librosa.effects.trim(y, top_db=top_db, frame_length=self.n_fft, hop_length=self.hop_length)
        return y_trimmed

    def _mean_variance_normalize(self, y):
        return (y - np.mean(y)) / (np.std(y) + 1e-8)

    def process_audio(self, file_path):
        y, _ = librosa.load(file_path, sr=self.sr)
        y = self._apply_highpass_filter(y)
        y = self._truncate_silence(y, top_db=25)
        
        if len(y) == 0:
            raise ValueError(f"Silence only: {file_path}")
            
        y = self._mean_variance_normalize(y)
        
        # 2D Features
        mel_spec = librosa.feature.melspectrogram(y=y, sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length, n_mels=self.n_mels)
        log_mel_spec = librosa.power_to_db(mel_spec, ref=np.max)
        
        # 1D Features
        mfccs_mean = np.mean(librosa.feature.mfcc(S=log_mel_spec, n_mfcc=self.n_mfcc), axis=1)
        centroid_mean = np.mean(librosa.feature.spectral_centroid(y=y, sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length))
        
        stft_mag = np.abs(librosa.stft(y, n_fft=self.n_fft, hop_length=self.hop_length))
        frame_kurtosis = np.nan_to_num(kurtosis(stft_mag, axis=0, fisher=True, bias=False))
        kurtosis_mean = np.mean(frame_kurtosis)
        
        return {
            "2d_spectrogram": log_mel_spec,
            "1d_statistics": np.hstack([mfccs_mean, centroid_mean, kurtosis_mean]),
        }

def pad_or_truncate(spectrogram, max_frames):
    if spectrogram.shape[1] > max_frames:
        return spectrogram[:, :max_frames]
    elif spectrogram.shape[1] < max_frames:
        pad_width = max_frames - spectrogram.shape[1]
        return np.pad(spectrogram, pad_width=((0, 0), (0, pad_width)), mode='constant')
    return spectrogram

# Load the model
try:
    model = tf.keras.models.load_model('best_v2f_generalist.keras')
except Exception as e:
    print("Warning: Could not load model. Ensure the path is correct.", e)
    model = None

preprocessor = MachineListenerPreprocessor()

def predict(audio_filepath):
    if model is None:
        return "Model not loaded properly."
    if audio_filepath is None:
        return "Please upload an audio file."
    
    try:
        # Extract features
        features = preprocessor.process_audio(audio_filepath)
        
        spec_2d = pad_or_truncate(features["2d_spectrogram"], MAX_TIME_FRAMES)
        stat_1d = features["1d_statistics"]
        
        # Add batch dimensions
        spec_2d_batch = np.expand_dims(spec_2d, axis=0)
        # Note: If your model expects a specific shape, e.g., (batch, channels, height, width), adjust dimensions below.
        spec_2d_batch = np.expand_dims(spec_2d_batch, axis=-1)  
        stat_1d_batch = np.expand_dims(stat_1d, axis=0)
        
        # Predict
        predictions = model.predict([spec_2d_batch, stat_1d_batch])
        
        predicted_class_idx = np.argmax(predictions, axis=-1)[0]
        predicted_label = LABEL_MAP_INVERSE.get(predicted_class_idx, "Unknown")
        
        confidence = float(np.max(predictions))
        
        return f"Prediction: {predicted_label} (Confidence: {confidence:.2f})"
        
    except Exception as e:
        return f"Error processing file: {str(e)}"

# Create Gradio interface
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()