Trixy3/trainer/data_pipeline.py

"""
Data pipeline management and preprocessing for ML trainers.

This module provides comprehensive data handling including audio preprocessing,
data augmentation, dataset creation, and batch loading with support for
various audio formats and training requirements.
"""

import os
import logging
import random
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Union, Callable, Any
from dataclasses import dataclass
import warnings

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, random_split, WeightedRandomSampler
import torchaudio
import torchaudio.transforms as T
from sklearn.model_selection import train_test_split

# Suppress torchaudio warnings
warnings.filterwarnings("ignore", category=UserWarning, module="torchaudio")


@dataclass
class AudioProcessingConfig:
    """Configuration for audio processing."""
    sample_rate: int = 16000
    target_length: float = 1.5  # seconds
    n_mels: int = 40
    n_fft: int = 512
    hop_length: int = 160  # 10ms at 16kHz
    win_length: int = 400  # 25ms at 16kHz
    fmin: float = 80.0
    fmax: float = 8000.0
    power: float = 2.0
    normalized: bool = True
    feature_type: str = "mel_spectrogram"  # mel_spectrogram, mfcc, raw
    
    @property
    def target_samples(self) -> int:
        """Get target number of samples."""
        return int(self.sample_rate * self.target_length)


class AudioProcessor:
    """Audio processing utilities for the training pipeline."""
    
    def __init__(self, config: AudioProcessingConfig, device: str = "cpu"):
        """
        Initialize audio processor.
        
        Args:
            config: Audio processing configuration
            device: Device for processing ('cpu', 'cuda')
        """
        self.config = config
        self.device = device
        
        # Initialize transforms
        self.mel_spectrogram = T.MelSpectrogram(
            sample_rate=config.sample_rate,
            n_fft=config.n_fft,
            hop_length=config.hop_length,
            win_length=config.win_length,
            n_mels=config.n_mels,
            fmin=config.fmin,
            fmax=config.fmax,
            power=config.power,
            normalized=config.normalized
        ).to(device)
        
        self.mfcc = T.MFCC(
            sample_rate=config.sample_rate,
            n_mfcc=config.n_mels,
            melkwargs={
                'n_fft': config.n_fft,
                'hop_length': config.hop_length,
                'win_length': config.win_length,
                'n_mels': config.n_mels,
                'fmin': config.fmin,
                'fmax': config.fmax
            }
        ).to(device)
        
        self.amplitude_to_db = T.AmplitudeToDB().to(device)
        
    def load_audio(self, filepath: str) -> Tuple[torch.Tensor, int]:
        """
        Load audio file and return waveform and sample rate.
        
        Args:
            filepath: Path to audio file
            
        Returns:
            Tuple of (waveform, sample_rate)
        """
        try:
            waveform, sr = torchaudio.load(filepath)
            return waveform, sr
        except Exception as e:
            raise ValueError(f"Error loading audio file {filepath}: {str(e)}")
    
    def resample(self, waveform: torch.Tensor, orig_sr: int) -> torch.Tensor:
        """
        Resample audio to target sample rate.
        
        Args:
            waveform: Input waveform
            orig_sr: Original sample rate
            
        Returns:
            Resampled waveform
        """
        if orig_sr != self.config.sample_rate:
            resampler = T.Resample(orig_sr, self.config.sample_rate)
            waveform = resampler(waveform)
        return waveform
    
    def normalize_length(self, waveform: torch.Tensor, 
                        method: str = "pad_trim") -> torch.Tensor:
        """
        Normalize audio length to target duration.
        
        Args:
            waveform: Input waveform
            method: Method for length normalization ('pad_trim', 'stretch', 'random_crop')
            
        Returns:
            Length-normalized waveform
        """
        target_samples = self.config.target_samples
        current_samples = waveform.shape[-1]
        
        if method == "pad_trim":
            if current_samples > target_samples:
                # Trim from center
                start = (current_samples - target_samples) // 2
                waveform = waveform[..., start:start + target_samples]
            elif current_samples < target_samples:
                # Pad with zeros
                pad_amount = target_samples - current_samples
                pad_left = pad_amount // 2
                pad_right = pad_amount - pad_left
                waveform = F.pad(waveform, (pad_left, pad_right))
        
        elif method == "random_crop":
            if current_samples > target_samples:
                # Random crop
                start = random.randint(0, current_samples - target_samples)
                waveform = waveform[..., start:start + target_samples]
            elif current_samples < target_samples:
                # Pad with zeros
                pad_amount = target_samples - current_samples
                waveform = F.pad(waveform, (0, pad_amount))
        
        elif method == "stretch":
            # Time-stretch to exact length (may change pitch)
            if current_samples != target_samples:
                stretch_factor = target_samples / current_samples
                # Simple linear interpolation for stretching
                indices = torch.linspace(0, current_samples - 1, target_samples)
                waveform = torch.nn.functional.interpolate(
                    waveform.unsqueeze(0), 
                    size=target_samples, 
                    mode='linear', 
                    align_corners=True
                ).squeeze(0)
        
        return waveform
    
    def extract_features(self, waveform: torch.Tensor) -> torch.Tensor:
        """
        Extract features from waveform.
        
        Args:
            waveform: Input waveform
            
        Returns:
            Feature tensor
        """
        waveform = waveform.to(self.device)
        
        if self.config.feature_type == "mel_spectrogram":
            features = self.mel_spectrogram(waveform)
            features = self.amplitude_to_db(features)
        elif self.config.feature_type == "mfcc":
            features = self.mfcc(waveform)
        elif self.config.feature_type == "raw":
            features = waveform
        else:
            raise ValueError(f"Unknown feature type: {self.config.feature_type}")
        
        return features
    
    def process_audio(self, filepath: str, 
                     normalize_method: str = "pad_trim") -> torch.Tensor:
        """
        Complete audio processing pipeline.
        
        Args:
            filepath: Path to audio file
            normalize_method: Method for length normalization
            
        Returns:
            Processed feature tensor
        """
        # Load audio
        waveform, sr = self.load_audio(filepath)
        
        # Convert to mono if stereo
        if waveform.shape[0] > 1:
            waveform = torch.mean(waveform, dim=0, keepdim=True)
        
        # Resample
        waveform = self.resample(waveform, sr)
        
        # Normalize length
        waveform = self.normalize_length(waveform, normalize_method)
        
        # Extract features
        features = self.extract_features(waveform)
        
        return features
    
    def trim_silence(self, waveform: torch.Tensor, 
                    threshold: float = 0.01) -> torch.Tensor:
        """
        Trim silence from beginning and end of audio.
        
        Args:
            waveform: Input waveform
            threshold: Silence threshold (relative to max amplitude)
            
        Returns:
            Trimmed waveform
        """
        # Calculate energy
        energy = torch.abs(waveform)
        max_energy = torch.max(energy)
        
        # Find non-silent regions
        non_silent = energy > (threshold * max_energy)
        
        if torch.any(non_silent):
            # Find first and last non-silent samples
            non_silent_indices = torch.where(non_silent.squeeze())[0]
            start_idx = non_silent_indices[0].item()
            end_idx = non_silent_indices[-1].item() + 1
            
            return waveform[..., start_idx:end_idx]
        else:
            # If all silent, return original
            return waveform


