"""
Prediction Script for Music Genre Classification

Project: Music Classification Dataset
Author: Molla Samser
Company: RSK World
Designer & Tester: Rima Khatun
Website: https://rskworld.in
Email: help@rskworld.in, support@rskworld.in
Phone: +91 93305 39277
Description: Make predictions on new music files
License: Educational Purpose Only
"""

import os
import sys
import numpy as np
import joblib
import argparse
from typing import Dict
import warnings
warnings.filterwarnings('ignore')

# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.audio_processor import AudioProcessor
from utils.feature_extractor import FeatureExtractor


class MusicPredictor:
    """
    Music Genre Predictor
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    Email: help@rskworld.in
    """
    
    def __init__(self, model_dir: str = 'saved_models', model_type: str = 'random_forest'):
        """
        Initialize predictor
        
        Args:
            model_dir: Directory containing saved models
            model_type: Type of model to load
        """
        self.audio_processor = AudioProcessor()
        self.feature_extractor = FeatureExtractor()
        
        # Load model and encoders
        model_path = os.path.join(model_dir, f'{model_type}_model.pkl')
        encoder_path = os.path.join(model_dir, 'label_encoder.pkl')
        scaler_path = os.path.join(model_dir, 'scaler.pkl')
        
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model not found at {model_path}")
        
        self.model = joblib.load(model_path)
        self.label_encoder = joblib.load(encoder_path)
        self.scaler = joblib.load(scaler_path)
        
        print(f"✓ Loaded {model_type} model successfully")
    
    def predict(self, audio_path: str) -> Dict[str, any]:
        """
        Predict genre for audio file
        
        Args:
            audio_path: Path to audio file
            
        Returns:
            Dictionary with prediction results
        """
        # Load audio
        audio = self.audio_processor.load_audio(audio_path)
        
        # Extract features
        features = self.feature_extractor.extract_all_features(audio)
        flat_features = self.feature_extractor.flatten_features(features)
        
        # Scale features
        features_scaled = self.scaler.transform([flat_features])
        
        # Make prediction
        prediction_encoded = self.model.predict(features_scaled)[0]
        prediction_proba = self.model.predict_proba(features_scaled)[0]
        
        # Decode prediction
        predicted_genre = self.label_encoder.inverse_transform([prediction_encoded])[0]
        confidence = prediction_proba[prediction_encoded] * 100
        
        # Get all genre probabilities
        genre_probabilities = {
            genre: prob * 100 
            for genre, prob in zip(self.label_encoder.classes_, prediction_proba)
        }
        
        return {
            'predicted_genre': predicted_genre,
            'confidence': confidence,
            'all_probabilities': genre_probabilities,
            'features': features
        }
    
    def predict_batch(self, audio_paths: list) -> list:
        """
        Predict genres for multiple audio files
        
        Args:
            audio_paths: List of audio file paths
            
        Returns:
            List of prediction results
        """
        results = []
        
        for audio_path in audio_paths:
            try:
                result = self.predict(audio_path)
                result['file_path'] = audio_path
                results.append(result)
            except Exception as e:
                print(f"Error processing {audio_path}: {str(e)}")
                continue
        
        return results


def print_prediction_result(result: Dict[str, any], audio_path: str):
    """
    Print prediction result in a formatted way
    
    Args:
        result: Prediction result dictionary
        audio_path: Path to audio file
    """
    print("\n" + "="*70)
    print(f"Prediction Results for: {os.path.basename(audio_path)}")
    print("="*70)
    
    print(f"\n🎵 Predicted Genre: {result['predicted_genre'].upper()}")
    print(f"📊 Confidence: {result['confidence']:.2f}%")
    
    print("\n📈 Genre Probabilities:")
    sorted_probs = sorted(
        result['all_probabilities'].items(), 
        key=lambda x: x[1], 
        reverse=True
    )
    
    for genre, prob in sorted_probs:
        bar_length = int(prob / 2)
        bar = "█" * bar_length
        print(f"  {genre:12s} | {bar} {prob:5.2f}%")
    
    print("\n🎼 Audio Features:")
    print(f"  Tempo: {result['features']['tempo']:.2f} BPM")
    print(f"  Spectral Centroid: {result['features']['spectral_centroid']:.2f}")
    print(f"  Zero Crossing Rate: {result['features']['zero_crossing_rate']:.6f}")
    print(f"  RMS Energy: {result['features']['rms_energy']:.6f}")
    
    print("\n" + "="*70)


def main():
    """
    Main prediction function
    
    Author: Molla Samser
    Company: RSK World
    Website: https://rskworld.in
    """
    parser = argparse.ArgumentParser(description='Predict Music Genre')
    parser.add_argument('--audio', type=str, required=True,
                        help='Path to audio file')
    parser.add_argument('--model', type=str, default='random_forest',
                        choices=['random_forest', 'svm', 'knn'],
                        help='Model type to use')
    parser.add_argument('--model_dir', type=str, default='saved_models',
                        help='Directory containing saved models')
    
    args = parser.parse_args()
    
    print("="*70)
    print("Music Genre Classification - Prediction")
    print("Author: Molla Samser | Company: RSK World")
    print("Website: https://rskworld.in")
    print("Email: help@rskworld.in | Phone: +91 93305 39277")
    print("="*70)
    
    # Check if audio file exists
    if not os.path.exists(args.audio):
        print(f"\n❌ Error: Audio file not found at {args.audio}")
        return
    
    # Initialize predictor
    try:
        predictor = MusicPredictor(
            model_dir=args.model_dir,
            model_type=args.model
        )
    except FileNotFoundError as e:
        print(f"\n❌ Error: {str(e)}")
        print("Please train a model first using train_model.py")
        return
    
    # Make prediction
    print(f"\n🔍 Analyzing audio file: {args.audio}")
    result = predictor.predict(args.audio)
    
    # Print results
    print_prediction_result(result, args.audio)
    
    print("\n© 2026 RSK World - Founded by Molla Samser")
    print("Designer & Tester: Rima Khatun")


if __name__ == "__main__":
    main()