"""
Example inference script for Tiny-AST MAD Military Audio Classifier
"""

from transformers import ASTForAudioClassification, ASTFeatureExtractor
import librosa
import torch
import numpy as np

def classify_military_audio(audio_path, model_name="akashpaul123/tiny-ast-mad-military-audio-classifier"):
    """
    Classify military audio using the fine-tuned Tiny-AST model
    
    Args:
        audio_path (str): Path to audio file
        model_name (str): Hugging Face model name
    
    Returns:
        dict: Classification results
    """
    
    # Load model and feature extractor
    model = ASTForAudioClassification.from_pretrained(model_name)
    feature_extractor = ASTFeatureExtractor.from_pretrained(model_name)
    
    # Load and preprocess audio
    audio, sr = librosa.load(audio_path, sr=16000, duration=10.0)
    
    # Extract features
    inputs = feature_extractor(audio, sampling_rate=16000, return_tensors="pt")
    
    # Predict
    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=-1)
        predicted_class = torch.argmax(probabilities, dim=-1).item()
        confidence = probabilities[0][predicted_class].item()
    
    # Class mapping
    classes = ['Communication', 'Footsteps', 'Gunshot', 'Shelling', 
               'Vehicle', 'Helicopter', 'Fighter']
    
    return {
        'predicted_class': classes[predicted_class],
        'class_id': predicted_class,
        'confidence': confidence,
        'all_probabilities': {cls: prob.item() for cls, prob in zip(classes, probabilities[0])}
    }

# Example usage
if __name__ == "__main__":
    # Replace with your audio file path
    result = classify_military_audio("path/to/your/military_audio.wav")
    
    print(f"Predicted class: {result['predicted_class']}")
    print(f"Confidence: {result['confidence']:.4f}")
    print("\nAll class probabilities:")
    for class_name, prob in result['all_probabilities'].items():
        print(f"  {class_name}: {prob:.4f}")