class AudioAugmentation:
    """Audio data augmentation utilities."""
    
    def __init__(self, sample_rate: int = 16000, device: str = "cpu"):
        """Initialize audio augmentation."""
        self.sample_rate = sample_rate
        self.device = device
    
    def add_noise(self, waveform: torch.Tensor, 
                  noise_factor: float = 0.1) -> torch.Tensor:
        """Add random noise to audio."""
        noise = torch.randn_like(waveform) * noise_factor
        return waveform + noise
    
    def time_shift(self, waveform: torch.Tensor, 
                  shift_factor: float = 0.1) -> torch.Tensor:
        """Apply random time shift."""
        shift_amount = int(waveform.shape[-1] * shift_factor * (2 * random.random() - 1))
        return torch.roll(waveform, shift_amount, dims=-1)
    
    def speed_change(self, waveform: torch.Tensor, 
                    speed_factor: float = 0.1) -> torch.Tensor:
        """Change playback speed (affects both tempo and pitch)."""
        speed_change = 1.0 + speed_factor * (2 * random.random() - 1)
        
        # Simple resampling for speed change
        new_length = int(waveform.shape[-1] / speed_change)
        indices = torch.linspace(0, waveform.shape[-1] - 1, new_length)
        
        # Linear interpolation
        waveform_resampled = torch.nn.functional.interpolate(
            waveform.unsqueeze(0), 
            size=new_length, 
            mode='linear', 
            align_corners=True
        ).squeeze(0)
        
        return waveform_resampled
    
    def pitch_shift(self, waveform: torch.Tensor, 
                   pitch_factor: float = 0.1) -> torch.Tensor:
        """Apply pitch shift (simplified version)."""
        # This is a simplified pitch shift - in practice, you'd use more sophisticated methods
        pitch_change = 1.0 + pitch_factor * (2 * random.random() - 1)
        
        # Create a simple frequency domain shift
        fft = torch.fft.fft(waveform)
        shifted_fft = torch.zeros_like(fft)
        
        n = fft.shape[-1]
        shift_bins = int(n * (pitch_change - 1.0))
        
        if shift_bins > 0:
            shifted_fft[..., shift_bins:] = fft[..., :-shift_bins]
        elif shift_bins < 0:
            shifted_fft[..., :shift_bins] = fft[..., -shift_bins:]
        else:
            shifted_fft = fft
        
        return torch.fft.ifft(shifted_fft).real
    
    def volume_change(self, waveform: torch.Tensor, 
                     volume_factor: float = 0.2) -> torch.Tensor:
        """Apply random volume change."""
        volume_change = 1.0 + volume_factor * (2 * random.random() - 1)
        return waveform * volume_change
    
    def apply_augmentations(self, waveform: torch.Tensor, 
                          augmentation_config: Dict[str, float]) -> torch.Tensor:
        """Apply multiple augmentations based on configuration."""
        augmented = waveform.clone()
        
        if augmentation_config.get('noise_factor', 0) > 0:
            if random.random() < 0.5:  # 50% chance
                augmented = self.add_noise(augmented, augmentation_config['noise_factor'])
        
        if augmentation_config.get('speed_factor', 0) > 0:
            if random.random() < 0.3:  # 30% chance
                augmented = self.speed_change(augmented, augmentation_config['speed_factor'])
        
        if augmentation_config.get('pitch_factor', 0) > 0:
            if random.random() < 0.3:  # 30% chance
                augmented = self.pitch_shift(augmented, augmentation_config['pitch_factor'])
        
        if augmentation_config.get('volume_factor', 0) > 0:
            if random.random() < 0.4:  # 40% chance
                augmented = self.volume_change(augmented, augmentation_config['volume_factor'])
        
        if augmentation_config.get('time_shift_factor', 0) > 0:
            if random.random() < 0.3:  # 30% chance
                augmented = self.time_shift(augmented, augmentation_config['time_shift_factor'])
        
        return augmented


