import os

# --- ACCELERATION LAYER: Configure CPU thread pooling before loading torch ---
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"

import gradio as gr
import torch
import torch.nn as nn
import numpy as np
import soundfile as sf
from transformers import AutoFeatureExtractor, Wav2Vec2ForSequenceClassification

# Enforce thread execution limits natively in PyTorch
torch.set_num_threads(1)
torch.set_num_interop_threads(1)

MODEL_ID = "facebook/mms-lid-1024"

print("--> Initializing processor and model...")
processor = AutoFeatureExtractor.from_pretrained(MODEL_ID)
base_model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_ID)
base_model.eval()

print("--> Applying PyTorch Dynamic Quantization (INT8)...")
quantized_model = torch.quantization.quantize_dynamic(
    base_model,
    {nn.Linear},
    dtype=torch.qint8
)
quantized_model.eval()
print("--> Model ready and optimized for single-thread execution!")

def detect_language(audio_path):
    if not audio_path:
        return {"No Audio Provided": 1.0}
        
    try:
        # Blazing fast audio loading via soundfile instead of librosa
        audio_array, native_sr = sf.read(audio_path)
        
        # If stereo, extract just the first channel
        if len(audio_array.shape) > 1:
            audio_array = audio_array[:, 0]
            
        # MMS requires a strict 16kHz sampling rate. 
        # If the input SR doesn't match, we safely resample using numpy interpolation.
        if native_sr != 16000:
            duration = len(audio_array) / native_sr
            num_samples = int(duration * 16000)
            audio_array = np.interp(
                np.linspace(0, duration, num_samples),
                np.linspace(0, duration, len(audio_array)),
                audio_array
            )

        # Build feature maps
        inputs = processor(audio_array, sampling_rate=16000, return_tensors="pt", padding=True)

        # Hyper-fast inference
        with torch.no_grad():
            outputs = quantized_model(**inputs)
            logits = outputs.logits
            probs = torch.nn.functional.softmax(logits, dim=-1).numpy()[0]

        # Extract predictions
        label_probs = {
            str(quantized_model.config.id2label[i]): float(probs[i])
            for i in range(len(probs))
        }
        
        sorted_labels = dict(sorted(label_probs.items(), key=lambda item: item[1], reverse=True)[:3])
        return sorted_labels

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

# Standalone interface mapping
app = gr.Interface(
    fn=detect_language,
    inputs=gr.Audio(type="filepath", label="Audio Input"),
    outputs=gr.Label(label="Identified Languages"),
    title="🌍 Ultra-Fast MMS-LID-1024 Identifier",
    description="Optimized with single-thread thread pools and fast I/O processing to eliminate free-tier container latency."
)

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860, show_error=True)