"""
WandB Utilities
===============
This module provides utilities for improved WandB integration.
It fixes common issues with metric logging and visualization.

Key features:
1. Automatic metric definition for proper visualization
2. Consistent naming conventions
3. Gradient and learning rate tracking
4. Custom chart creation
"""

import os
import wandb
from typing import Dict, Any, Optional, List
import torch
import numpy as np


class WandBLogger:
    """
    Enhanced WandB logger with proper metric tracking and visualization.
    
    This class ensures all metrics are properly logged and visualized
    in the WandB dashboard.
    """
    
    def __init__(
        self,
        project: str,
        name: str,
        config: Dict[str, Any],
        tags: Optional[List[str]] = None,
        entity: Optional[str] = None,
        mode: str = "online",
        resume: str = "allow",
        rank: int = 0
    ):
        """
        Initialize WandB logger with comprehensive configuration.
        
        Args:
            project: WandB project name
            name: Run name
            config: Configuration dictionary
            tags: List of tags for the run
            entity: WandB entity (team/user)
            mode: "online", "offline", or "disabled"
            resume: "allow", "must", "never", "auto", or run ID
            rank: Process rank for distributed training
        """
        self.rank = rank
        self.enabled = rank == 0  # Only log from main process
        
        if not self.enabled:
            return
        
        # Initialize WandB
        self.run = wandb.init(
            project=project,
            name=name,
            config=config,
            tags=tags or [],
            entity=entity,
            mode=mode,
            resume=resume,
            # Fix: Ensure run ID is consistent for resuming
            id=config.get('wandb_run_id', None),
            # Fix: Add more metadata
            job_type="training",
            notes=config.get('notes', '')
        )
        
        # Define custom metrics for proper visualization
        self._define_metrics()
        
        # Log initial configuration
        self._log_initial_info(config)
    
    def _define_metrics(self):
        """
        Define custom metrics for proper step tracking.
        
        This ensures graphs are displayed correctly in WandB.
        """
        # Define step metrics
        wandb.define_metric("epoch")
        wandb.define_metric("global_step")
        
        # Training metrics (tracked per step)
        wandb.define_metric("training/*", step_metric="global_step")
        wandb.define_metric("time/*", step_metric="global_step")
        wandb.define_metric("gradients/*", step_metric="global_step")
        
        # Validation metrics (tracked per epoch)
        wandb.define_metric("validation/*", step_metric="epoch")
        wandb.define_metric("best/*", step_metric="epoch")
        
        # Learning rate metrics (special handling for multiple param groups)
        wandb.define_metric("lr/*", step_metric="global_step")
    
    def _log_initial_info(self, config: Dict[str, Any]):
        """Log initial training information."""
        # Log hyperparameters as summary metrics
        wandb.run.summary["total_epochs"] = config.get('epochs', 0)
        wandb.run.summary["batch_size"] = config.get('batch_size', 0)
        wandb.run.summary["learning_rate"] = config.get('base_lr', 0)
        wandb.run.summary["model_type"] = config.get('model_name', 'unknown')
    
    def log_training_step(
        self,
        metrics: Dict[str, float],
        epoch: int,
        step: int,
        global_step: int,
        optimizer: Optional[torch.optim.Optimizer] = None,
        grad_norm: Optional[float] = None
    ):
        """
        Log training metrics for a single step.
        
        Args:
            metrics: Dictionary of metric values
            epoch: Current epoch
            step: Step within epoch
            global_step: Global step count
            optimizer: Optimizer for learning rate logging
            grad_norm: Gradient norm if available
        """
        if not self.enabled:
            return
        
        log_dict = {
            "epoch": epoch,
            "global_step": global_step,
        }
        
        # Add training metrics with proper prefix
        for key, value in metrics.items():
            if not key.startswith('training/'):
                key = f"training/{key}"
            log_dict[key] = value
        
        # Log learning rates from optimizer
        if optimizer is not None:
            # Log individual parameter group learning rates
            for i, param_group in enumerate(optimizer.param_groups):
                log_dict[f"lr/group_{i}"] = param_group['lr']
                # Also log parameter group names if available
                if 'name' in param_group:
                    log_dict[f"lr/{param_group['name']}"] = param_group['lr']
            
            # Log the main learning rate
            log_dict["training/learning_rate"] = optimizer.param_groups[0]['lr']
        
        # Log gradient norm if available
        if grad_norm is not None:
            log_dict["gradients/norm"] = grad_norm
        
        wandb.log(log_dict, step=global_step)
    
    def log_validation_epoch(
        self,
        metrics: Dict[str, float],
        epoch: int,
        best_metrics: Optional[Dict[str, float]] = None
    ):
        """
        Log validation metrics for an epoch.
        
        Args:
            metrics: Validation metrics
            epoch: Current epoch
            best_metrics: Best metrics so far
        """
        if not self.enabled:
            return
        
        log_dict = {"epoch": epoch}
        
        # Add validation metrics
        for key, value in metrics.items():
            if not key.startswith('validation/'):
                key = f"validation/{key}"
            log_dict[key] = value
        
        # Add best metrics if provided
        if best_metrics:
            for key, value in best_metrics.items():
                if not key.startswith('best/'):
                    key = f"best/{key}"
                log_dict[key] = value
        
        wandb.log(log_dict, step=epoch)
    
    def log_images(
        self,
        images: Dict[str, np.ndarray],
        epoch: int,
        captions: Optional[Dict[str, str]] = None
    ):
        """
        Log images to WandB.
        
        Args:
            images: Dictionary of image arrays
            epoch: Current epoch
            captions: Optional captions for images
        """
        if not self.enabled:
            return
        
        wandb_images = {}
        for key, img in images.items():
            caption = captions.get(key, key) if captions else key
            wandb_images[key] = wandb.Image(img, caption=caption)
        
        wandb.log(wandb_images, step=epoch)
    
    def create_custom_charts(self):
        """
        Create custom charts for better visualization.
        
        This method creates preset chart configurations.
        """
        if not self.enabled:
            return
        
        # Training overview chart
        wandb.log({
            "charts/training_overview": wandb.plot.line_series(
                xs=range(100),  # Placeholder
                ys=[[0]*100, [0]*100, [0]*100],  # Placeholder
                keys=["Loss", "IoU", "Learning Rate"],
                title="Training Overview",
                xname="Step"
            )
        })
    
    def log_model_gradients(
        self,
        model: torch.nn.Module,
        global_step: int,
        log_histogram: bool = False
    ):
        """
        Log model gradient statistics.
        
        Args:
            model: PyTorch model
            global_step: Global step count
            log_histogram: Whether to log gradient histograms
        """
        if not self.enabled:
            return
        
        gradients = []
        gradient_norms = []
        
        for name, param in model.named_parameters():
            if param.grad is not None:
                grad = param.grad.data
                gradients.append(grad.norm().item())
                gradient_norms.append((name, grad.norm().item()))
                
                if log_histogram:
                    wandb.log({
                        f"gradients/histogram/{name}": wandb.Histogram(grad.cpu().numpy())
                    }, step=global_step)
        
        # Log aggregate statistics
        if gradients:
            wandb.log({
                "gradients/mean_norm": np.mean(gradients),
                "gradients/max_norm": np.max(gradients),
                "gradients/min_norm": np.min(gradients),
            }, step=global_step)
    
    def finish(self, final_metrics: Optional[Dict[str, float]] = None):
        """
        Finish the WandB run.
        
        Args:
            final_metrics: Final metrics to log as summary
        """
        if not self.enabled:
            return
        
        # Log final metrics to summary
        if final_metrics:
            for key, value in final_metrics.items():
                wandb.run.summary[key] = value
        
        wandb.finish()


