"""
Logging utilities for training metrics.

Provides a unified Logger class that handles:
- TensorBoard logging
- Console output
- File-based logs
"""

from pathlib import Path
from torch.utils.tensorboard import SummaryWriter


class Logger:
    """
    Unified logger for training metrics.
    
    Wraps TensorBoard SummaryWriter and provides a clean interface
    for logging scalars, histograms, and other metrics.
    """
    
    def __init__(self, log_dir: Path):
        """
        Args:
            log_dir: Directory for TensorBoard logs
        """
        self.log_dir = log_dir
        self.writer = SummaryWriter(log_dir / "tensorboard")
    
    def log_scalar(self, tag: str, value: float, step: int) -> None:
        """
        Log a scalar value.
        
        Args:
            tag: Metric name (e.g., "losses/policy_loss")
            value: Scalar value
            step: Global step/timestep
        """
        self.writer.add_scalar(tag, value, step)
    
    def log_histogram(self, tag: str, values, step: int) -> None:
        """
        Log a histogram.
        
        Args:
            tag: Metric name
            values: Array/tensor of values
            step: Global step
        """
        self.writer.add_histogram(tag, values, step)
    
    def close(self) -> None:
        """Close the logger and flush all pending writes."""
        self.writer.close()