class AudioDataset(Dataset):
    """Generic audio dataset for training."""
    
    def __init__(self, audio_files: List[str], labels: List[int], 
                 audio_processor: AudioProcessor,
                 augmentation: Optional[AudioAugmentation] = None,
                 augmentation_config: Optional[Dict[str, float]] = None,
                 apply_augmentation: bool = True):
        """
        Initialize audio dataset.
        
        Args:
            audio_files: List of audio file paths
            labels: List of corresponding labels
            audio_processor: Audio processor instance
            augmentation: Audio augmentation instance
            augmentation_config: Augmentation configuration
            apply_augmentation: Whether to apply augmentation during training
        """
        self.audio_files = audio_files
        self.labels = labels
        self.audio_processor = audio_processor
        self.augmentation = augmentation
        self.augmentation_config = augmentation_config or {}
        self.apply_augmentation = apply_augmentation
        
        assert len(audio_files) == len(labels), "Audio files and labels must have same length"
    
    def __len__(self) -> int:
        return len(self.audio_files)
    
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        """Get item from dataset."""
        audio_file = self.audio_files[idx]
        label = self.labels[idx]
        
        try:
            # Load and process audio
            waveform, sr = self.audio_processor.load_audio(audio_file)
            
            # Convert to mono
            if waveform.shape[0] > 1:
                waveform = torch.mean(waveform, dim=0, keepdim=True)
            
            # Resample
            waveform = self.audio_processor.resample(waveform, sr)
            
            # Apply augmentation if enabled and available
            if (self.apply_augmentation and self.augmentation is not None 
                and self.augmentation_config):
                waveform = self.augmentation.apply_augmentations(
                    waveform, self.augmentation_config
                )
            
            # Normalize length and extract features
            waveform = self.audio_processor.normalize_length(waveform)
            features = self.audio_processor.extract_features(waveform)
            
            # Ensure consistent shape
            if features.dim() == 3:  # (1, freq, time)
                features = features.squeeze(0)  # (freq, time)
            
            return features, label
            
        except Exception as e:
            # Return zero tensor and label on error
            logging.warning(f"Error processing {audio_file}: {str(e)}")
            dummy_features = torch.zeros(
                self.audio_processor.config.n_mels, 
                self.audio_processor.config.target_samples // self.audio_processor.config.hop_length + 1
            )
            return dummy_features, label


