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.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):
    """
    AlibiPositionalBias class represents a positional bias module for neural networks.

    Args:
        heads (int): Number of heads in the neural network.
        num_heads (int): Number of heads in the neural network.

    Attributes:
        slopes (Tensor): Tensor containing the slopes for the bias.
        bias (Tensor): Tensor containing the bias values.

    Methods:
        get_bias(i, j, device): Returns the bias tensor for the given indices.
        forward(i, j): Computes and returns the bias tensor for the given indices.

    """

    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):
    """
    LearnedAlibiPositionalBias is a subclass of AlibiPositionalBias that introduces learned biases.

    Args:
        heads (int): Number of attention heads.
        num_heads (int): Number of heads per layer.

    Attributes:
        learned_logslopes (nn.Parameter): Learned logarithmic slopes.

    """

    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
