"""
Logging utilities for CASWiT training.

Provides WandB integration and logging helpers.
"""

from typing import Optional
try:
    import wandb
except ImportError:
    wandb = None


def setup_wandb_logging(project: str, entity: str, run_name: str, config: dict, 
                       use_wandb: bool = True) -> bool:
    """
    Setup Weights & Biases logging.
    
    Args:
        project: WandB project name
        entity: WandB entity/username
        run_name: Run name
        config: Configuration dictionary
        use_wandb: Whether to use WandB
        
    Returns:
        True if WandB is initialized, False otherwise
    """
    if not use_wandb or wandb is None:
        return False
    
    wandb.init(project=project, entity=entity, config=config, name=run_name)
    return True


def log_metrics(metrics: dict, step: Optional[int] = None):
    """Log metrics to WandB."""
    if wandb is not None:
        wandb.log(metrics, step=step)