class GradientDetector:
    """Tracks per-layer gradient norms and detects pathologies."""

    def __init__(self, vanishing_threshold=0.01, exploding_threshold=100.0,
                 stagnant_steps=50, ema_alpha=0.1):
        self.vanishing_threshold = vanishing_threshold
        self.exploding_threshold = exploding_threshold
        self.stagnant_steps = stagnant_steps
        self.alpha = ema_alpha

        self._ema = {}          # exponential moving average per layer
        self._stagnant_count = {}  # consecutive near-zero steps per layer
        self._stats = {}        # latest {mean, std, norm, status} per layer

    def update(self, layer_name, mean, std, norm):
        """Update stats for a layer and return detected status."""
        # Update EMA baseline
        if layer_name not in self._ema:
            self._ema[layer_name] = norm if norm > 0 else 1e-8
        else:
            self._ema[layer_name] = (
                self.alpha * norm + (1 - self.alpha) * self._ema[layer_name]
            )

        baseline = self._ema[layer_name]
        ratio = norm / baseline if baseline > 1e-10 else 0.0

        # Stagnation counter
        if norm < 1e-7:
            self._stagnant_count[layer_name] = self._stagnant_count.get(layer_name, 0) + 1
        else:
            self._stagnant_count[layer_name] = 0

        # Classify status
        stagnant_count = self._stagnant_count.get(layer_name, 0)
        if stagnant_count >= self.stagnant_steps:
            status = 'stagnant'
        elif ratio < self.vanishing_threshold:
            status = 'vanishing'
        elif ratio > self.exploding_threshold:
            status = 'exploding'
        else:
            status = 'ok'

        self._stats[layer_name] = {
            'mean': mean,
            'std': std,
            'norm': norm,
            'ratio': ratio,
            'stagnant_count': stagnant_count,
            'status': status,
        }
        return status

    def get_stats(self):
        return dict(self._stats)

    def reset(self):
        self._ema.clear()
        self._stagnant_count.clear()
        self._stats.clear()