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

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

# Memory optimizations for TensorFlow
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

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

# ==========================================
# 1. PREPROCESSOR
# ==========================================
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


# ==========================================
# 2. CUSTOM LAYER FOR MODEL LOADING
# ==========================================
@tf.keras.utils.register_keras_serializable()
class SpecAugmentLayer(tf.keras.layers.Layer):
    def __init__(self, freq_mask_param=15, time_mask_param=30, **kwargs):
        super().__init__(**kwargs)
        self.freq_mask_param = freq_mask_param
        self.time_mask_param = time_mask_param

    def call(self, inputs, training=None):
        return inputs # During inference, SpecAugment does nothing

    def get_config(self):
        config = super().get_config()
        config.update({
            "freq_mask_param": self.freq_mask_param,
            "time_mask_param": self.time_mask_param,
        })
        return config


# ==========================================
# 3. INFERENCE LOGIC
# ==========================================
def main():
    parser = argparse.ArgumentParser(description="Classify machine audio.")
    parser.add_argument("audio_path", help="Path to the .wav audio file")
    parser.add_argument("--model", default="best_v2f_generalist.keras", help="Path to model")
    args = parser.parse_args()

    if not os.path.exists(args.audio_path):
        print(f"Error: Audio file not found at {args.audio_path}")
        return

    # Enable custom layers & lambda loading
    try:
        import keras
        keras.config.enable_unsafe_deserialization()
    except:
        pass

    print("Loading model...")
    custom_objects = {'SpecAugmentLayer': SpecAugmentLayer}
    model = tf.keras.models.load_model(args.model, custom_objects=custom_objects)
    
    print("Extracting features...")
    preprocessor = MachineListenerPreprocessor()
    features = preprocessor.process_audio(args.audio_path)

    # Prepare inputs
    spec_2d = pad_or_truncate(features["2d_spectrogram"], MAX_TIME_FRAMES)
    stat_1d = features["1d_statistics"]

    # Add batch and channel dimensions
    spec_2d_batch = np.expand_dims(spec_2d, axis=0)        # (1, 128, 313)
    spec_2d_batch = np.expand_dims(spec_2d_batch, axis=-1) # (1, 128, 313, 1)
    stat_1d_batch = np.expand_dims(stat_1d, axis=0)        # (1, 22)

    print("Running prediction...")
    # model(inputs, training=False) is much faster and lighter than model.predict()
    predictions = model([spec_2d_batch, stat_1d_batch], training=False).numpy()

    predicted_class_idx = np.argmax(predictions, axis=-1)[0]
    predicted_label = LABEL_MAP_INVERSE.get(predicted_class_idx, "Unknown")
    confidence = float(np.max(predictions))

    print("\n" + "="*40)
    print(f"File: {os.path.basename(args.audio_path)}")
    print(f"Result: {predicted_label}")
    print(f"Confidence: {confidence:.2%}")
    print("="*40)

if __name__ == "__main__":
    main()