"""Gradient monitoring utilities for detecting vanishing/exploding gradients."""

import torch
import torch.nn as nn
from typing import Dict, List, Optional, Tuple
import warnings


def compute_gradient_norms(model: nn.Module) -> Dict[str, float]:
    """Compute gradient L2 norm for each parameter in the model.
    
    Args:
        model: PyTorch model with gradients computed
        
    Returns:
        Dictionary mapping parameter name to gradient norm
    """
    grad_norms = {}
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norms[name] = param.grad.norm().item()
    return grad_norms


def compute_layer_gradient_norms(model: nn.Module) -> Dict[str, float]:
    """Compute aggregated gradient norms per layer/module.
    
    Groups parameters by their parent module (e.g., 'layers.0', 'layers.1', etc.)
    
    Args:
        model: PyTorch model with gradients computed
        
    Returns:
        Dictionary mapping layer prefix to total gradient norm
    """
    layer_norms = {}
    
    for name, param in model.named_parameters():
        if param.grad is None:
            continue
            
        # Extract layer prefix (first two components of the name)
        parts = name.split('.')
        if len(parts) >= 2:
            layer_prefix = '.'.join(parts[:2])
        else:
            layer_prefix = parts[0]
        
        grad_norm = param.grad.norm().item() ** 2
        
        if layer_prefix not in layer_norms:
            layer_norms[layer_prefix] = 0.0
        layer_norms[layer_prefix] += grad_norm
    
    # Take square root to get L2 norm
    return {k: v ** 0.5 for k, v in layer_norms.items()}


def check_vanishing_gradients(
    model: nn.Module, 
    threshold: float = 1e-7
) -> List[str]:
    """Check for parameters with vanishing gradients.
    
    Args:
        model: PyTorch model with gradients computed
        threshold: Gradient norm threshold below which gradient is considered vanishing
        
    Returns:
        List of parameter names with vanishing gradients
    """
    vanishing = []
    for name, param in model.named_parameters():
        if param.grad is not None:
            norm = param.grad.norm().item()
            if norm < threshold:
                vanishing.append(name)
    return vanishing


def check_exploding_gradients(
    model: nn.Module,
    threshold: float = 100.0
) -> List[str]:
    """Check for parameters with exploding gradients.
    
    Args:
        model: PyTorch model with gradients computed
        threshold: Gradient norm threshold above which gradient is considered exploding
        
    Returns:
        List of parameter names with exploding gradients
    """
    exploding = []
    for name, param in model.named_parameters():
        if param.grad is not None:
            norm = param.grad.norm().item()
            if norm > threshold or torch.isnan(param.grad).any() or torch.isinf(param.grad).any():
                exploding.append(name)
    return exploding


def gradient_health_check(
    model: nn.Module,
    vanishing_threshold: float = 1e-7,
    exploding_threshold: float = 100.0,
    warn: bool = True
) -> Tuple[List[str], List[str], Dict[str, float]]:
    """Comprehensive gradient health check.
    
    Args:
        model: PyTorch model with gradients computed
        vanishing_threshold: Threshold for vanishing gradients
        exploding_threshold: Threshold for exploding gradients
        warn: Whether to emit warnings for issues
        
    Returns:
        Tuple of (vanishing_params, exploding_params, layer_norms)
    """
    vanishing = check_vanishing_gradients(model, vanishing_threshold)
    exploding = check_exploding_gradients(model, exploding_threshold)
    layer_norms = compute_layer_gradient_norms(model)
    
    if warn and vanishing:
        warnings.warn(
            f"Vanishing gradients detected in {len(vanishing)} parameters: {vanishing[:5]}..."
            if len(vanishing) > 5 else f"Vanishing gradients in: {vanishing}"
        )
    
    if warn and exploding:
        warnings.warn(
            f"Exploding gradients detected in {len(exploding)} parameters: {exploding[:5]}..."
            if len(exploding) > 5 else f"Exploding gradients in: {exploding}"
        )
    
    return vanishing, exploding, layer_norms


class GradientMonitor:
    """Callback-style monitor for tracking gradients during training.
    
    Usage:
        monitor = GradientMonitor(model, check_interval=100)
        
        for step, batch in enumerate(dataloader):
            loss = model(batch)
            loss.backward()
            
            monitor.step(step)  # Logs every check_interval steps
            
            optimizer.step()
            optimizer.zero_grad()
    """
    
    def __init__(
        self,
        model: nn.Module,
        check_interval: int = 100,
        vanishing_threshold: float = 1e-7,
        exploding_threshold: float = 100.0,
        log_fn: Optional[callable] = None,
    ):
        """Initialize gradient monitor.
        
        Args:
            model: Model to monitor
            check_interval: How often to check gradients (in steps)
            vanishing_threshold: Threshold for vanishing gradient detection
            exploding_threshold: Threshold for exploding gradient detection
            log_fn: Optional logging function. If None, uses print.
        """
        self.model = model
        self.check_interval = check_interval
        self.vanishing_threshold = vanishing_threshold
        self.exploding_threshold = exploding_threshold
        self.log_fn = log_fn or print
        
        self.history: List[Dict] = []
    
    def step(self, step_num: int) -> Optional[Dict]:
        """Check gradients at this step if interval reached.
        
        Args:
            step_num: Current training step number
            
        Returns:
            Gradient statistics dict if check performed, None otherwise
        """
        if step_num % self.check_interval != 0:
            return None
        
        vanishing, exploding, layer_norms = gradient_health_check(
            self.model,
            self.vanishing_threshold,
            self.exploding_threshold,
            warn=False  # We'll log ourselves
        )
        
        stats = {
            'step': step_num,
            'num_vanishing': len(vanishing),
            'num_exploding': len(exploding),
            'layer_norms': layer_norms,
            'min_norm': min(layer_norms.values()) if layer_norms else 0,
            'max_norm': max(layer_norms.values()) if layer_norms else 0,
        }
        
        self.history.append(stats)
        
        # Log summary
        self.log_fn(
            f"[Step {step_num}] Gradient stats: "
            f"min={stats['min_norm']:.2e}, max={stats['max_norm']:.2e}, "
            f"vanishing={stats['num_vanishing']}, exploding={stats['num_exploding']}"
        )
        
        if vanishing:
            self.log_fn(f"  Vanishing in: {vanishing[:3]}{'...' if len(vanishing) > 3 else ''}")
        if exploding:
            self.log_fn(f"  Exploding in: {exploding[:3]}{'...' if len(exploding) > 3 else ''}")
        
        return stats
    
    def get_history(self) -> List[Dict]:
        """Return all recorded gradient statistics."""
        return self.history
