"""
Baby Cry AI - Dataset Explorer
Step 2: Explore and analyze the available datasets
"""

import os
import librosa
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
import json

class BabyCryDataExplorer:
    def __init__(self, data_path="data"):
        self.data_path = Path(data_path)
        self.categories = []
        self.file_info = {}
        
    def explore_dataset(self):
        """Explore the structure and content of baby cry datasets"""
        print("🔍 Exploring Baby Cry Dataset...")
        print("=" * 50)
        
        if not self.data_path.exists():
            print(f"❌ Dataset path not found: {self.data_path}")
            print("📥 Please download datasets first. See README.md for instructions.")
            return False
            
        # Get all categories
        self.categories = [d for d in os.listdir(self.data_path) 
                          if os.path.isdir(os.path.join(self.data_path, d))]
        
        if not self.categories:
            print("❌ No categories found in dataset directory")
            return False
        
        print(f"📁 Found {len(self.categories)} categories:")
        for category in self.categories:
            category_path = self.data_path / category
            files = [f for f in os.listdir(category_path) if f.endswith(('.wav', '.mp3', '.m4a', '.flac'))]
            self.file_info[category] = {
                'count': len(files),
                'files': files[:5]  # Store first 5 files for analysis
            }
            print(f"  • {category}: {len(files)} audio files")
            
        return True
    
    def analyze_audio_features(self):
        """Analyze audio features of sample files"""
        print("\n🎵 Analyzing Audio Features...")
        print("=" * 50)
        
        features_summary = {}
        
        for category in self.categories[:3]:  # Analyze first 3 categories
            if self.file_info[category]['files']:
                sample_file = self.file_info[category]['files'][0]
                file_path = self.data_path / category / sample_file
                
                try:
                    # Load audio
                    y, sr = librosa.load(file_path)
                    
                    # Extract features
                    duration = len(y) / sr
                    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
                    spectral_centroids = librosa.feature.spectral_centroid(y=y, sr=sr)
                    zero_crossing_rate = librosa.feature.zero_crossing_rate(y)
                    rms = librosa.feature.rms(y=y)[0]
                    
                    features_summary[category] = {
                        'duration': duration,
                        'sample_rate': sr,
                        'mfcc_mean': np.mean(mfccs, axis=1).tolist(),
                        'spectral_centroid_mean': float(np.mean(spectral_centroids)),
                        'zcr_mean': float(np.mean(zero_crossing_rate)),
                        'rms_mean': float(np.mean(rms))
                    }
                    
                    print(f"\n📊 {category.upper()} Sample Analysis:")
                    print(f"  • Duration: {duration:.2f} seconds")
                    print(f"  • Sample Rate: {sr} Hz")
                    print(f"  • MFCC Shape: {mfccs.shape}")
                    print(f"  • Spectral Centroid Mean: {np.mean(spectral_centroids):.2f}")
                    print(f"  • Zero Crossing Rate: {np.mean(zero_crossing_rate):.4f}")
                    print(f"  • RMS Energy: {np.mean(rms):.4f}")
                    
                except Exception as e:
                    print(f"❌ Error analyzing {category}: {e}")
        
        return features_summary
    
    def create_data_summary(self):
        """Create a summary of the dataset"""
        total_files = sum(info['count'] for info in self.file_info.values())
        
        summary = {
            'total_categories': len(self.categories),
            'total_files': total_files,
            'categories': self.categories,
            'files_per_category': {cat: info['count'] for cat, info in self.file_info.items()},
            'average_files_per_category': total_files / len(self.categories) if self.categories else 0
        }
        
        print(f"\n📈 Dataset Summary:")
        print(f"  • Total Categories: {summary['total_categories']}")
        print(f"  • Total Files: {summary['total_files']}")
        print(f"  • Average Files per Category: {summary['average_files_per_category']:.1f}")
        
        return summary
    
    def save_analysis_report(self, filename="data_analysis_report.json"):
        """Save analysis results to a JSON file"""
        summary = self.create_data_summary()
        features = self.analyze_audio_features()
        
        report = {
            'summary': summary,
            'audio_features': features,
            'timestamp': pd.Timestamp.now().isoformat()
        }
        
        with open(filename, 'w') as f:
            json.dump(report, f, indent=2)
        
        print(f"\n💾 Analysis report saved to: {filename}")
        return report

if __name__ == "__main__":
    # Initialize explorer
    explorer = BabyCryDataExplorer("data")
    
    # Explore dataset
    if explorer.explore_dataset():
        features = explorer.analyze_audio_features()
        summary = explorer.create_data_summary()
        report = explorer.save_analysis_report()
        
        print("\n✅ Dataset exploration complete!")
        print("Ready to proceed with model development.")
    else:
        print("\n❌ Dataset exploration failed. Please check your data directory.")