class DatasetBuilder:
    """Builder for creating training datasets from file structures."""
    
    def __init__(self, logger: Optional[logging.Logger] = None):
        """Initialize dataset builder."""
        self.logger = logger or logging.getLogger(__name__)
    
    def scan_directory_structure(self, data_dir: str, 
                                supported_extensions: List[str] = None) -> Dict[str, List[str]]:
        """
        Scan directory structure and organize files by class.
        
        Args:
            data_dir: Root data directory
            supported_extensions: List of supported file extensions
            
        Returns:
            Dictionary mapping class names to file lists
        """
        if supported_extensions is None:
            supported_extensions = ['.wav', '.mp3', '.flac', '.m4a']
        
        data_dir = Path(data_dir)
        if not data_dir.exists():
            raise FileNotFoundError(f"Data directory not found: {data_dir}")
        
        class_files = {}
        
        # For each subdirectory (class)
        for class_dir in data_dir.iterdir():
            if class_dir.is_dir():
                class_name = class_dir.name
                files = []
                
                # Find all audio files in class directory
                for ext in supported_extensions:
                    files.extend(list(class_dir.glob(f"**/*{ext}")))
                
                if files:
                    class_files[class_name] = [str(f) for f in files]
                    self.logger.info(f"Found {len(files)} files for class '{class_name}'")
        
        return class_files
    
    def create_wakeword_dataset(self, data_dir: str,
                               audio_processor: AudioProcessor,
                               wakeword_classes: Dict[int, str],
                               val_split: float = 0.2,
                               test_split: float = 0.1,
                               augmentation_config: Dict[str, float] = None) -> Tuple[AudioDataset, AudioDataset, AudioDataset]:
        """
        Create wakeword detection dataset.
        
        Args:
            data_dir: Data directory containing wakeword audio files
            audio_processor: Audio processor instance
            wakeword_classes: Mapping of class IDs to class names
            val_split: Validation split ratio
            test_split: Test split ratio
            augmentation_config: Augmentation configuration
            
        Returns:
            Tuple of (train_dataset, val_dataset, test_dataset)
        """
        # Scan directory structure
        class_files = self.scan_directory_structure(data_dir)
        
        # Map class names to IDs
        class_name_to_id = {name: class_id for class_id, name in wakeword_classes.items()}
        
        all_files = []
        all_labels = []
        
        for class_name, files in class_files.items():
            if class_name in class_name_to_id:
                class_id = class_name_to_id[class_name]
                all_files.extend(files)
                all_labels.extend([class_id] * len(files))
                self.logger.info(f"Added {len(files)} files for class '{class_name}' (ID: {class_id})")
            else:
                self.logger.warning(f"Unknown class '{class_name}' - skipping")
        
        if not all_files:
            raise ValueError("No valid audio files found")
        
        # Split data
        train_files, test_files, train_labels, test_labels = train_test_split(
            all_files, all_labels, test_size=test_split, random_state=42, stratify=all_labels
        )
        
        train_files, val_files, train_labels, val_labels = train_test_split(
            train_files, train_labels, test_size=val_split/(1-test_split), 
            random_state=42, stratify=train_labels
        )
        
        # Create augmentation
        augmentation = AudioAugmentation(audio_processor.config.sample_rate)
        
        # Create datasets
        train_dataset = AudioDataset(
            train_files, train_labels, audio_processor, 
            augmentation, augmentation_config, apply_augmentation=True
        )
        val_dataset = AudioDataset(
            val_files, val_labels, audio_processor,
            None, None, apply_augmentation=False
        )
        test_dataset = AudioDataset(
            test_files, test_labels, audio_processor,
            None, None, apply_augmentation=False
        )
        
        self.logger.info(f"Created datasets - Train: {len(train_dataset)}, "
                        f"Val: {len(val_dataset)}, Test: {len(test_dataset)}")
        
        return train_dataset, val_dataset, test_dataset
    
    def create_voice_recognition_dataset(self, data_dir: str,
                                       audio_processor: AudioProcessor,
                                       speaker_mapping: Dict[str, str],
                                       val_split: float = 0.2,
                                       test_split: float = 0.1,
                                       augmentation_config: Dict[str, float] = None) -> Tuple[AudioDataset, AudioDataset, AudioDataset, Dict[str, str]]:
        """
        Create voice recognition dataset.
        
        Args:
            data_dir: Data directory containing speaker audio files
            audio_processor: Audio processor instance
            speaker_mapping: Mapping of directory names to speaker names
            val_split: Validation split ratio
            test_split: Test split ratio
            augmentation_config: Augmentation configuration
            
        Returns:
            Tuple of (train_dataset, val_dataset, test_dataset, final_speaker_mapping)
        """
        # Scan directory structure
        class_files = self.scan_directory_structure(data_dir)
        
        # Create speaker ID mapping
        speaker_id_mapping = {}
        all_files = []
        all_labels = []
        
        for idx, (dir_name, files) in enumerate(class_files.items()):
            speaker_name = speaker_mapping.get(dir_name, dir_name)
            speaker_id = f"speaker{str(idx).zfill(3)}"
            speaker_id_mapping[speaker_id] = speaker_name
            
            all_files.extend(files)
            all_labels.extend([idx] * len(files))
            
            self.logger.info(f"Added {len(files)} files for speaker '{speaker_name}' "
                           f"(ID: {speaker_id}, Label: {idx})")
        
        if not all_files:
            raise ValueError("No valid audio files found")
        
        # Split data
        train_files, test_files, train_labels, test_labels = train_test_split(
            all_files, all_labels, test_size=test_split, random_state=42, stratify=all_labels
        )
        
        train_files, val_files, train_labels, val_labels = train_test_split(
            train_files, train_labels, test_size=val_split/(1-test_split),
            random_state=42, stratify=train_labels
        )
        
        # Create augmentation
        augmentation = AudioAugmentation(audio_processor.config.sample_rate)
        
        # Create datasets
        train_dataset = AudioDataset(
            train_files, train_labels, audio_processor,
            augmentation, augmentation_config, apply_augmentation=True
        )
        val_dataset = AudioDataset(
            val_files, val_labels, audio_processor,
            None, None, apply_augmentation=False
        )
        test_dataset = AudioDataset(
            test_files, test_labels, audio_processor,
            None, None, apply_augmentation=False
        )
        
        self.logger.info(f"Created voice recognition datasets - Train: {len(train_dataset)}, "
                        f"Val: {len(val_dataset)}, Test: {len(test_dataset)}")
        
        return train_dataset, val_dataset, test_dataset, speaker_id_mapping


