src/data_cleaner.py

"""
Baby Cry AI - Data Cleaning Pipeline
Removes corrupted, duplicate, and low-quality audio files safely
"""

import os
import sys
import shutil
import json
from pathlib import Path
from datetime import datetime
import argparse

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

from data_quality_analyzer import DataQualityAnalyzer


class DataCleaner:
    """Clean dataset by removing problematic files"""
    
    def __init__(self, data_dir="../data", backup_dir=None, dry_run=True):
        """
        Initialize cleaner
        
        Args:
            data_dir: Path to data directory
            backup_dir: Path to backup directory (default: data_dir/../data_backup_TIMESTAMP)
            dry_run: If True, only show what would be removed without actually removing
        """
        self.data_dir = Path(data_dir)
        self.dry_run = dry_run
        
        if backup_dir:
            self.backup_dir = Path(backup_dir)
        else:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            self.backup_dir = self.data_dir.parent / f"data_backup_{timestamp}"
        
        self.removed_files = []
        self.backed_up_files = []
        
    def create_backup(self):
        """Create backup of data directory before cleaning"""
        if not self.dry_run:
            print(f"📦 Creating backup to: {self.backup_dir}")
            self.backup_dir.mkdir(parents=True, exist_ok=True)
            
            # Copy entire data directory structure
            for category_dir in self.data_dir.iterdir():
                if category_dir.is_dir():
                    backup_category = self.backup_dir / category_dir.name
                    backup_category.mkdir(exist_ok=True)
                    print(f"   Backing up {category_dir.name}...")
        else:
            print(f"📦 Would create backup to: {self.backup_dir}")
    
    def backup_file(self, file_path):
        """Backup a single file before removal"""
        if self.dry_run:
            return
        
        relative_path = file_path.relative_to(self.data_dir)
        backup_path = self.backup_dir / relative_path
        backup_path.parent.mkdir(parents=True, exist_ok=True)
        
        try:
            shutil.copy2(file_path, backup_path)
            self.backed_up_files.append(str(backup_path))
        except Exception as e:
            print(f"   ⚠️  Warning: Could not backup {file_path}: {e}")
    
    def remove_file(self, file_path, reason=""):
        """Remove a file (with backup)"""
        file_path = Path(file_path)
        
        if not file_path.exists():
            return False
        
        # Backup before removal
        self.backup_file(file_path)
        
        if self.dry_run:
            print(f"   [DRY RUN] Would remove: {file_path.name} ({reason})")
            self.removed_files.append({
                'file': str(file_path),
                'reason': reason,
                'removed': False
            })
            return True
        
        try:
            file_path.unlink()
            self.removed_files.append({
                'file': str(file_path),
                'reason': reason,
                'removed': True
            })
            return True
        except Exception as e:
            print(f"   ❌ Error removing {file_path}: {e}")
            return False
    
    def clean_corrupted_files(self, analyzer):
        """Remove corrupted files"""
        print("\n🗑️  Removing corrupted files...")
        count = 0
        
        for file_info in analyzer.results['corrupted_files']:
            file_path = Path(file_info['file'])
            if self.remove_file(file_path, f"corrupted: {file_info.get('issue', 'unknown')}"):
                count += 1
        
        print(f"   {'Would remove' if self.dry_run else 'Removed'}: {count} corrupted files")
        return count
    
    def clean_duplicate_files(self, analyzer):
        """Remove duplicate files (keep first occurrence)"""
        print("\n🗑️  Removing duplicate files...")
        count = 0
        
        # Track which files we've seen (keep first)
        seen_hashes = {}
        
        for dup_info in analyzer.results['duplicate_files']:
            file_path = Path(dup_info['file'])
            original = dup_info.get('duplicate_of')
            
            # Remove the duplicate (not the original)
            if original and str(file_path) != original:
                if self.remove_file(file_path, f"duplicate of {Path(original).name}"):
                    count += 1
        
        print(f"   {'Would remove' if self.dry_run else 'Removed'}: {count} duplicate files")
        return count
    
    def clean_short_files(self, analyzer, min_duration=3.0):
        """Remove files that are too short"""
        print(f"\n🗑️  Removing files shorter than {min_duration}s...")
        count = 0
        
        for file_info in analyzer.results['short_files']:
            file_path = Path(file_info['file'])
            duration = file_info.get('duration', 0)
            if self.remove_file(file_path, f"too short: {duration:.1f}s"):
                count += 1
        
        print(f"   {'Would remove' if self.dry_run else 'Removed'}: {count} short files")
        return count
    
    def clean_long_files(self, analyzer, max_duration=30.0):
        """Remove files that are too long"""
        print(f"\n🗑️  Removing files longer than {max_duration}s...")
        count = 0
        
        for file_info in analyzer.results['long_files']:
            file_path = Path(file_info['file'])
            duration = file_info.get('duration', 0)
            if self.remove_file(file_path, f"too long: {duration:.1f}s"):
                count += 1
        
        print(f"   {'Would remove' if self.dry_run else 'Removed'}: {count} long files")
        return count
    
    def clean_quiet_files(self, analyzer, min_level_db=-40):
        """Remove files that are too quiet"""
        print(f"\n🗑️  Removing files quieter than {min_level_db} dB...")
        count = 0
        
        for file_info in analyzer.results['quiet_files']:
            file_path = Path(file_info['file'])
            level = file_info.get('audio_level_db', -100)
            if self.remove_file(file_path, f"too quiet: {level:.1f} dB"):
                count += 1
        
        print(f"   {'Would remove' if self.dry_run else 'Removed'}: {count} quiet files")
        return count
    
    def clean_silent_files(self, analyzer, max_silence_ratio=0.80):
        """Remove files that are mostly silence"""
        print(f"\n🗑️  Removing files with >{max_silence_ratio*100:.0f}% silence...")
        count = 0
        
        for file_info in analyzer.results['silent_files']:
            file_path = Path(file_info['file'])
            silence = file_info.get('silence_ratio', 1.0)
            if self.remove_file(file_path, f"too silent: {silence*100:.0f}%"):
                count += 1
        
        print(f"   {'Would remove' if self.dry_run else 'Removed'}: {count} silent files")
        return count
    
    def clean_low_snr_files(self, analyzer, min_snr_db=10.0):
        """Remove files with low signal-to-noise ratio"""
        print(f"\n🗑️  Removing files with SNR < {min_snr_db} dB...")
        count = 0
        
        for file_info in analyzer.results['low_snr_files']:
            file_path = Path(file_info['file'])
            snr = file_info.get('snr_db', -100)
            if self.remove_file(file_path, f"low SNR: {snr:.1f} dB"):
                count += 1
        
        print(f"   {'Would remove' if self.dry_run else 'Removed'}: {count} low SNR files")
        return count
    
    def clean_all(self, analyzer, 
                  remove_corrupted=True,
                  remove_duplicates=True,
                  remove_short=True,
                  remove_long=True,
                  remove_quiet=True,
                  remove_silent=True,
                  remove_low_snr=False):  # Low SNR is optional, might be too aggressive
        """Run complete cleaning pipeline"""
        print("🧹 Starting Data Cleaning Pipeline")
        print("=" * 60)
        
        if self.dry_run:
            print("⚠️  DRY RUN MODE - No files will be removed")
        else:
            print("⚠️  LIVE MODE - Files will be permanently removed")
        
        # Create backup
        self.create_backup()
        
        total_removed = 0
        
        # Remove files based on options
        if remove_corrupted:
            total_removed += self.clean_corrupted_files(analyzer)
        
        if remove_duplicates:
            total_removed += self.clean_duplicate_files(analyzer)
        
        if remove_short:
            total_removed += self.clean_short_files(analyzer)
        
        if remove_long:
            total_removed += self.clean_long_files(analyzer)
        
        if remove_quiet:
            total_removed += self.clean_quiet_files(analyzer)
        
        if remove_silent:
            total_removed += self.clean_silent_files(analyzer)
        
        if remove_low_snr:
            total_removed += self.clean_low_snr_files(analyzer)
        
        # Save cleaning report
        self.save_report()
        
        print("\n" + "=" * 60)
        print(f"✅ Cleaning complete!")
        print(f"   {'Would remove' if self.dry_run else 'Removed'}: {total_removed} files")
        if not self.dry_run:
            print(f"   Backed up to: {self.backup_dir}")
        print("=" * 60)
        
        return total_removed
    
    def save_report(self):
        """Save cleaning report"""
        report_path = self.data_dir.parent / 'data_cleaning_report.json'
        
        report = {
            'timestamp': datetime.now().isoformat(),
            'dry_run': self.dry_run,
            'backup_dir': str(self.backup_dir),
            'total_removed': len(self.removed_files),
            'removed_files': self.removed_files[:500],  # Limit to first 500
            'backed_up_files': len(self.backed_up_files)
        }
        
        with open(report_path, 'w') as f:
            json.dump(report, f, indent=2)
        
        print(f"\n💾 Cleaning report saved to: {report_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Clean dataset by removing problematic files')
    parser.add_argument('--data-dir', type=str, default='../data',
                       help='Path to data directory')
    parser.add_argument('--backup-dir', type=str, default=None,
                       help='Path to backup directory')
    parser.add_argument('--execute', action='store_true',
                       help='Actually remove files (default is dry-run)')
    parser.add_argument('--min-snr', type=float, default=10.0,
                       help='Minimum acceptable SNR in dB')
    parser.add_argument('--remove-low-snr', action='store_true',
                       help='Also remove files with low SNR')
    
    args = parser.parse_args()
    
    # First, analyze data quality
    print("Step 1: Analyzing data quality...")
    analyzer = DataQualityAnalyzer(data_dir=args.data_dir, min_snr_db=args.min_snr)
    analyzer.analyze_dataset()
    
    # Then, clean based on analysis
    print("\nStep 2: Cleaning data...")
    cleaner = DataCleaner(
        data_dir=args.data_dir,
        backup_dir=args.backup_dir,
        dry_run=not args.execute
    )
    
    cleaner.clean_all(
        analyzer,
        remove_low_snr=args.remove_low_snr
    )
    
    if cleaner.dry_run:
        print("\n💡 To actually remove files, run with --execute flag")