"""
Environmental Sound Dataset - Audio Analysis Module

Project: Environmental Sound Dataset
Website: https://rskworld.in
Founded by: Molla Samser
Designer & Tester: Rima Khatun
Email: help@rskworld.in
Phone: +91 93305 39277
"""

import numpy as np
import librosa
import librosa.display
import matplotlib.pyplot as plt
from typing import Dict, Tuple, Optional
import os


def analyze_audio_file(
    audio_path: str,
    sample_rate: int = 22050
) -> Dict:
    """
    Analyze a single audio file and extract various features.
    
    Args:
        audio_path: Path to the audio file
        sample_rate: Target sample rate
    
    Returns:
        Dictionary containing audio analysis results
    """
    if not os.path.exists(audio_path):
        raise FileNotFoundError(f"Audio file not found: {audio_path}")
    
    # Load audio
    y, sr = librosa.load(audio_path, sr=sample_rate)
    
    # Calculate features
    duration = len(y) / sr
    tempo, beats = librosa.beat.beat_track(y=y, sr=sr)
    
    # Extract various features
    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    spectral_centroids = librosa.feature.spectral_centroid(y=y, sr=sr)[0]
    spectral_rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr)[0]
    zero_crossing_rate = librosa.feature.zero_crossing_rate(y)[0]
    chroma = librosa.feature.chroma_stft(y=y, sr=sr)
    mel_spectrogram = librosa.feature.melspectrogram(y=y, sr=sr)
    
    analysis = {
        'file_path': audio_path,
        'duration': duration,
        'sample_rate': sr,
        'tempo': float(tempo),
        'mean_mfcc': np.mean(mfccs, axis=1).tolist(),
        'mean_spectral_centroid': float(np.mean(spectral_centroids)),
        'mean_spectral_rolloff': float(np.mean(spectral_rolloff)),
        'mean_zero_crossing_rate': float(np.mean(zero_crossing_rate)),
        'mean_chroma': np.mean(chroma, axis=1).tolist(),
        'mean_mel_spectrogram': np.mean(mel_spectrogram, axis=1).tolist(),
        'max_amplitude': float(np.max(np.abs(y))),
        'rms_energy': float(np.sqrt(np.mean(y**2)))
    }
    
    return analysis


def plot_audio_features(audio_path: str, save_path: Optional[str] = None):
    """
    Create visualization plots for audio features.
    
    Args:
        audio_path: Path to the audio file
        save_path: Optional path to save the plot
    """
    y, sr = librosa.load(audio_path)
    
    fig, axes = plt.subplots(3, 1, figsize=(12, 10))
    
    # Waveform
    time = np.linspace(0, len(y) / sr, len(y))
    axes[0].plot(time, y)
    axes[0].set_title('Waveform')
    axes[0].set_xlabel('Time (s)')
    axes[0].set_ylabel('Amplitude')
    axes[0].grid(True)
    
    # Spectrogram
    D = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
    librosa.display.specshow(D, y_axis='hz', x_axis='time', sr=sr, ax=axes[1])
    axes[1].set_title('Spectrogram')
    axes[1].set_ylabel('Frequency (Hz)')
    
    # MFCC
    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    librosa.display.specshow(mfccs, x_axis='time', sr=sr, ax=axes[2])
    axes[2].set_title('MFCC')
    axes[2].set_ylabel('MFCC Coefficients')
    axes[2].set_xlabel('Time (s)')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {save_path}")
    else:
        plt.show()


def compare_audio_files(audio_paths: list, labels: list = None) -> Dict:
    """
    Compare multiple audio files and return comparative analysis.
    
    Args:
        audio_paths: List of paths to audio files
        labels: Optional list of labels for each file
    
    Returns:
        Dictionary with comparative analysis
    """
    if labels is None:
        labels = [f"File {i+1}" for i in range(len(audio_paths))]
    
    comparisons = {}
    
    for path, label in zip(audio_paths, labels):
        try:
            analysis = analyze_audio_file(path)
            comparisons[label] = analysis
        except Exception as e:
            print(f"Error analyzing {path}: {e}")
    
    return comparisons


def get_dataset_statistics(data_dir: str = './environmental-sounds') -> Dict:
    """
    Get overall statistics about the dataset.
    
    Args:
        data_dir: Root directory of the dataset
    
    Returns:
        Dictionary with dataset statistics
    """
    stats = {
        'total_files': 0,
        'total_duration': 0,
        'classes': {},
        'formats': {},
        'avg_duration': 0,
        'min_duration': float('inf'),
        'max_duration': 0
    }
    
    for split in ['train', 'test']:
        split_dir = os.path.join(data_dir, split)
        if not os.path.exists(split_dir):
            continue
        
        for class_name in os.listdir(split_dir):
            class_dir = os.path.join(split_dir, class_name)
            if not os.path.isdir(class_dir):
                continue
            
            if class_name not in stats['classes']:
                stats['classes'][class_name] = 0
            
            for audio_file in os.listdir(class_dir):
                if audio_file.endswith(('.wav', '.mp3', '.flac')):
                    stats['total_files'] += 1
                    stats['classes'][class_name] += 1
                    
                    ext = os.path.splitext(audio_file)[1]
                    stats['formats'][ext] = stats['formats'].get(ext, 0) + 1
                    
                    try:
                        y, sr = librosa.load(os.path.join(class_dir, audio_file))
                        duration = len(y) / sr
                        stats['total_duration'] += duration
                        stats['min_duration'] = min(stats['min_duration'], duration)
                        stats['max_duration'] = max(stats['max_duration'], duration)
                    except:
                        pass
    
    if stats['total_files'] > 0:
        stats['avg_duration'] = stats['total_duration'] / stats['total_files']
    
    if stats['min_duration'] == float('inf'):
        stats['min_duration'] = 0
    
    return stats


