"""
PyTorch hooks for capturing activations and gradients.

Used by analysis methods to intercept intermediate model outputs.
"""

from typing import Optional, Callable
import torch
import torch.nn as nn


class ActivationHook:
    """
    Forward hook to capture layer activations.
    
    Usage:
        hook = ActivationHook(model.layer)
        output = model(input)
        activations = hook.activation  # Captured tensor
        hook.remove()
    """
    
    def __init__(self, module: nn.Module):
        self.activation: Optional[torch.Tensor] = None
        self._handle = module.register_forward_hook(self._hook_fn)
    
    def _hook_fn(self, module: nn.Module, input: tuple, output: torch.Tensor):
        """Store activation, detached from computation graph."""
        self.activation = output.detach()
    
    def remove(self):
        """Remove the hook."""
        self._handle.remove()
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.remove()


class GradientHook:
    """
    Backward hook to capture gradients flowing through a layer.
    
    Usage:
        hook = GradientHook(model.layer)
        output = model(input)
        loss.backward()
        gradients = hook.gradient  # Captured gradient tensor
        hook.remove()
    """
    
    def __init__(self, module: nn.Module):
        self.gradient: Optional[torch.Tensor] = None
        self._handle = module.register_full_backward_hook(self._hook_fn)
    
    def _hook_fn(
        self, 
        module: nn.Module, 
        grad_input: tuple, 
        grad_output: tuple
    ):
        """Store gradient w.r.t. layer output."""
        self.gradient = grad_output[0].detach()
    
    def remove(self):
        """Remove the hook."""
        self._handle.remove()
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.remove()


class ActivationAndGradientHook:
    """
    Combined hook for capturing both activations and gradients.
    
    Useful for CAM-style methods that need both forward activations
    and backward gradients for the same layer.
    """
    
    def __init__(self, module: nn.Module):
        self.activation: Optional[torch.Tensor] = None
        self.gradient: Optional[torch.Tensor] = None
        
        self._fwd_handle = module.register_forward_hook(self._fwd_hook)
        self._bwd_handle = module.register_full_backward_hook(self._bwd_hook)
    
    def _fwd_hook(self, module: nn.Module, input: tuple, output: torch.Tensor):
        self.activation = output.detach()
    
    def _bwd_hook(self, module: nn.Module, grad_input: tuple, grad_output: tuple):
        self.gradient = grad_output[0].detach()
    
    def remove(self):
        """Remove both hooks."""
        self._fwd_handle.remove()
        self._bwd_handle.remove()
    
    def clear(self):
        """Clear stored tensors (useful between iterations)."""
        self.activation = None
        self.gradient = None
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.remove()