def setup_wandb_from_config(config: Dict[str, Any], rank: int = 0) -> WandBLogger:
    """
    Create WandB logger from configuration dictionary.
    
    This is a convenience function for setting up WandB from YAML configs.
    """
    # Extract WandB configuration
    wandb_config = config.get('wandb', {})
    
    # Use environment variable for API key (more secure)
    api_key = os.environ.get('WANDB_API_KEY')
    if api_key:
        wandb.login(key=api_key)
    
    # Create logger
    logger = WandBLogger(
        project=wandb_config.get('project', 'default-project'),
        name=config.get('exp_name', 'unnamed-experiment'),
        config=config,
        tags=wandb_config.get('tags', []),
        entity=wandb_config.get('entity', None),
        mode=wandb_config.get('mode', 'online'),
        resume=wandb_config.get('resume', 'allow'),
        rank=rank
    )
    
    return logger


def fix_metric_names(metrics: Dict[str, Any]) -> Dict[str, Any]:
    """
    Fix metric names to follow WandB conventions.
    
    This ensures consistent naming across different parts of the codebase.
    """
    fixed_metrics = {}
    
    for key, value in metrics.items():
        # Skip if already properly prefixed
        if '/' in key:
            fixed_metrics[key] = value
            continue
        
        # Add appropriate prefix based on metric name
        if key in ['loss', 'iou', 'accuracy', 'precision', 'recall']:
            fixed_metrics[f'training/{key}'] = value
        elif key.startswith('val_'):
            fixed_metrics[f'validation/{key[4:]}'] = value
        elif key.startswith('lr'):
            fixed_metrics[f'lr/{key}'] = value
        elif key.startswith('time'):
            fixed_metrics[f'time/{key}'] = value
        else:
            fixed_metrics[key] = value
    
    return fixed_metrics