if __name__ == '__main__':
    # Example usage
    print("Audio Analysis Module")
    print("=" * 50)
    
    # Get dataset statistics
    try:
        stats = get_dataset_statistics()
        print("\nDataset Statistics:")
        print(f"Total files: {stats['total_files']}")
        print(f"Total duration: {stats['total_duration']:.2f} seconds")
        print(f"Average duration: {stats['avg_duration']:.2f} seconds")
        print(f"Min duration: {stats['min_duration']:.2f} seconds")
        print(f"Max duration: {stats['max_duration']:.2f} seconds")
        print(f"\nClasses: {list(stats['classes'].keys())}")
        print(f"Formats: {stats['formats']}")
    except Exception as e:
        print(f"Error: {e}")

"""
Environmental Sound Dataset - Audio Quality Assessment Module

Project: Environmental Sound Dataset
Website: https://rskworld.in
Founded by: Molla Samser
Designer & Tester: Rima Khatun
Email: help@rskworld.in
Phone: +91 93305 39277

This module provides audio quality assessment and validation tools.
"""

import numpy as np
import librosa
from typing import Dict, Tuple, List
import os


class AudioQualityAssessor:
    """
    Assess audio quality and detect issues.
    """
    
    def __init__(self, sample_rate: int = 22050):
        """
        Initialize audio quality assessor.
        
        Args:
            sample_rate: Expected sample rate
        """
        self.sample_rate = sample_rate
    
    def assess_audio(self, audio: np.ndarray, sr: int = None) -> Dict:
        """
        Comprehensive audio quality assessment.
        
        Args:
            audio: Audio waveform
            sr: Sample rate (uses self.sample_rate if None)
        
        Returns:
            Dictionary with quality metrics
        """
        if sr is None:
            sr = self.sample_rate
        
        assessment = {
            'duration': len(audio) / sr,
            'sample_rate': sr,
            'n_samples': len(audio),
            'max_amplitude': float(np.max(np.abs(audio))),
            'rms_energy': float(np.sqrt(np.mean(audio**2))),
            'zero_crossing_rate': float(np.mean(librosa.feature.zero_crossing_rate(audio))),
            'issues': []
        }
        
        # Check for clipping
        if assessment['max_amplitude'] >= 0.99:
            assessment['issues'].append('clipping_detected')
        
        # Check for silence
        if assessment['rms_energy'] < 0.01:
            assessment['issues'].append('likely_silence')
        
        # Check for very short audio
        if assessment['duration'] < 0.1:
            assessment['issues'].append('too_short')
        
        # Check for DC offset
        dc_offset = np.mean(audio)
        if abs(dc_offset) > 0.1:
            assessment['issues'].append('dc_offset')
            assessment['dc_offset'] = float(dc_offset)
        
        # Calculate SNR estimate (simple)
        signal_power = np.mean(audio**2)
        noise_estimate = np.var(audio - np.mean(audio))
        if noise_estimate > 0:
            snr = 10 * np.log10(signal_power / noise_estimate)
            assessment['estimated_snr_db'] = float(snr)
        else:
            assessment['estimated_snr_db'] = float('inf')
        
        # Quality score (0-100)
        quality_score = 100
        if 'clipping_detected' in assessment['issues']:
            quality_score -= 20
        if 'likely_silence' in assessment['issues']:
            quality_score -= 30
        if 'too_short' in assessment['issues']:
            quality_score -= 15
        if 'dc_offset' in assessment['issues']:
            quality_score -= 10
        
        assessment['quality_score'] = max(0, quality_score)
        
        return assessment
    
    def validate_audio_file(self, filepath: str) -> Tuple[bool, Dict]:
        """
        Validate an audio file.
        
        Args:
            filepath: Path to audio file
        
        Returns:
            Tuple of (is_valid, assessment_dict)
        """
        try:
            audio, sr = librosa.load(filepath, sr=self.sample_rate)
            assessment = self.assess_audio(audio, sr)
            
            # Consider valid if no critical issues
            critical_issues = ['too_short', 'likely_silence']
            is_valid = not any(issue in assessment['issues'] for issue in critical_issues)
            
            return is_valid, assessment
        
        except Exception as e:
            return False, {'error': str(e), 'issues': ['load_error']}
    
    def batch_validate(self, file_paths: List[str]) -> Dict:
        """
        Validate multiple audio files.
        
        Args:
            file_paths: List of audio file paths
        
        Returns:
            Dictionary with validation results
        """
        results = {
            'total': len(file_paths),
            'valid': 0,
            'invalid': 0,
            'errors': 0,
            'details': []
        }
        
        for filepath in file_paths:
            is_valid, assessment = self.validate_audio_file(filepath)
            
            result = {
                'file': filepath,
                'valid': is_valid,
                'assessment': assessment
            }
            
            results['details'].append(result)
            
            if 'error' in assessment:
                results['errors'] += 1
            elif is_valid:
                results['valid'] += 1
            else:
                results['invalid'] += 1
        
        return results
    
    def normalize_audio(self, audio: np.ndarray, target_level: float = -3.0) -> np.ndarray:
        """
        Normalize audio to target level in dB.
        
        Args:
            audio: Audio waveform
            target_level: Target level in dB
        
        Returns:
            Normalized audio
        """
        # Calculate current RMS
        rms = np.sqrt(np.mean(audio**2))
        
        if rms > 0:
            # Calculate target RMS
            target_rms = 10 ** (target_level / 20)
            
            # Calculate gain
            gain = target_rms / rms
            
            # Apply gain
            normalized = audio * gain
            
            # Prevent clipping
            if np.max(np.abs(normalized)) > 1.0:
                normalized = normalized / np.max(np.abs(normalized))
            
            return normalized
        
        return audio
    
    def remove_dc_offset(self, audio: np.ndarray) -> np.ndarray:
        """
        Remove DC offset from audio.
        
        Args:
            audio: Audio waveform
        
        Returns:
            Audio with DC offset removed
        """
        return audio - np.mean(audio)
    
    def enhance_audio(self, audio: np.ndarray) -> np.ndarray:
        """
        Enhance audio quality (normalize, remove DC offset).
        
        Args:
            audio: Audio waveform
        
        Returns:
            Enhanced audio
        """
        enhanced = self.remove_dc_offset(audio)
        enhanced = self.normalize_audio(enhanced)
        return enhanced


