import torch


def make_backward_hook(layer_name, detector, reporter, step_counter):
    """Returns a backward hook that computes gradient stats in-hook."""
    def hook(module, grad_input, grad_output):
        # Use grad_output[0] — gradient flowing into this layer
        for g in grad_output:
            if g is None:
                continue
            with torch.no_grad():
                g_flat = g.detach().float().reshape(-1)
                mean = g_flat.mean().item()
                std = g_flat.std().item()
                norm = g_flat.norm().item()
            # Stats computed — raw tensor NOT stored
            status = detector.update(layer_name, mean, std, norm)
            if status != 'ok':
                reporter.alert(layer_name, status, norm, step_counter[0])
            break  # only process first non-None grad
    return hook


def make_forward_hook(layer_name, detector):
    """Returns a forward hook that checks for dead activations."""
    def hook(module, input, output):
        with torch.no_grad():
            if isinstance(output, torch.Tensor):
                out = output.detach().float()
                dead_ratio = (out.abs() < 1e-7).float().mean().item()
                # Store dead neuron ratio for reporting
                stats = detector.get_stats()
                if layer_name in stats:
                    stats[layer_name]['dead_ratio'] = dead_ratio
    return hook