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


class Dropout(nn.Module):
    def __init__(
        self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **_kwargs
    ):
        super().__init__()
        if p < 0.0 or p > 1.0:
            raise ValueError(
                f"dropout probability has to be between 0 and 1, but got {p}"
            )
        self.p = p
        self.inplace = inplace
        self.force_drop = force_drop

    def forward(self, x, **_kwargs):
        return F.dropout(
            x,
            p=self.p,
            training=True if self.force_drop else self.training,
            inplace=self.inplace,
        )


class Conv1d(nn.Conv1d):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 1,
        stride: int = 1,
        dilation: int = 1,
        groups: int = 1,
        padding_mode: str = "zeros",
        bias: bool = True,
        padding=None,
        causal: bool = False,
        **_kwargs,
    ):
        self.causal = causal
        if padding is None:
            if causal:
                padding = 0
                self.left_padding = dilation * (kernel_size - 1)
            else:
                padding = int((kernel_size * dilation - dilation) / 2)

        super().__init__(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            padding_mode=padding_mode,
            bias=bias,
        )

        self.in_channels = in_channels

    def forward(self, x):
        if self.causal:
            x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2)
        return super().forward(x)


class ConvTranspose1d(nn.ConvTranspose1d):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
        output_padding: int = 0,
        groups: int = 1,
        bias: bool = True,
        dilation: int = 1,
        padding=None,
        padding_mode: str = "zeros",
        causal: bool = False,
        **_kwargs,
    ):
        if padding is None:
            padding = 0 if causal else (kernel_size - stride) // 2
        if causal:
            assert padding == 0, "padding is not allowed in causal ConvTranspose1d."
            assert kernel_size == 2 * stride, (
                "kernel_size must be equal to 2*stride in Causal ConvTranspose1d."
            )

        super().__init__(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            output_padding=output_padding,
            groups=groups,
            bias=bias,
            dilation=dilation,
            padding_mode=padding_mode,
        )

        self.causal = causal
        self.stride = stride

    def forward(self, x):
        x = super().forward(x)
        if self.causal:
            x = x[:, :, : -self.stride]
        return x


class Mlp(nn.Module):
    def __init__(
        self,
        hidden_size,
        ffn_hidden_size=4096,
        act_layer=nn.GELU,
        dropout=0.0,
        **_kwargs,
    ):
        super().__init__()
        self.fc1 = nn.Linear(hidden_size, ffn_hidden_size)
        self.act = act_layer()
        self.fc2 = nn.Linear(ffn_hidden_size, hidden_size)
        self.drop = Dropout(dropout)

    def forward(self, x, _mask=None):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        return self.drop(x)


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


@torch.autocast(enabled=False, device_type="cuda")
def apply_rotary_pos_emb(pos, t):
    if pos.dim() == 3:
        pos = pos.unsqueeze(1)
    return t * pos.cos() + rotate_half(t) * pos.sin()


class RotaryEmbedding(nn.Module):
    def __init__(self, dim, theta=50000):
        super().__init__()
        self.register_buffer(
            "inv_freq",
            1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)),
            persistent=False,
        )
        self._theta = float(theta)

    def _apply(self, fn):
        inv_freq = self.inv_freq
        super()._apply(fn)
        self.inv_freq = inv_freq.to(device=self.inv_freq.device, dtype=torch.float32)
        return self

    @torch.autocast(enabled=False, device_type="cuda")
    def forward(self, t):
        inv_freq = self.inv_freq
        if inv_freq.device != t.device:
            raise RuntimeError(
                "RotaryEmbedding buffer device mismatch: "
                f"inv_freq={inv_freq.device} input={t.device}."
            )
        t = t.to(dtype=inv_freq.dtype)
        if t.dim() == 1:
            freqs = torch.einsum("i , j -> i j", t, inv_freq)
        else:
            freqs = torch.einsum("bi, j -> bij", t, inv_freq)
        return torch.cat((freqs, freqs), dim=-1)


