"""
Advanced Feature Extraction for Music 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: Advanced audio features for improved classification
License: Educational Purpose Only
"""

import numpy as np
import librosa
from scipy import stats
from typing import Dict
import warnings
warnings.filterwarnings('ignore')


class AdvancedFeatureExtractor:
    """
    Advanced feature extraction for music analysis
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    Email: help@rskworld.in
    """
    
    def __init__(self, sample_rate=22050):
        """Initialize Advanced Feature Extractor"""
        self.sample_rate = sample_rate
    
    def extract_spectral_contrast(self, audio: np.ndarray) -> np.ndarray:
        """
        Extract spectral contrast features
        
        Args:
            audio: Audio time series
            
        Returns:
            Spectral contrast features
        """
        contrast = librosa.feature.spectral_contrast(y=audio, sr=self.sample_rate)
        return np.mean(contrast.T, axis=0)
    
    def extract_tonnetz(self, audio: np.ndarray) -> np.ndarray:
        """
        Extract tonal centroid features (tonnetz)
        
        Args:
            audio: Audio time series
            
        Returns:
            Tonnetz features
        """
        tonnetz = librosa.feature.tonnetz(y=audio, sr=self.sample_rate)
        return np.mean(tonnetz.T, axis=0)
    
    def extract_spectral_flatness(self, audio: np.ndarray) -> float:
        """
        Extract spectral flatness (measure of tonality)
        
        Args:
            audio: Audio time series
            
        Returns:
            Mean spectral flatness
        """
        flatness = librosa.feature.spectral_flatness(y=audio)
        return np.mean(flatness)
    
    def extract_spectral_flux(self, audio: np.ndarray) -> float:
        """
        Extract spectral flux (rate of change in power spectrum)
        
        Args:
            audio: Audio time series
            
        Returns:
            Mean spectral flux
        """
        spec = np.abs(librosa.stft(audio))
        flux = np.sqrt(np.sum(np.diff(spec, axis=1)**2, axis=0))
        return np.mean(flux)
    
    def extract_harmonic_percussive(self, audio: np.ndarray) -> Dict[str, float]:
        """
        Separate harmonic and percussive components
        
        Args:
            audio: Audio time series
            
        Returns:
            Dictionary with harmonic and percussive ratios
        """
        y_harmonic, y_percussive = librosa.effects.hpss(audio)
        
        return {
            'harmonic_ratio': np.sum(np.abs(y_harmonic)) / (np.sum(np.abs(audio)) + 1e-10),
            'percussive_ratio': np.sum(np.abs(y_percussive)) / (np.sum(np.abs(audio)) + 1e-10)
        }
    
    def extract_mel_spectrogram_stats(self, audio: np.ndarray) -> Dict[str, float]:
        """
        Extract statistics from mel spectrogram
        
        Args:
            audio: Audio time series
            
        Returns:
            Dictionary with mel spectrogram statistics
        """
        mel_spec = librosa.feature.melspectrogram(y=audio, sr=self.sample_rate)
        mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)
        
        return {
            'mel_mean': np.mean(mel_spec_db),
            'mel_std': np.std(mel_spec_db),
            'mel_max': np.max(mel_spec_db),
            'mel_min': np.min(mel_spec_db),
            'mel_median': np.median(mel_spec_db),
            'mel_skew': stats.skew(mel_spec_db.flatten()),
            'mel_kurtosis': stats.kurtosis(mel_spec_db.flatten())
        }
    
    def extract_rhythm_features(self, audio: np.ndarray) -> Dict[str, any]:
        """
        Extract rhythm-related features
        
        Args:
            audio: Audio time series
            
        Returns:
            Dictionary with rhythm features
        """
        # Tempogram
        tempo, beats = librosa.beat.beat_track(y=audio, sr=self.sample_rate)
        
        # Onset strength
        onset_env = librosa.onset.onset_strength(y=audio, sr=self.sample_rate)
        
        return {
            'tempo': float(tempo),
            'beat_count': len(beats),
            'onset_strength_mean': np.mean(onset_env),
            'onset_strength_std': np.std(onset_env)
        }
    
    def extract_energy_features(self, audio: np.ndarray) -> Dict[str, float]:
        """
        Extract energy-related features
        
        Args:
            audio: Audio time series
            
        Returns:
            Dictionary with energy features
        """
        # Short-time energy
        frame_length = 2048
        hop_length = 512
        
        frames = librosa.util.frame(audio, frame_length=frame_length, hop_length=hop_length)
        energy = np.sum(frames**2, axis=0)
        
        return {
            'energy_mean': np.mean(energy),
            'energy_std': np.std(energy),
            'energy_max': np.max(energy),
            'energy_entropy': stats.entropy(energy + 1e-10)
        }
    
    def extract_all_advanced_features(self, audio: np.ndarray) -> Dict[str, any]:
        """
        Extract all advanced features
        
        Author: Molla Samser
        Company: RSK World
        Website: https://rskworld.in
        
        Args:
            audio: Audio time series
            
        Returns:
            Dictionary containing all advanced features
        """
        features = {}
        
        # Spectral features
        features['spectral_contrast'] = self.extract_spectral_contrast(audio)
        features['tonnetz'] = self.extract_tonnetz(audio)
        features['spectral_flatness'] = self.extract_spectral_flatness(audio)
        features['spectral_flux'] = self.extract_spectral_flux(audio)
        
        # Harmonic/Percussive
        hp_features = self.extract_harmonic_percussive(audio)
        features.update(hp_features)
        
        # Mel spectrogram stats
        mel_stats = self.extract_mel_spectrogram_stats(audio)
        features.update(mel_stats)
        
        # Rhythm features
        rhythm_features = self.extract_rhythm_features(audio)
        features.update(rhythm_features)
        
        # Energy features
        energy_features = self.extract_energy_features(audio)
        features.update(energy_features)
        
        return features
    
    def flatten_features(self, features: Dict[str, any]) -> np.ndarray:
        """Flatten feature dictionary to 1D array"""
        feature_list = []
        
        for key, value in features.items():
            if isinstance(value, np.ndarray):
                feature_list.extend(value)
            else:
                feature_list.append(value)
        
        return np.array(feature_list)


# Example usage
if __name__ == "__main__":
    """
    Demo script for AdvancedFeatureExtractor
    
    Author: Molla Samser
    Company: RSK World
    Website: https://rskworld.in
    """
    extractor = AdvancedFeatureExtractor()
    
    # Create dummy audio
    dummy_audio = np.random.randn(22050 * 30)  # 30 seconds
    
    print("Advanced Feature Extraction")
    print("Author: Molla Samser | RSK World")
    print("Website: https://rskworld.in")
    print("\nExtracting advanced features...")
    
    features = extractor.extract_all_advanced_features(dummy_audio)
    
    print(f"\nExtracted features:")
    print(f"- Spectral Contrast: {features['spectral_contrast'].shape}")
    print(f"- Tonnetz: {features['tonnetz'].shape}")
    print(f"- Spectral Flatness: {features['spectral_flatness']:.6f}")
    print(f"- Tempo: {features['tempo']:.2f} BPM")
    print(f"- Harmonic Ratio: {features['harmonic_ratio']:.4f}")
    print(f"- Percussive Ratio: {features['percussive_ratio']:.4f}")
    
    flat_features = extractor.flatten_features(features)
    print(f"\nTotal feature dimensions: {flat_features.shape[0]}")
    
    print("\n© 2026 RSK World - Molla Samser")