"""
TensorBoard Logger - PyTorch Neural Networks
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: TensorBoard logging utilities
"""

from torch.utils.tensorboard import SummaryWriter
import torch
import numpy as np


class TensorBoardLogger:
    """
    TensorBoard logger for training visualization
    
    Project: PyTorch Neural Networks
    Author: RSK World
    Website: https://rskworld.in
    """
    
    def __init__(self, log_dir='runs/experiment'):
        """
        Initialize TensorBoard logger
        
        Args:
            log_dir: Directory to save logs
        """
        self.writer = SummaryWriter(log_dir)
        self.global_step = 0
    
    def log_scalar(self, tag, value, step=None):
        """
        Log scalar value
        
        Args:
            tag: Tag name
            value: Scalar value
            step: Step number (uses global step if None)
        """
        if step is None:
            step = self.global_step
        self.writer.add_scalar(tag, value, step)
    
    def log_histogram(self, tag, values, step=None):
        """
        Log histogram
        
        Args:
            tag: Tag name
            values: Values to histogram
            step: Step number
        """
        if step is None:
            step = self.global_step
        self.writer.add_histogram(tag, values, step)
    
    def log_model_graph(self, model, input_sample):
        """
        Log model graph
        
        Args:
            model: Model to log
            input_sample: Sample input
        """
        self.writer.add_graph(model, input_sample)
    
    def log_images(self, tag, images, step=None):
        """
        Log images
        
        Args:
            tag: Tag name
            images: Image tensor
            step: Step number
        """
        if step is None:
            step = self.global_step
        self.writer.add_images(tag, images, step)
    
    def log_learning_rate(self, optimizer, step=None):
        """
        Log learning rate
        
        Args:
            optimizer: Optimizer
            step: Step number
        """
        if step is None:
            step = self.global_step
        for i, param_group in enumerate(optimizer.param_groups):
            self.writer.add_scalar(f'LearningRate/group_{i}', param_group['lr'], step)
    
    def log_training_metrics(self, metrics, step=None):
        """
        Log training metrics
        
        Args:
            metrics: Dictionary of metrics
            step: Step number
        """
        if step is None:
            step = self.global_step
        
        for key, value in metrics.items():
            self.log_scalar(f'Train/{key}', value, step)
    
    def log_validation_metrics(self, metrics, step=None):
        """
        Log validation metrics
        
        Args:
            metrics: Dictionary of metrics
            step: Step number
        """
        if step is None:
            step = self.global_step
        
        for key, value in metrics.items():
            self.log_scalar(f'Val/{key}', value, step)
    
    def increment_step(self):
        """Increment global step"""
        self.global_step += 1
    
    def close(self):
        """Close writer"""
        self.writer.close()