# GENERAL attention with flash attention flavor
from typing import Optional, Tuple

import einops
import torch
import torch.nn as nn
from flash_attn import flash_attn_func

from src.model.dino import apply_rotary_pos_emb, rotate_half


class PlainAttention(nn.Module):
    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout_rate: float = 0.1,
        query_bias: bool = True,
        key_bias: bool = True,
        value_bias: bool = True,
        proj_bias: bool = True,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = self.embed_dim // self.num_heads

        self.dropout = dropout_rate

        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=query_bias)
        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=key_bias)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=value_bias)
        self.o_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=proj_bias)

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor | None = None,
        v: torch.Tensor | None = None,
        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        k = k if k is not None else q
        v = v if v is not None else q
        assert k is not None and v is not None

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

        if position_embeddings is not None:
            cos, sin = position_embeddings
            q, k = apply_rotary_pos_emb(q, k, cos, sin)  # type: ignore

        # flash attn handles sftmx scaling too
        o = flash_attn_func(
            q.transpose(1, 2),
            k.transpose(1, 2),  # type: ignore
            v.transpose(1, 2),  # type: ignore
            dropout_p=self.dropout if self.training else 0.0,
            causal=False,
        )

        return self.o_proj(einops.rearrange(o, "b s h d -> b s (h d)").contiguous())  # type: ignore


class CrossAttention(nn.Module):
    def __init__(self, query_dim: int, key_dim: int, value_dim: int):
        super().__init__()
        self.attention = nn.MultiheadAttention(
            embed_dim=query_dim,
            vdim=value_dim,
            num_heads=1,
            dropout=0.0,
            batch_first=True,
        )

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
    ):
        _, attn_score = self.attention(q, k, v, average_attn_weights=True)
        return torch.einsum("b i j, b j d -> b i d", attn_score, v)


class RoPE(nn.Module):
    def __init__(
        self,
        dim: int,
        theta: int = 100,
    ):
        super().__init__()
        self.dim = dim
        self.theta = theta
        self.freqs = nn.Parameter(torch.empty(2, self.dim))

    def _device_weight_init(self):
        # Create freqs in 1d
        freqs_1d = self.theta ** torch.linspace(0, -1, self.dim // 4)
        # duplicate freqs for rotation pairs of channels
        freqs_1d = torch.cat([freqs_1d, freqs_1d])
        # First half of channels do x, second half do y
        freqs_2d = torch.zeros(2, self.dim)
        freqs_2d[0, : self.dim // 2] = freqs_1d
        freqs_2d[1, -self.dim // 2 :] = freqs_1d
        # it's an angular freq here
        self.freqs.data.copy_(freqs_2d * 2 * torch.pi)

    def forward(self, x: torch.Tensor, coords: torch.Tensor) -> torch.Tensor:
        angle = coords @ self.freqs
        return x * angle.cos() + rotate_half(x) * angle.sin()
