from torch import nn

class ActivationExtractor(nn.Module):

  """
  A utility class for extracting activations from intermediate layers of a PyTorch model.

  This class uses forward hooks to capture the output (activations) of a specified layer during the forward pass of the model.
  It is designed to work as a context manager for safe and automatic hook management.

  Args:
    model: The PyTorch model from which activations are to be extracted.
    target_layer: The specific layer iin the model whose activation needs to be captured.

  Attributes:
    activation: Captures the output (activations) of the target layer during the forward pass.

  Usage:
    >>> model = torchvision.models.resnet50(pretrained=True)
    >>> target_layer = model.layer4[0].conv1
    >>> with ActivationExtractor(model, target_layer) as extractor:
    ...   output = model(input_tensor)
    ...   activations = extractor.activation

  """

  def __init__(self, model, target_layer):

    """
    Initialize the ActivationExtractor with the model and target layer.

    Args:
      model: the PyTorch model
      target_layer: the layer from which activations are to be extracted.
    """

    super().__init__()

    self.model = model
    self.target_layer = target_layer
    self.activation = None # Stores activation from target layer

  def hook_fn(self, module, input, output):

    """
    The forward hook function to capture the output of the target layer.

    Args:
      module: The target layer (passed automatically by PyTorch)
      input: Input to the target layer (not used here but required by PyTorch)
      output: Output from the target layer (stored as th activation).
    """
    self.activation = output # Save the activation output during the forward pass.


  def __enter__(self):
    """
    Enter the context and register the forward hook on the target layer.

    Returns:
      self: The instance of the ActivationExtractor with the hook registered.
    """

    # Register the forward hook to capture activations
    self.hook = self.target_layer.register_forward_hook(self.hook_fn)

    return self


  def __exit__(self, exc_type, exc_val, exc_tb):
    """
    Exit the context and remove the forward hook to clean up resources.

    Args:
      exc_type: Exception type (if any) raised inside the context.
      exc_val: Exception value (if any) raised inside the context.
      exc_tb: Traceback (if any) raised inside the context.

    """
    # Remove the forward hook to avoid memory leaks.
    self.hook.remove()