src/utils.py

"""
Utility functions for Speaker Profiling
"""

import os
import logging
import random
import numpy as np
import torch
import librosa
from pathlib import Path
from omegaconf import OmegaConf
from typing import Union, Optional, Tuple


def setup_logging(
    name: str = "speaker_profiling",
    level: int = logging.INFO,
    log_file: Optional[str] = None
) -> logging.Logger:
    """
    Setup logging configuration
    
    Args:
        name: Logger name
        level: Logging level
        log_file: Optional path to log file
    
    Returns:
        Configured logger instance
    """
    logger = logging.getLogger(name)
    logger.setLevel(level)
    
    if logger.handlers:
        logger.handlers.clear()
    
    formatter = logging.Formatter(
        fmt="%(asctime)s | %(levelname)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    )
    
    console_handler = logging.StreamHandler()
    console_handler.setLevel(level)
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    
    if log_file:
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setLevel(level)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
    
    return logger


def get_logger(name: str = "speaker_profiling") -> logging.Logger:
    """Get existing logger or create new one"""
    logger = logging.getLogger(name)
    if not logger.handlers:
        return setup_logging(name)
    return logger


def load_config(config_path: str) -> OmegaConf:
    """
    Load configuration from yaml file
    
    Args:
        config_path: Path to yaml config file
    
    Returns:
        OmegaConf configuration object
    """
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"Config file not found: {config_path}")
    return OmegaConf.load(config_path)


def set_seed(seed: int) -> None:
    """
    Set random seed for reproducibility
    
    Args:
        seed: Random seed value
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def load_audio(
    audio_path: Union[str, Path],
    sampling_rate: int = 16000,
    mono: bool = True
) -> Tuple[np.ndarray, int]:
    """
    Load audio file
    
    Args:
        audio_path: Path to audio file
        sampling_rate: Target sampling rate
        mono: Whether to convert to mono
    
    Returns:
        Tuple of (audio array, sampling rate)
    """
    audio, sr = librosa.load(audio_path, sr=sampling_rate, mono=mono)
    return audio, sr


def preprocess_audio(
    audio: np.ndarray,
    sampling_rate: int = 16000,
    max_duration: float = 10.0,
    trim_db: int = 20,
    normalize: bool = True,
    center_crop: bool = True
) -> np.ndarray:
    """
    Preprocess audio for model input
    
    Args:
        audio: Raw audio array
        sampling_rate: Audio sampling rate
        max_duration: Maximum duration in seconds
        trim_db: Threshold for silence trimming
        normalize: Whether to normalize audio
        center_crop: If True, center crop; else random crop (for training)
    
    Returns:
        Preprocessed audio array
    """
    max_length = int(sampling_rate * max_duration)
    
    audio, _ = librosa.effects.trim(audio, top_db=trim_db)
    
    if normalize:
        audio = audio / (np.max(np.abs(audio)) + 1e-8)
    
    if len(audio) < max_length:
        audio = np.pad(audio, (0, max_length - len(audio)))
    elif len(audio) > max_length:
        if center_crop:
            start = (len(audio) - max_length) // 2
        else:
            start = np.random.randint(0, len(audio) - max_length + 1)
        audio = audio[start:start + max_length]
    
    return audio


def load_and_preprocess_audio(
    audio_path: Union[str, Path],
    sampling_rate: int = 16000,
    max_duration: float = 10.0,
    trim_db: int = 20,
    normalize: bool = True,
    center_crop: bool = True
) -> np.ndarray:
    """
    Load and preprocess audio file in one step
    
    Args:
        audio_path: Path to audio file
        sampling_rate: Target sampling rate
        max_duration: Maximum duration in seconds
        trim_db: Threshold for silence trimming
        normalize: Whether to normalize audio
        center_crop: If True, center crop; else random crop
    
    Returns:
        Preprocessed audio array
    """
    audio, _ = load_audio(audio_path, sampling_rate)
    return preprocess_audio(
        audio, 
        sampling_rate, 
        max_duration, 
        trim_db, 
        normalize, 
        center_crop
    )


def load_model_checkpoint(
    model: torch.nn.Module,
    checkpoint_path: str,
    device: str = 'cpu'
) -> torch.nn.Module:
    """
    Load model from checkpoint
    
    Args:
        model: PyTorch model instance
        checkpoint_path: Path to checkpoint directory
        device: Device to load model on
    
    Returns:
        Model with loaded weights
    """
    logger = get_logger()
    
    safetensors_path = os.path.join(checkpoint_path, 'model.safetensors')
    pytorch_path = os.path.join(checkpoint_path, 'pytorch_model.bin')
    
    if os.path.exists(safetensors_path):
        from safetensors.torch import load_file
        state_dict = load_file(safetensors_path)
        logger.info(f"Loading checkpoint from {safetensors_path}")
    elif os.path.exists(pytorch_path):
        state_dict = torch.load(pytorch_path, map_location=device)
        logger.info(f"Loading checkpoint from {pytorch_path}")
    else:
        raise FileNotFoundError(
            f"No checkpoint found in {checkpoint_path}. "
            f"Expected 'model.safetensors' or 'pytorch_model.bin'"
        )
    
    model.load_state_dict(state_dict)
    return model


def get_device(device_str: str = 'cuda') -> torch.device:
    """
    Get torch device, fallback to CPU if CUDA not available
    
    Args:
        device_str: Desired device string ('cuda' or 'cpu')
    
    Returns:
        torch.device instance
    """
    if device_str == 'cuda' and torch.cuda.is_available():
        return torch.device('cuda')
    return torch.device('cpu')


def count_parameters(model: torch.nn.Module) -> Tuple[int, int]:
    """
    Count model parameters
    
    Args:
        model: PyTorch model
    
    Returns:
        Tuple of (total_params, trainable_params)
    """
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total, trainable


def format_number(num: int) -> str:
    """Format large numbers with commas"""
    return f"{num:,}"