#!/usr/bin/env python3
"""
Feature Extraction 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 argparse
import json
import librosa
import numpy as np
from pathlib import Path
from tqdm import tqdm


def extract_mfcc(audio, sr, n_mfcc=13):
    """Extract MFCC features from audio."""
    mfccs = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=n_mfcc)
    return mfccs.T  # Transpose to get time frames as rows


def extract_mel_spectrogram(audio, sr, n_mels=80):
    """Extract mel spectrogram from audio."""
    mel_spec = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=n_mels)
    mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)
    return mel_spec_db.T


def extract_pitch(audio, sr):
    """Extract pitch (F0) from audio."""
    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)
    return np.mean(pitch_values) if pitch_values else 0.0


def extract_features(audio_path):
    """
    Extract comprehensive features from an audio file.
    
    Returns:
        Dictionary containing extracted features
    """
    try:
        # Load audio
        audio, sr = librosa.load(audio_path, sr=22050, mono=True)
        
        # Extract features
        features = {
            'file_path': str(audio_path),
            'duration': librosa.get_duration(y=audio, sr=sr),
            'sample_rate': sr,
            'pitch': float(extract_pitch(audio, sr)),
            'mfcc': extract_mfcc(audio, sr).tolist(),
            'mel_spectrogram': extract_mel_spectrogram(audio, sr).tolist(),
            '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])),
            'spectral_rolloff': float(np.mean(librosa.feature.spectral_rolloff(y=audio, sr=sr)[0])),
        }
        
        return features
    except Exception as e:
        print(f"Error extracting features from {audio_path}: {str(e)}")
        return None


def process_directory(input_dir, output_dir):
    """
    Extract features from all audio files in a directory.
    
    Args:
        input_dir: Input directory containing audio files
        output_dir: Output directory for feature files
    """
    input_path = Path(input_dir)
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Supported audio formats
    audio_extensions = ['.wav', '.flac', '.mp3']
    
    # Find all audio files
    audio_files = []
    for ext in audio_extensions:
        audio_files.extend(input_path.rglob(f'*{ext}'))
    
    print(f"Found {len(audio_files)} audio files to process")
    
    # Extract features from each file
    all_features = {}
    for audio_file in tqdm(audio_files, desc="Extracting features"):
        features = extract_features(audio_file)
        if features:
            # Save individual feature file
            relative_path = audio_file.relative_to(input_path)
            feature_file = output_path / relative_path.with_suffix('.json')
            feature_file.parent.mkdir(parents=True, exist_ok=True)
            
            with open(feature_file, 'w') as f:
                json.dump(features, f, indent=2)
            
            # Store in all_features dict
            all_features[str(relative_path)] = {
                'duration': features['duration'],
                'pitch': features['pitch'],
                'sample_rate': features['sample_rate']
            }
    
    # Save summary file
    summary_file = output_path / 'features_summary.json'
    with open(summary_file, 'w') as f:
        json.dump(all_features, f, indent=2)
    
    print(f"\nExtracted features from {len(all_features)} files")
    print(f"Summary saved to {summary_file}")


def main():
    parser = argparse.ArgumentParser(description='Extract features from audio files')
    parser.add_argument('--audio', type=str, required=True, help='Input directory containing audio files')
    parser.add_argument('--output', type=str, required=True, help='Output directory for feature files')
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("Voice Cloning Dataset - Feature Extraction")
    print("Developer: Molla Samser")
    print("Website: https://rskworld.in")
    print("=" * 60)
    
    process_directory(args.audio, args.output)
    
    print("\nFeature extraction complete!")


if __name__ == '__main__':
    main()