#!/usr/bin/env python3
"""
Dataset Analysis Script for Voice Cloning Dataset
Developer: Molla Samser
Email: help@rskworld.in
Phone: +91 93305 39277
Website: https://rskworld.in
Year: 2026
"""

import os
import json
import argparse
from pathlib import Path
from collections import defaultdict
import librosa
import numpy as np


def analyze_audio_file(audio_path):
    """Analyze a single audio file and return statistics."""
    try:
        audio, sr = librosa.load(audio_path, sr=None, mono=True)
        duration = librosa.get_duration(y=audio, sr=sr)
        
        # Calculate statistics
        stats = {
            'file': str(audio_path),
            'sample_rate': int(sr),
            'duration': float(duration),
            'samples': len(audio),
            'rms_energy': float(np.sqrt(np.mean(audio**2))),
            'zero_crossing_rate': float(np.mean(librosa.feature.zero_crossing_rate(audio)[0])),
            'spectral_centroid': float(np.mean(librosa.feature.spectral_centroid(y=audio, sr=sr)[0])),
        }
        
        # Extract pitch if possible
        try:
            pitches, magnitudes = librosa.piptrack(y=audio, sr=sr)
            pitch_values = []
            for t in range(pitches.shape[1]):
                index = magnitudes[:, t].argmax()
                pitch = pitches[index, t]
                if pitch > 0:
                    pitch_values.append(pitch)
            if pitch_values:
                stats['pitch_mean'] = float(np.mean(pitch_values))
                stats['pitch_std'] = float(np.std(pitch_values))
        except:
            pass
        
        return stats
    except Exception as e:
        return {'file': str(audio_path), 'error': str(e)}


def analyze_dataset(dataset_dir):
    """Analyze entire dataset and generate statistics."""
    dataset_path = Path(dataset_dir)
    analysis = {
        'dataset_path': str(dataset_path),
        'speakers': {},
        'overall_statistics': {
            'total_files': 0,
            'total_duration': 0.0,
            'sample_rates': defaultdict(int),
            'durations': [],
            'formats': defaultdict(int)
        }
    }
    
    # Analyze each speaker
    for speaker_dir in dataset_path.iterdir():
        if speaker_dir.is_dir() and speaker_dir.name.startswith('speaker_'):
            speaker_id = speaker_dir.name
            speaker_stats = {
                'speaker_id': speaker_id,
                'files': [],
                'total_duration': 0.0,
                'file_count': 0
            }
            
            recordings_dir = speaker_dir / 'recordings'
            if recordings_dir.exists():
                audio_extensions = ['.wav', '.flac', '.mp3']
                for ext in audio_extensions:
                    for audio_file in recordings_dir.rglob(f'*{ext}'):
                        file_stats = analyze_audio_file(audio_file)
                        if 'error' not in file_stats:
                            speaker_stats['files'].append(file_stats)
                            speaker_stats['total_duration'] += file_stats['duration']
                            speaker_stats['file_count'] += 1
                            
                            # Update overall statistics
                            analysis['overall_statistics']['total_files'] += 1
                            analysis['overall_statistics']['total_duration'] += file_stats['duration']
                            analysis['overall_statistics']['sample_rates'][file_stats['sample_rate']] += 1
                            analysis['overall_statistics']['durations'].append(file_stats['duration'])
                            analysis['overall_statistics']['formats'][ext[1:].upper()] += 1
            
            analysis['speakers'][speaker_id] = speaker_stats
    
    # Calculate summary statistics
    durations = analysis['overall_statistics']['durations']
    if durations:
        analysis['overall_statistics']['duration_stats'] = {
            'min': float(np.min(durations)),
            'max': float(np.max(durations)),
            'mean': float(np.mean(durations)),
            'median': float(np.median(durations)),
            'std': float(np.std(durations))
        }
        analysis['overall_statistics']['total_duration_hours'] = float(
            analysis['overall_statistics']['total_duration'] / 3600
        )
        analysis['overall_statistics']['total_duration_minutes'] = float(
            analysis['overall_statistics']['total_duration'] / 60
        )
    
    return analysis


def main():
    parser = argparse.ArgumentParser(description='Analyze voice cloning dataset')
    parser.add_argument('--dataset', type=str, default='audio',
                       help='Dataset directory (default: audio)')
    parser.add_argument('--output', type=str, default='dataset_analysis.json',
                       help='Output JSON file (default: dataset_analysis.json)')
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("Voice Cloning Dataset - Analysis")
    print("Developer: Molla Samser")
    print("Website: https://rskworld.in")
    print("=" * 60)
    
    # Analyze dataset
    print(f"\nAnalyzing dataset in {args.dataset}...")
    analysis = analyze_dataset(args.dataset)
    
    # Print summary
    stats = analysis['overall_statistics']
    print(f"\nDataset Analysis Summary:")
    print(f"  Total Files: {stats['total_files']}")
    print(f"  Total Duration: {stats.get('total_duration_hours', 0):.2f} hours")
    print(f"  Total Duration: {stats.get('total_duration_minutes', 0):.2f} minutes")
    
    if stats.get('duration_stats'):
        ds = stats['duration_stats']
        print(f"\n  Duration Statistics:")
        print(f"    Min: {ds['min']:.2f}s")
        print(f"    Max: {ds['max']:.2f}s")
        print(f"    Mean: {ds['mean']:.2f}s")
        print(f"    Median: {ds['median']:.2f}s")
    
    print(f"\n  Sample Rates:")
    for sr, count in sorted(stats['sample_rates'].items()):
        print(f"    {sr} Hz: {count} files")
    
    print(f"\n  Formats:")
    for fmt, count in sorted(stats['formats'].items()):
        print(f"    {fmt}: {count} files")
    
    # Save analysis
    with open(args.output, 'w') as f:
        json.dump(analysis, f, indent=2)
    
    print(f"\nAnalysis saved to {args.output}")
    print("\nAnalysis complete!")


if __name__ == '__main__':
    main()