class MultiHeadAttention(nn.Module):
    """Multi-head attention"""

    def __init__(
        self,
        hidden_size: int,
        num_heads: int = 8,
        qkv_bias: bool = False,
        qk_norm: bool = False,
        attn_drop: float = 0.0,
        dropout: float = 0.0,
        norm_layer: str = "LayerNorm",
        rotary_bias: bool = False,
        rotary_theta: float | None = 50000,
        **_kwargs,
    ):
        super().__init__()
        assert hidden_size % num_heads == 0, (
            "hidden_size should be divisible by num_heads"
        )
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.scale = self.head_dim**-0.5
        self.rotary_bias = rotary_bias

        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)

        norm_layer = getattr(nn, norm_layer)
        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()

        self.attn_drop = Dropout(attn_drop)
        self.o_proj = nn.Linear(hidden_size, hidden_size)
        self.o_dropout = Dropout(dropout)

        if self.rotary_bias:
            self.rotary = RotaryEmbedding(self.head_dim, theta=rotary_theta)

    def forward(self, q, k=None, v=None, mask=None, pos_ids=None, **_kwargs):
        k = k or q
        v = v or q
        B, L, _ = q.shape
        _, S, _ = v.shape
        if mask is not None:
            if mask.ndim == 2:  # [B, L]
                assert L == S
                mask = rearrange(mask, "b j -> b 1 1 j")
                mask = mask.expand(-1, self.num_heads, L, -1)
            elif mask.ndim == 3:  # [B, L, S]
                assert mask.size(1) == L and mask.size(2) == S
                mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)

        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
        q, k = self.q_norm(q), self.k_norm(k)

        # Apply rotary
        if self.rotary_bias:
            if L == S:
                if pos_ids is None:
                    rotary_emb = self.rotary(torch.arange(L, device=q.device))
                else:
                    rotary_emb = self.rotary(pos_ids)
                q, k = (apply_rotary_pos_emb(rotary_emb, tensor) for tensor in (q, k))
            else:
                q_rotary_emb = self.rotary(torch.arange(L, device=q.device))
                k_rotary_emb = self.rotary(torch.arange(S, device=k.device))
                q = apply_rotary_pos_emb(q_rotary_emb, q)
                k = apply_rotary_pos_emb(k_rotary_emb, k)

        attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device)

        if mask is not None:
            attn_bias.masked_fill_(mask.logical_not(), float("-inf"))

        out = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=attn_bias,
            dropout_p=self.attn_drop.p if self.training else 0.0,
        )

        out = rearrange(out, "b h n d -> b n (h d)")
        return self.o_dropout(self.o_proj(out))

    def decode_step(self, x, *, cache, positions: torch.Tensor):
        if x.size(1) <= 0:
            raise ValueError("MultiHeadAttention.decode_step expects a non-empty input.")
        if positions.ndim != 1 or positions.size(0) != x.size(1):
            raise ValueError(
                "MultiHeadAttention.decode_step positions must match the decode block length."
            )

        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
        q, k = self.q_norm(q), self.k_norm(k)
        block_len = q.size(2)

        if self.rotary_bias:
            rotary_emb = self.rotary(positions)
            q = apply_rotary_pos_emb(rotary_emb, q)
            k = apply_rotary_pos_emb(rotary_emb, k)

        cached_k, cached_v = cache
        cached_k.index_copy_(2, positions, k)
        cached_v.index_copy_(2, positions, v)

        cache_capacity = cached_k.size(2)
        key_positions = torch.arange(
            cache_capacity,
            device=x.device,
            dtype=torch.long,
        ).unsqueeze(0)
        query_positions = positions.unsqueeze(1)
        causal_mask = key_positions <= query_positions
        valid_mask = key_positions <= positions[-1]
        attn_bias = torch.zeros(
            q.size(0),
            self.num_heads,
            block_len,
            cache_capacity,
            dtype=q.dtype,
            device=q.device,
        )
        attn_bias.masked_fill_(
            (causal_mask & valid_mask).unsqueeze(0).unsqueeze(0).logical_not(),
            float("-inf"),
        )

        out = F.scaled_dot_product_attention(
            q,
            cached_k,
            cached_v,
            attn_mask=attn_bias,
            dropout_p=self.attn_drop.p if self.training else 0.0,
        )
        out = rearrange(out, "b h n d -> b n (h d)")
        return self.o_dropout(self.o_proj(out)), cache