"""
Gradient-weighted Attention Rollout for Vision Transformers (Chefer et al., CVPR 2021).
Calculates class-discriminative attention maps by weighting attention probabilities with backward gradients.
"""

import torch
import numpy as np


class ViTCheferRollout:
    def __init__(self, model):
        self.model = model
        self.attentions = []
        self.attention_gradients = []
        self.hooks = []

    def _register_hooks(self):
        self._remove_hooks()
        self.attentions = []
        self.attention_gradients = []

        for name, module in self.model.named_modules():
            if "attention.attention" in name and not hasattr(module, "_chefer_hooked"):
                def get_attn_hook(module, input, output):
                    # output in HF ViT self-attention is (context_layer, attn_probs)
                    if isinstance(output, tuple) and len(output) > 1:
                        self.attentions.append(output[1])
                    elif hasattr(output, "attentions"):
                        self.attentions.append(output.attentions)

                def get_attn_grad_hook(module, grad_input, grad_output):
                    # grad_output[1] is gradient w.r.t attention probabilities
                    if len(grad_output) > 1 and grad_output[1] is not None:
                        self.attention_gradients.append(grad_output[1])

                h1 = module.register_forward_hook(get_attn_hook)
                h2 = module.register_full_backward_hook(get_attn_grad_hook)
                self.hooks.extend([h1, h2])

    def _remove_hooks(self):
        for h in self.hooks:
            h.remove()
        self.hooks = []

    def attribute(self, input_tensor, target_class=0):
        """
        Compute Chefer Gradient-Weighted Attention Rollout.
        
        Args:
            input_tensor: (1, 3, 384, 384) preprocessed input tensor
            target_class: output index (0 for single binary logit)
            
        Returns:
            24x24 numpy array of patch attributions
        """
        self.model.eval()
        self._register_hooks()
        device = input_tensor.device

        # Enable gradients for backward pass
        input_tensor.requires_grad_(True)
        self.model.zero_grad()

        # Forward pass with output_attentions=True
        outputs = self.model(input_tensor, output_attentions=True)
        logits = outputs.logits  # (1, 1)

        # Backward pass on logit
        if logits.dim() > 1 and logits.size(1) > target_class:
            score = logits[:, target_class]
        else:
            score = logits.squeeze()

        score.backward()

        attentions = outputs.attentions  # Tuple of 12 (B, heads, 577, 577) tensors
        
        # Cleanup hooks
        self._remove_hooks()

        if not attentions:
            # Fallback if hooks failed to catch: return flat matrix
            return np.zeros((24, 24))

        seq_len = attentions[0].size(-1)  # 577
        result = torch.eye(seq_len, device=device)

        # Reverse gradient order because backward hook captures in reverse
        gradients = self.attention_gradients[::-1] if self.attention_gradients else [None] * len(attentions)

        for i, attn in enumerate(attentions):
            # attn: (1, heads, 577, 577)
            if i < len(gradients) and gradients[i] is not None:
                grad = gradients[i]
                # Elementwise product: gradient * attention
                cam = (grad * attn).clamp(min=0).mean(dim=1)  # (1, 577, 577)
            else:
                # Fallback to plain attention if gradient missing
                cam = attn.mean(dim=1)

            cam = cam[0]  # (577, 577)
            cam = cam + torch.eye(seq_len, device=device)
            cam = cam / cam.sum(dim=-1, keepdim=True)

            result = torch.matmul(cam, result)

        # Extract CLS token -> patch tokens
        patch_attn = result[0, 1:].reshape(24, 24).detach().cpu().numpy()
        return patch_attn


def chefer_gradient_rollout(wrapper, input_tensor) -> np.ndarray:
    """Convenience function for running Chefer rollout on ViTWrapper."""
    explainer = ViTCheferRollout(wrapper.model)
    return explainer.attribute(input_tensor, target_class=0)