class AudioDatasetValidator:
    """
    Validate entire audio dataset.
    """
    
    def __init__(self, data_dir: str = './environmental-sounds'):
        """
        Initialize dataset validator.
        
        Args:
            data_dir: Root directory of dataset
        """
        self.data_dir = data_dir
        self.assessor = AudioQualityAssessor()
    
    def validate_dataset(self) -> Dict:
        """
        Validate entire dataset.
        
        Returns:
            Dictionary with validation results
        """
        results = {
            'train': {'valid': 0, 'invalid': 0, 'total': 0, 'issues': {}},
            'test': {'valid': 0, 'invalid': 0, 'total': 0, 'issues': {}}
        }
        
        for split in ['train', 'test']:
            split_dir = os.path.join(self.data_dir, split)
            
            if not os.path.exists(split_dir):
                continue
            
            for class_name in os.listdir(split_dir):
                class_dir = os.path.join(split_dir, class_name)
                
                if not os.path.isdir(class_dir):
                    continue
                
                for audio_file in os.listdir(class_dir):
                    if audio_file.endswith(('.wav', '.mp3', '.flac')):
                        filepath = os.path.join(class_dir, audio_file)
                        results[split]['total'] += 1
                        
                        is_valid, assessment = self.assessor.validate_audio_file(filepath)
                        
                        if is_valid:
                            results[split]['valid'] += 1
                        else:
                            results[split]['invalid'] += 1
                            
                            # Track issues
                            for issue in assessment.get('issues', []):
                                if issue not in results[split]['issues']:
                                    results[split]['issues'][issue] = 0
                                results[split]['issues'][issue] += 1
        
        return results
    
    def generate_report(self, save_path: str = None) -> str:
        """
        Generate validation report.
        
        Args:
            save_path: Optional path to save report
        
        Returns:
            Report as string
        """
        results = self.validate_dataset()
        
        report = "Audio Dataset Validation Report\n"
        report += "=" * 50 + "\n\n"
        
        for split in ['train', 'test']:
            report += f"{split.upper()} SET:\n"
            report += f"  Total files: {results[split]['total']}\n"
            report += f"  Valid: {results[split]['valid']}\n"
            report += f"  Invalid: {results[split]['invalid']}\n"
            
            if results[split]['invalid'] > 0:
                report += f"  Issues found:\n"
                for issue, count in results[split]['issues'].items():
                    report += f"    - {issue}: {count}\n"
            
            report += "\n"
        
        if save_path:
            with open(save_path, 'w') as f:
                f.write(report)
            print(f"Report saved to {save_path}")
        
        return report


if __name__ == '__main__':
    print("Audio Quality Assessment Module")
    print("=" * 50)
    print("Features:")
    print("  - AudioQualityAssessor: Assess audio quality")
    print("  - AudioDatasetValidator: Validate entire dataset")
    print()
    print("Example usage:")
    print("  assessor = AudioQualityAssessor()")
    print("  is_valid, assessment = assessor.validate_audio_file('audio.wav')")