#!/usr/bin/env python3
"""
Simple Audio Classifier using Madverse Music AI Model
Usage: python classify_audio.py <audio_file_path>
"""

import sys
import os
from pathlib import Path

def load_model():
    """Load the Madverse Music AI model"""
    try:
        from sonics import HFAudioClassifier
        print("🔄 Loading Madverse Music AI model...")
        model = HFAudioClassifier.from_pretrained("awsaf49/sonics-spectttra-alpha-120s")
        print("✅ Model loaded successfully!")
        return model
    except Exception as e:
        print(f"❌ Error loading model: {e}")
        return None

def classify_audio(model, audio_path):
    """Classify a single audio file"""
    try:
        import librosa
        import torch
        
        print(f"🎵 Analyzing: {audio_path}")
        
        # Load audio file (model uses 16kHz sample rate)
        audio, sr = librosa.load(audio_path, sr=16000)
        
        # Convert to tensor and add batch dimension
        audio_tensor = torch.FloatTensor(audio).unsqueeze(0)
        
        # Set model to evaluation mode
        model.eval()
        
        # Get prediction
        with torch.no_grad():
            output = model(audio_tensor)
            
            # Convert logit to probability using sigmoid
            prob = torch.sigmoid(output).item()
            
            # Classify: prob < 0.5 = Real, prob >= 0.5 = Fake
            if prob < 0.5:
                result = "Real"
                emoji = "🎤"
                description = "Human-created music"
                confidence = (1 - prob) * 2  # Convert to 0-1 scale
            else:
                result = "Fake"
                emoji = "🤖"
                description = "AI-generated music"
                confidence = (prob - 0.5) * 2  # Convert to 0-1 scale
            
        print(f"{emoji} Result: {result} ({description})")
        print(f"   Confidence: {confidence:.2f} | Raw output: {output.item():.3f}")
        return result
        
    except Exception as e:
        print(f"❌ Error classifying {audio_path}: {e}")
        return None

def classify_multiple_files(model, directory_path="."):
    """Classify all audio files in a directory"""
    audio_extensions = ['.wav', '.mp3', '.flac', '.m4a', '.ogg']
    directory = Path(directory_path)
    
    audio_files = []
    for ext in audio_extensions:
        audio_files.extend(list(directory.glob(f'*{ext}')))
    
    if not audio_files:
        print(f"No audio files found in {directory}")
        return
    
    print(f"Found {len(audio_files)} audio file(s) to classify:")
    print("=" * 50)
    
    results = {}
    for audio_file in audio_files:
        result = classify_audio(model, str(audio_file))
        if result:
            results[str(audio_file)] = result
        print()
    
    # Summary
    if results:
        print("📊 Summary:")
        print("=" * 30)
        real_count = sum(1 for r in results.values() if r.lower() == 'real')
        fake_count = len(results) - real_count
        print(f"🎤 Real music: {real_count}")
        print(f"🤖 AI-generated: {fake_count}")

def main():
    """Main function"""
    print("🎵 Madverse Music: AI Audio Classifier")
    print("=" * 40)
    
    # Load model
    model = load_model()
    if not model:
        return
    
    if len(sys.argv) > 1:
        # Classify specific file
        audio_path = sys.argv[1]
        if os.path.exists(audio_path):
            classify_audio(model, audio_path)
        else:
            print(f"❌ File not found: {audio_path}")
    else:
        # Classify all files in current directory
        print("\nNo specific file provided. Scanning current directory...")
        classify_multiple_files(model)

if __name__ == "__main__":
    main()