import torch
from ...component import Component


class Alibi(Component):
	"""ALiBi (Attention with Linear Biases) positional encoding.
	
	ALiBi doesn't add positional embeddings to the input. Instead, it modifies
	the attention mechanism by adding linear biases. This component stores
	the bias matrix that should be added to attention scores.
	
	For use with attention, the bias should be added to attention scores.
	This component can be used to generate the bias matrix.
	
	Args:
		num_heads: Number of attention heads (if None, uses cfg.num_heads)
		slope_base: Base slope value (default: 2^(-8/head))
	"""
	def __init__(self, num_heads: int = None, slope_base: float = None):
		super().__init__(name="Alibi")
		self.num_heads = num_heads
		self.slope_base = slope_base
		self.bias_cache = {}

	def build(self, cfg):
		super().build(cfg)
		num_heads = self.num_heads if self.num_heads is not None else cfg.num_heads
		if num_heads is None:
			raise ValueError("num_heads must be specified in cfg or component init")
		

		if 'slopes' not in self._buffers:
			if 'slopes' in self.__dict__:
				delattr(self, 'slopes')
			if self.slope_base is None:
				slopes = []
				for h in range(1, num_heads + 1):
					slope = 2 ** (-8.0 / h)
					slopes.append(slope)
				self.slopes = torch.tensor(slopes)
			else:
				self.slopes = torch.tensor([self.slope_base] * num_heads)
			
			self.register_buffer('slopes', self.slopes, persistent=False)

	def get_bias(self, seq_len: int, device: torch.device = None) -> torch.Tensor:
		"""Get ALiBi bias matrix for a given sequence length.
		
		Args:
			seq_len: Sequence length
			device: Device to create bias on
		
		Returns:
			Bias tensor of shape (1, num_heads, 1, seq_len)
		"""
		context_position = torch.arange(seq_len, device=device)[:, None]
		memory_position = torch.arange(seq_len, device=device)[None, :]
		relative_position = memory_position - context_position 
		relative_position = torch.abs(relative_position).unsqueeze(0).expand(self.num_heads, -1, -1)
		
		# Slopes are (num_heads,)
		slopes = self.slopes.unsqueeze(1).unsqueeze(1)
		bias = slopes * relative_position
		# Returns (1, num_heads, seq_len, seq_len)
		return bias.unsqueeze(0)

	def forward(self, x, **kwargs):
		"""
		Returns x and the ALiBi bias matrix.
		"""
		seq_len = x.size(1)
		bias = self.get_bias(seq_len, x.device)
		return x, bias