class DataLoaderFactory:
    """Factory for creating optimized data loaders."""
    
    @staticmethod
    def create_dataloader(dataset: Dataset, 
                         batch_size: int,
                         shuffle: bool = True,
                         num_workers: int = 4,
                         pin_memory: bool = True,
                         drop_last: bool = False,
                         balanced_sampling: bool = False) -> DataLoader:
        """
        Create optimized data loader.
        
        Args:
            dataset: Dataset to load from
            batch_size: Batch size
            shuffle: Whether to shuffle data
            num_workers: Number of worker processes
            pin_memory: Whether to pin memory
            drop_last: Whether to drop last incomplete batch
            balanced_sampling: Whether to use balanced sampling for imbalanced datasets
            
        Returns:
            Configured DataLoader
        """
        sampler = None
        
        if balanced_sampling and hasattr(dataset, 'labels'):
            # Create weighted sampler for balanced sampling
            labels = dataset.labels
            class_counts = torch.bincount(torch.tensor(labels))
            class_weights = 1.0 / class_counts.float()
            sample_weights = class_weights[labels]
            
            sampler = WeightedRandomSampler(
                weights=sample_weights,
                num_samples=len(dataset),
                replacement=True
            )
            shuffle = False  # Don't shuffle when using sampler
        
        return DataLoader(
            dataset=dataset,
            batch_size=batch_size,
            shuffle=shuffle,
            sampler=sampler,
            num_workers=num_workers,
            pin_memory=pin_memory,
            drop_last=drop_last,
            persistent_workers=num_workers > 0
        )


