"""
Feature Extractor Module 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: Extract audio features for music genre classification
License: Educational Purpose Only
"""

import librosa
import numpy as np
from typing import Dict, List, Tuple
import warnings
warnings.filterwarnings('ignore')


class FeatureExtractor:
    """
    Feature extraction class for audio analysis
    
    Author: Molla Samser (RSK World)
    Website: https://rskworld.in
    """
    
    def __init__(self, sample_rate: int = 22050):
        """
        Initialize FeatureExtractor
        
        Args:
            sample_rate: Sample rate of audio files
        """
        self.sample_rate = sample_rate
        
    def extract_mfcc(self, audio: np.ndarray, n_mfcc: int = 13) -> np.ndarray:
        """
        Extract MFCC features
        
        Args:
            audio: Audio time series
            n_mfcc: Number of MFCCs to extract
            
        Returns:
            MFCC features array
        """
        mfcc = librosa.feature.mfcc(y=audio, sr=self.sample_rate, n_mfcc=n_mfcc)
        return np.mean(mfcc.T, axis=0)
    
    def extract_spectral_centroid(self, audio: np.ndarray) -> float:
        """
        Extract spectral centroid
        
        Args:
            audio: Audio time series
            
        Returns:
            Mean spectral centroid value
        """
        spectral_centroid = librosa.feature.spectral_centroid(y=audio, sr=self.sample_rate)
        return np.mean(spectral_centroid)
    
    def extract_spectral_rolloff(self, audio: np.ndarray) -> float:
        """
        Extract spectral rolloff
        
        Args:
            audio: Audio time series
            
        Returns:
            Mean spectral rolloff value
        """
        spectral_rolloff = librosa.feature.spectral_rolloff(y=audio, sr=self.sample_rate)
        return np.mean(spectral_rolloff)
    
    def extract_zero_crossing_rate(self, audio: np.ndarray) -> float:
        """
        Extract zero crossing rate
        
        Args:
            audio: Audio time series
            
        Returns:
            Mean zero crossing rate
        """
        zcr = librosa.feature.zero_crossing_rate(audio)
        return np.mean(zcr)
    
    def extract_chroma(self, audio: np.ndarray) -> np.ndarray:
        """
        Extract chroma features
        
        Args:
            audio: Audio time series
            
        Returns:
            Chroma features array
        """
        chroma = librosa.feature.chroma_stft(y=audio, sr=self.sample_rate)
        return np.mean(chroma.T, axis=0)
    
    def extract_tempo(self, audio: np.ndarray) -> float:
        """
        Extract tempo (BPM)
        
        Args:
            audio: Audio time series
            
        Returns:
            Tempo in beats per minute
        """
        tempo, _ = librosa.beat.beat_track(y=audio, sr=self.sample_rate)
        return float(tempo)
    
    def extract_spectral_bandwidth(self, audio: np.ndarray) -> float:
        """
        Extract spectral bandwidth
        
        Args:
            audio: Audio time series
            
        Returns:
            Mean spectral bandwidth
        """
        spectral_bandwidth = librosa.feature.spectral_bandwidth(y=audio, sr=self.sample_rate)
        return np.mean(spectral_bandwidth)
    
    def extract_rms_energy(self, audio: np.ndarray) -> float:
        """
        Extract RMS energy
        
        Args:
            audio: Audio time series
            
        Returns:
            Mean RMS energy
        """
        rms = librosa.feature.rms(y=audio)
        return np.mean(rms)
    
    def extract_all_features(self, audio: np.ndarray) -> Dict[str, any]:
        """
        Extract all features at once
        
        Args:
            audio: Audio time series
            
        Returns:
            Dictionary containing all extracted features
        """
        features = {
            'mfcc': self.extract_mfcc(audio),
            'spectral_centroid': self.extract_spectral_centroid(audio),
            'spectral_rolloff': self.extract_spectral_rolloff(audio),
            'zero_crossing_rate': self.extract_zero_crossing_rate(audio),
            'chroma': self.extract_chroma(audio),
            'tempo': self.extract_tempo(audio),
            'spectral_bandwidth': self.extract_spectral_bandwidth(audio),
            'rms_energy': self.extract_rms_energy(audio)
        }
        
        return features
    
    def flatten_features(self, features: Dict[str, any]) -> np.ndarray:
        """
        Flatten feature dictionary to 1D array
        
        Args:
            features: Dictionary of features
            
        Returns:
            Flattened feature 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 FeatureExtractor
    
    Author: Molla Samser
    Company: RSK World
    Website: https://rskworld.in
    """
    from audio_processor import AudioProcessor
    
    # Initialize processors
    audio_processor = AudioProcessor()
    feature_extractor = FeatureExtractor()
    
    # Example audio file
    audio_path = "../data/audio/jazz/jazz_001.wav"
    
    # Load audio
    print(f"Loading audio from: {audio_path}")
    audio = audio_processor.load_audio(audio_path)
    
    # Extract all features
    print("\nExtracting features...")
    features = feature_extractor.extract_all_features(audio)
    
    # Display features
    print("\nExtracted Features:")
    print(f"- MFCC shape: {features['mfcc'].shape}")
    print(f"- Spectral Centroid: {features['spectral_centroid']:.2f}")
    print(f"- Spectral Rolloff: {features['spectral_rolloff']:.2f}")
    print(f"- Zero Crossing Rate: {features['zero_crossing_rate']:.6f}")
    print(f"- Chroma shape: {features['chroma'].shape}")
    print(f"- Tempo: {features['tempo']:.2f} BPM")
    print(f"- Spectral Bandwidth: {features['spectral_bandwidth']:.2f}")
    print(f"- RMS Energy: {features['rms_energy']:.6f}")
    
    # Flatten features
    flat_features = feature_extractor.flatten_features(features)
    print(f"\nFlattened features shape: {flat_features.shape}")