import math

import torch
import torch.nn.functional as F
from einops import rearrange
from torch import Tensor, nn

from zeta.nn.biases.base import BaseBias

######## Helpers

def exists(val):
    return val is not None


def pad_at_dim(t, pad, dim=-1, value=0.):
    dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
    zeros = ((0, 0) * dims_from_right)
    return F.pad(t, (*zeros, *pad), value=value)





class AlibiPositionalBias(BaseBias):
    def __init__(self, heads, num_heads, **kwargs):
        super().__init__()
        self.heads = heads
        self.num_heads = num_heads

        slopes = Tensor(self.__get_slopes(heads))
        slopes = rearrange(slopes, 'h -> h 1 1')
        self.register_buffer('slopes', slopes, persistent=False)
        self.register_buffer('bias', None, persistent=False)


    def get_bias(self, i, j, device):
        torch.arange(j - i, j, device=device)
        j_arange = torch.arange(j, device = device)
        bias = -torch.abs(rearrange(j_arange, 'j -> 1 1 j'))
        return bias
    
    @staticmethod
    def _get_slopes(heads):
        def get_slopes_power_of_2(n):
            start = (2**(-2**-(math.log2(n)-3)))
            ratio = start

            return [start*ratio**i for i in range(n)]
        
        if math.log2(heads).is_integer():
            return get_slopes_power_of_2(heads)
        
        closest_power_of_2 = 2 ** math.floor(math.log2(heads))
        return get_slopes_power_of_2(closest_power_of_2) + get_slopes_power_of_2(2 * closest_power_of_2)[0::2][:heads-closest_power_of_2]
    
    @property
    def device(self):
        return next(self.buffers()).device
    
    def forward(self, i, j):
        h, device = self.num_heads, self.device
        
        if exists(self.bias) and self.bias.shape[-1] >= j and self.bias.shape[-2] >= i:
            return self.bias[..., :i, :j]
        
        bias = self.get_bias(i, j, device)
        bias = bias * self._get_slopes

        num_heads_unalibied = h - bias.shape[0]
        bias = pad_at_dim(bias, (0, num_heads_unalibied), dim=0)
        self.register_buffer('bias', bias, persistent=False)

        return self.bias
    
class LearnedAlibiPositionalBias(AlibiPositionalBias):
    def __init__(self, heads, num_heads):
        super().__init__(heads, num_heads)
        log_slopes = torch.log(self.slopes)
        self.learned_logslopes = nn.Parameter(log_slopes)

    def forward(self, i, j):
        h, device = self.heads, self.device

        def get_slopes(param):
            return pad_at_dim(param.exp(), (0, h - param.shape[0]), dim=-2)
        
        if exists(self.bias) and self.bias.shape[-1] >= j and self.bias.shape[-2] >= i:
            bias = self.bias[..., :i, :j]
        else:
            bias = self.get_bias(i, j, device)
            self.register_buffer('bias', bias, persistent=False)

        slopes = get_slopes(self.learned_logslopes)
        bias = bias * slopes

        return bias 
    