class DataPreprocessor:
    """Comprehensive data preprocessing pipeline."""
    
    def __init__(self, audio_config: AudioProcessingConfig, 
                 logger: Optional[logging.Logger] = None):
        """Initialize data preprocessor."""
        self.audio_config = audio_config
        self.logger = logger or logging.getLogger(__name__)
        self.audio_processor = AudioProcessor(audio_config)
        self.dataset_builder = DatasetBuilder(logger)
    
    def preprocess_wakeword_data(self, raw_data_dir: str, 
                               output_dir: str,
                               wakeword_classes: Dict[int, str]) -> str:
        """
        Preprocess raw wakeword data.
        
        This method handles the requirements from the project description:
        1. Trim audio files to consistent length
        2. Split negative/background files into chunks
        3. Sort into train/val/test directories
        
        Args:
            raw_data_dir: Directory containing raw audio files
            output_dir: Output directory for processed data
            wakeword_classes: Mapping of class IDs to names
            
        Returns:
            Path to processed data directory
        """
        raw_data_path = Path(raw_data_dir)
        output_path = Path(output_dir)
        
        self.logger.info(f"Preprocessing wakeword data from {raw_data_dir}")
        
        # Process each class
        for class_id, class_name in wakeword_classes.items():
            class_raw_dir = raw_data_path / class_name
            if not class_raw_dir.exists():
                self.logger.warning(f"Class directory not found: {class_raw_dir}")
                continue
            
            self.logger.info(f"Processing class: {class_name}")
            
            # Get all audio files
            audio_files = []
            for ext in ['.wav', '.mp3', '.flac']:
                audio_files.extend(list(class_raw_dir.glob(f"**/*{ext}")))
            
            if not audio_files:
                self.logger.warning(f"No audio files found for class: {class_name}")
                continue
            
            processed_files = []
            
            for audio_file in audio_files:
                try:
                    # Load and process audio
                    waveform, sr = self.audio_processor.load_audio(str(audio_file))
                    
                    # Convert to mono
                    if waveform.shape[0] > 1:
                        waveform = torch.mean(waveform, dim=0, keepdim=True)
                    
                    # Resample
                    waveform = self.audio_processor.resample(waveform, sr)
                    
                    if class_name in ['negative', 'background']:
                        # Split long files into chunks
                        chunks = self._split_into_chunks(waveform, self.audio_config.target_samples)
                        for i, chunk in enumerate(chunks):
                            chunk_filename = f"{audio_file.stem}_chunk_{i:03d}.wav"
                            processed_files.append((chunk, chunk_filename))
                    else:
                        # For wakeword files, trim silence and normalize length
                        waveform = self.audio_processor.trim_silence(waveform)
                        waveform = self.audio_processor.normalize_length(waveform)
                        processed_files.append((waveform, audio_file.name))
                    
                except Exception as e:
                    self.logger.error(f"Error processing {audio_file}: {str(e)}")
                    continue
            
            # Save processed files
            class_output_dir = output_path / "chunked" / class_name
            class_output_dir.mkdir(parents=True, exist_ok=True)
            
            for waveform, filename in processed_files:
                output_file = class_output_dir / filename
                torchaudio.save(str(output_file), waveform, self.audio_config.sample_rate)
            
            self.logger.info(f"Processed {len(processed_files)} files for class: {class_name}")
        
        # Split into train/val/test
        self._split_processed_data(output_path / "chunked", output_path)
        
        return str(output_path)
    
    def preprocess_voice_recognition_data(self, raw_data_dir: str,
                                        output_dir: str,
                                        target_chunk_length: float = None) -> str:
        """
        Preprocess raw voice recognition data.
        
        Args:
            raw_data_dir: Directory containing raw speaker audio files
            output_dir: Output directory for processed data
            target_chunk_length: Target chunk length (uses wakeword length if None)
            
        Returns:
            Path to processed data directory
        """
        if target_chunk_length is None:
            target_chunk_length = self.audio_config.target_length
        
        raw_data_path = Path(raw_data_dir)
        output_path = Path(output_dir)
        
        self.logger.info(f"Preprocessing voice recognition data from {raw_data_dir}")
        
        # Process each speaker directory
        for speaker_dir in raw_data_path.iterdir():
            if not speaker_dir.is_dir():
                continue
            
            speaker_name = speaker_dir.name
            self.logger.info(f"Processing speaker: {speaker_name}")
            
            # Get all audio files
            audio_files = []
            for ext in ['.wav', '.mp3', '.flac']:
                audio_files.extend(list(speaker_dir.glob(f"**/*{ext}")))
            
            if not audio_files:
                self.logger.warning(f"No audio files found for speaker: {speaker_name}")
                continue
            
            processed_files = []
            target_samples = int(self.audio_config.sample_rate * target_chunk_length)
            
            for audio_file in audio_files:
                try:
                    # Load and process audio
                    waveform, sr = self.audio_processor.load_audio(str(audio_file))
                    
                    # Convert to mono
                    if waveform.shape[0] > 1:
                        waveform = torch.mean(waveform, dim=0, keepdim=True)
                    
                    # Resample
                    waveform = self.audio_processor.resample(waveform, sr)
                    
                    # Trim silence
                    waveform = self.audio_processor.trim_silence(waveform)
                    
                    # Split into chunks
                    chunks = self._split_into_chunks(waveform, target_samples)
                    
                    for i, chunk in enumerate(chunks):
                        chunk_filename = f"{audio_file.stem}_chunk_{i:03d}.wav"
                        processed_files.append((chunk, chunk_filename))
                
                except Exception as e:
                    self.logger.error(f"Error processing {audio_file}: {str(e)}")
                    continue
            
            # Save processed files
            speaker_output_dir = output_path / "chunked" / speaker_name
            speaker_output_dir.mkdir(parents=True, exist_ok=True)
            
            for waveform, filename in processed_files:
                output_file = speaker_output_dir / filename
                torchaudio.save(str(output_file), waveform, self.audio_config.sample_rate)
            
            self.logger.info(f"Processed {len(processed_files)} files for speaker: {speaker_name}")
        
        # Split into train/val/test
        self._split_processed_data(output_path / "chunked", output_path)
        
        return str(output_path)
    
    def _split_into_chunks(self, waveform: torch.Tensor, 
                          chunk_size: int, 
                          overlap: float = 0.0) -> List[torch.Tensor]:
        """Split waveform into chunks of specified size."""
        chunks = []
        total_samples = waveform.shape[-1]
        
        if total_samples <= chunk_size:
            # If shorter than chunk size, pad and return
            padded = F.pad(waveform, (0, chunk_size - total_samples))
            chunks.append(padded)
        else:
            # Split into overlapping chunks
            step_size = int(chunk_size * (1 - overlap))
            start = 0
            
            while start + chunk_size <= total_samples:
                chunk = waveform[..., start:start + chunk_size]
                chunks.append(chunk)
                start += step_size
            
            # Add final chunk if there's remaining audio
            if start < total_samples:
                final_chunk = waveform[..., -chunk_size:]
                chunks.append(final_chunk)
        
        return chunks
    
    def _split_processed_data(self, chunked_dir: Path, output_dir: Path):
        """Split processed data into train/val/test directories."""
        train_dir = output_dir / "train"
        val_dir = output_dir / "val"
        test_dir = output_dir / "test"
        
        # Create directories
        for split_dir in [train_dir, val_dir, test_dir]:
            split_dir.mkdir(parents=True, exist_ok=True)
        
        # Process each class/speaker
        for class_dir in chunked_dir.iterdir():
            if not class_dir.is_dir():
                continue
            
            class_name = class_dir.name
            
            # Get all files
            files = list(class_dir.glob("*.wav"))
            
            if len(files) < 3:
                self.logger.warning(f"Too few files for class {class_name} - skipping split")
                continue
            
            # Split files
            random.shuffle(files)
            
            total_files = len(files)
            test_count = max(1, int(total_files * 0.1))
            val_count = max(1, int(total_files * 0.2))
            train_count = total_files - test_count - val_count
            
            test_files = files[:test_count]
            val_files = files[test_count:test_count + val_count]
            train_files = files[test_count + val_count:]
            
            # Copy files to appropriate directories
            for split_name, split_files, split_dir in [
                ("train", train_files, train_dir),
                ("val", val_files, val_dir),
                ("test", test_files, test_dir)
            ]:
                class_split_dir = split_dir / class_name
                class_split_dir.mkdir(exist_ok=True)
                
                for file in split_files:
                    import shutil
                    shutil.copy2(file, class_split_dir / file.name)
            
            self.logger.info(f"Split {class_name}: {len(train_files)} train, "
                           f"{len(val_files)} val, {len(test_files)} test")