src/utils/helpers.py

"""
Utility helper functions.
Single Responsibility: Common helper utilities.
"""
import torch
import numpy as np
import random
from pathlib import Path
from typing import Optional


def set_seed(seed: int = 42) -> 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(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def get_device(prefer_gpu: bool = True) -> torch.device:
    """
    Get computation device (GPU or CPU).

    Args:
        prefer_gpu: Whether to prefer GPU if available

    Returns:
        torch.device object
    """
    if prefer_gpu and torch.cuda.is_available():
        device = torch.device('cuda')
        print(f"Using GPU: {torch.cuda.get_device_name(0)}")
    else:
        device = torch.device('cpu')
        print("Using CPU")

    return device


def ensure_dir(directory: Path) -> Path:
    """
    Ensure directory exists, create if needed.

    Args:
        directory: Directory path

    Returns:
        Path object
    """
    directory = Path(directory)
    directory.mkdir(parents=True, exist_ok=True)
    return directory


def count_parameters(model: torch.nn.Module) -> int:
    """
    Count trainable parameters in a PyTorch model.

    Args:
        model: PyTorch model

    Returns:
        Number of trainable parameters
    """
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def print_data_summary(df, name: str = "DataFrame") -> None:
    """
    Print summary statistics of a DataFrame.

    Args:
        df: Pandas DataFrame
        name: Name to display
    """
    print(f"\n{name} Summary:")
    print(f"  Shape: {df.shape}")
    print(f"  Columns: {list(df.columns)}")
    print(f"  Missing values: {df.isnull().sum().sum()}")
    if "date" in df.columns:
        print(f"  Date range: {df['date'].min()} to {df['date'].max()}")
    if "station_id" in df.columns:
        print(f"  Unique stations: {df['station_id'].nunique()}")
    print()