src/data_quality_analyzer.py

"""
Baby Cry AI - Data Quality Analyzer
Analyzes audio quality metrics, identifies corrupted/duplicate files, and generates quality reports
"""

import os
import sys
import json
import hashlib
import numpy as np
import librosa
from pathlib import Path
from collections import defaultdict
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from audio_processor import AudioProcessor


class DataQualityAnalyzer:
    """Analyze data quality and identify problematic files"""
    
    def __init__(self, data_dir="../data", min_snr_db=10.0):
        """
        Initialize analyzer
        
        Args:
            data_dir: Path to data directory
            min_snr_db: Minimum acceptable SNR in dB
        """
        self.data_dir = Path(data_dir)
        self.min_snr_db = min_snr_db
        self.processor = AudioProcessor()
        
        # Quality thresholds
        self.thresholds = {
            'min_duration': 3.0,
            'max_duration': 30.0,
            'min_audio_level_db': -40,
            'max_silence_ratio': 0.80,
            'min_snr_db': min_snr_db,
            'min_cry_likelihood': 0.3
        }
        
        # Results storage
        self.results = {
            'total_files': 0,
            'valid_files': 0,
            'corrupted_files': [],
            'low_quality_files': [],
            'duplicate_files': [],
            'short_files': [],
            'long_files': [],
            'quiet_files': [],
            'silent_files': [],
            'low_snr_files': [],
            'category_stats': defaultdict(lambda: {
                'total': 0,
                'valid': 0,
                'corrupted': 0,
                'low_quality': 0,
                'duplicates': 0
            })
        }
        
        # For duplicate detection
        self.audio_hashes = {}
        self.feature_hashes = {}
    
    def calculate_snr(self, y, sr):
        """
        Calculate Signal-to-Noise Ratio (SNR) in dB
        
        Uses spectral subtraction approach:
        - Signal: energy in typical cry frequency range (300-3000 Hz)
        - Noise: energy in low frequencies (< 100 Hz) and high frequencies (> 5000 Hz)
        """
        try:
            # Compute power spectral density
            freqs = np.fft.rfftfreq(len(y), 1/sr)
            fft = np.fft.rfft(y)
            psd = np.abs(fft) ** 2
            
            # Signal: 300-3000 Hz (typical cry range)
            signal_mask = (freqs >= 300) & (freqs <= 3000)
            signal_power = np.sum(psd[signal_mask])
            
            # Noise: < 100 Hz and > 5000 Hz
            noise_mask = (freqs < 100) | (freqs > 5000)
            noise_power = np.sum(psd[noise_mask])
            
            # Avoid division by zero
            if noise_power < 1e-10:
                noise_power = 1e-10
            
            snr_linear = signal_power / noise_power
            snr_db = 10 * np.log10(snr_linear)
            
            return float(snr_db)
        except Exception:
            return -100.0  # Invalid SNR
    
    def calculate_audio_hash(self, file_path):
        """Calculate hash of audio file for duplicate detection"""
        try:
            # Use file size and first few seconds for quick hash
            stat = os.stat(file_path)
            file_size = stat.st_size
            
            # Load first 2 seconds for content hash
            y, sr = librosa.load(file_path, sr=None, duration=2.0)
            if y is None or len(y) == 0:
                return None
            
            # Create hash from file size and audio content
            content_hash = hashlib.md5(y.tobytes()).hexdigest()
            combined = f"{file_size}_{content_hash}"
            return hashlib.md5(combined.encode()).hexdigest()
        except Exception:
            return None
    
    def calculate_feature_hash(self, features):
        """Calculate hash of features for similarity detection"""
        try:
            # Convert features to string representation
            feature_str = json.dumps(features, sort_keys=True)
            return hashlib.md5(feature_str.encode()).hexdigest()
        except Exception:
            return None
    
    def analyze_file(self, file_path, category):
        """Analyze a single audio file"""
        file_path = Path(file_path)
        self.results['total_files'] += 1
        self.results['category_stats'][category]['total'] += 1
        
        issues = []
        quality_metrics = {}
        
        # Check if file exists and is readable
        if not file_path.exists():
            issues.append('file_not_found')
            self.results['corrupted_files'].append({
                'file': str(file_path),
                'category': category,
                'issue': 'file_not_found'
            })
            self.results['category_stats'][category]['corrupted'] += 1
            return False, issues, quality_metrics
        
        # Try to load audio
        try:
            y, sr = librosa.load(str(file_path), sr=None, duration=None)
            if y is None or len(y) == 0:
                raise ValueError("Empty audio")
        except Exception as e:
            issues.append('corrupted')
            self.results['corrupted_files'].append({
                'file': str(file_path),
                'category': category,
                'issue': 'corrupted',
                'error': str(e)
            })
            self.results['category_stats'][category]['corrupted'] += 1
            return False, issues, quality_metrics
        
        # Get audio stats
        stats = self.processor.get_audio_stats(str(file_path))
        if not stats.get('valid', False):
            issues.append('invalid_stats')
            self.results['corrupted_files'].append({
                'file': str(file_path),
                'category': category,
                'issue': 'invalid_stats',
                'error': stats.get('error', 'Unknown')
            })
            self.results['category_stats'][category]['corrupted'] += 1
            return False, issues, quality_metrics
        
        # Extract quality metrics
        duration = stats.get('duration_seconds', 0)
        audio_level_db = stats.get('audio_level_db', -100)
        silence_ratio = stats.get('silence_ratio', 1.0)
        cry_likelihood = stats.get('cry_likelihood', 0.0)
        
        # Calculate SNR
        snr_db = self.calculate_snr(y, sr)
        
        quality_metrics = {
            'duration': duration,
            'audio_level_db': audio_level_db,
            'silence_ratio': silence_ratio,
            'snr_db': snr_db,
            'cry_likelihood': cry_likelihood,
            'sample_rate': sr
        }
        
        # Check duration
        if duration < self.thresholds['min_duration']:
            issues.append('too_short')
            self.results['short_files'].append({
                'file': str(file_path),
                'category': category,
                'duration': duration
            })
        
        if duration > self.thresholds['max_duration']:
            issues.append('too_long')
            self.results['long_files'].append({
                'file': str(file_path),
                'category': category,
                'duration': duration
            })
        
        # Check audio level
        if audio_level_db < self.thresholds['min_audio_level_db']:
            issues.append('too_quiet')
            self.results['quiet_files'].append({
                'file': str(file_path),
                'category': category,
                'audio_level_db': audio_level_db
            })
        
        # Check silence ratio
        if silence_ratio > self.thresholds['max_silence_ratio']:
            issues.append('too_silent')
            self.results['silent_files'].append({
                'file': str(file_path),
                'category': category,
                'silence_ratio': silence_ratio
            })
        
        # Check SNR
        if snr_db < self.thresholds['min_snr_db']:
            issues.append('low_snr')
            self.results['low_snr_files'].append({
                'file': str(file_path),
                'category': category,
                'snr_db': snr_db
            })
        
        # Check cry likelihood
        if cry_likelihood < self.thresholds['min_cry_likelihood']:
            issues.append('low_cry_likelihood')
        
        # Check for duplicates
        audio_hash = self.calculate_audio_hash(str(file_path))
        if audio_hash:
            if audio_hash in self.audio_hashes:
                issues.append('duplicate')
                duplicate_info = {
                    'file': str(file_path),
                    'category': category,
                    'duplicate_of': self.audio_hashes[audio_hash]
                }
                self.results['duplicate_files'].append(duplicate_info)
                self.results['category_stats'][category]['duplicates'] += 1
            else:
                self.audio_hashes[audio_hash] = str(file_path)
        
        # Determine if file is valid
        is_valid = len(issues) == 0
        
        if is_valid:
            self.results['valid_files'] += 1
            self.results['category_stats'][category]['valid'] += 1
        else:
            self.results['low_quality_files'].append({
                'file': str(file_path),
                'category': category,
                'issues': issues,
                'metrics': quality_metrics
            })
            self.results['category_stats'][category]['low_quality'] += 1
        
        return is_valid, issues, quality_metrics
    
    def analyze_dataset(self):
        """Analyze entire dataset"""
        print("🔍 Starting Data Quality Analysis...")
        print("=" * 60)
        
        if not self.data_dir.exists():
            print(f"❌ Data directory not found: {self.data_dir}")
            return
        
        # Get all categories
        categories = [d for d in os.listdir(self.data_dir) 
                     if os.path.isdir(self.data_dir / d)]
        
        print(f"\n📁 Found {len(categories)} categories: {categories}")
        
        # Analyze each category
        for category in categories:
            category_path = self.data_dir / category
            audio_files = [f for f in os.listdir(category_path) 
                          if f.endswith(('.wav', '.mp3', '.m4a', '.flac'))]
            
            print(f"\n📂 Analyzing {category}... ({len(audio_files)} files)")
            
            for filename in audio_files:
                file_path = category_path / filename
                self.analyze_file(file_path, category)
        
        # Generate summary
        self.generate_report()
    
    def generate_report(self):
        """Generate quality analysis report"""
        print("\n" + "=" * 60)
        print("📊 DATA QUALITY ANALYSIS REPORT")
        print("=" * 60)
        
        total = self.results['total_files']
        valid = self.results['valid_files']
        invalid = total - valid
        
        print(f"\n📈 Overall Statistics:")
        print(f"   Total files: {total}")
        print(f"   Valid files: {valid} ({valid/total*100:.1f}%)")
        print(f"   Invalid files: {invalid} ({invalid/total*100:.1f}%)")
        
        print(f"\n❌ Issues Found:")
        print(f"   Corrupted: {len(self.results['corrupted_files'])}")
        print(f"   Low quality: {len(self.results['low_quality_files'])}")
        print(f"   Duplicates: {len(self.results['duplicate_files'])}")
        print(f"   Too short: {len(self.results['short_files'])}")
        print(f"   Too long: {len(self.results['long_files'])}")
        print(f"   Too quiet: {len(self.results['quiet_files'])}")
        print(f"   Too silent: {len(self.results['silent_files'])}")
        print(f"   Low SNR: {len(self.results['low_snr_files'])}")
        
        print(f"\n📊 Per-Category Statistics:")
        for category, stats in sorted(self.results['category_stats'].items()):
            total_cat = stats['total']
            valid_cat = stats['valid']
            valid_pct = (valid_cat / total_cat * 100) if total_cat > 0 else 0
            print(f"   {category:15}: {valid_cat:4}/{total_cat:4} valid ({valid_pct:5.1f}%) | "
                  f"Corrupted: {stats['corrupted']:3} | "
                  f"Low quality: {stats['low_quality']:3} | "
                  f"Duplicates: {stats['duplicates']:3}")
        
        # Save detailed report
        report_path = self.data_dir.parent / 'data_quality_report.json'
        with open(report_path, 'w') as f:
            json.dump({
                'timestamp': datetime.now().isoformat(),
                'thresholds': self.thresholds,
                'summary': {
                    'total_files': total,
                    'valid_files': valid,
                    'invalid_files': invalid,
                    'valid_percentage': valid/total*100 if total > 0 else 0
                },
                'issues': {
                    'corrupted': len(self.results['corrupted_files']),
                    'low_quality': len(self.results['low_quality_files']),
                    'duplicates': len(self.results['duplicate_files']),
                    'short': len(self.results['short_files']),
                    'long': len(self.results['long_files']),
                    'quiet': len(self.results['quiet_files']),
                    'silent': len(self.results['silent_files']),
                    'low_snr': len(self.results['low_snr_files'])
                },
                'category_stats': dict(self.results['category_stats']),
                'corrupted_files': self.results['corrupted_files'][:100],  # Limit to first 100
                'low_quality_files': self.results['low_quality_files'][:100],
                'duplicate_files': self.results['duplicate_files'][:100]
            }, f, indent=2)
        
        print(f"\n💾 Detailed report saved to: {report_path}")
        
        # Recommendations
        print(f"\n💡 Recommendations:")
        if len(self.results['corrupted_files']) > 0:
            print(f"   - Remove {len(self.results['corrupted_files'])} corrupted files")
        if len(self.results['duplicate_files']) > 0:
            print(f"   - Remove {len(self.results['duplicate_files'])} duplicate files")
        if len(self.results['low_quality_files']) > 0:
            print(f"   - Review {len(self.results['low_quality_files'])} low-quality files")
        if len(self.results['short_files']) > 0:
            print(f"   - Consider removing {len(self.results['short_files'])} files that are too short")
        if len(self.results['low_snr_files']) > 0:
            print(f"   - Review {len(self.results['low_snr_files'])} files with low SNR")
        
        print("\n" + "=" * 60)
    
    def get_files_to_remove(self):
        """Get list of files recommended for removal"""
        files_to_remove = []
        
        # Corrupted files
        files_to_remove.extend([f['file'] for f in self.results['corrupted_files']])
        
        # Duplicate files (keep first, remove others)
        seen_hashes = set()
        for dup in self.results['duplicate_files']:
            if dup.get('duplicate_of'):
                files_to_remove.append(dup['file'])
        
        return list(set(files_to_remove))  # Remove duplicates from list
    
    def get_low_quality_files(self):
        """Get list of low-quality files for review"""
        return self.results['low_quality_files']


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Analyze data quality')
    parser.add_argument('--data-dir', type=str, default='../data',
                       help='Path to data directory')
    parser.add_argument('--min-snr', type=float, default=10.0,
                       help='Minimum acceptable SNR in dB')
    
    args = parser.parse_args()
    
    analyzer = DataQualityAnalyzer(data_dir=args.data_dir, min_snr_db=args.min_snr)
    analyzer.analyze_dataset()
    
    # Show files to remove
    files_to_remove = analyzer.get_files_to_remove()
    if files_to_remove:
        print(f"\n🗑️  Files recommended for removal: {len(files_to_remove)}")
        print("   (Use data_cleaner.py to remove them safely)")