import logging
import os
import warnings

import torch
from torch import Tensor, nn
from torch.nn import init
from torch.nn.functional import softmax

logger = logging.getLogger("dinov2")


XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
try:
   if XFORMERS_ENABLED:
      from xformers.ops import memory_efficient_attention, unbind  # type: ignore

      XFORMERS_AVAILABLE = True
   else:
      warnings.warn("xFormers is disabled (Attention)")
      raise ImportError
except ImportError:
   XFORMERS_AVAILABLE = False
   warnings.warn("xFormers is not available (Attention)")


class Attention(nn.Module):
   def __init__(
      self,
      dim: int,
      num_heads: int = 8,
      qkv_bias: bool = False,
      proj_bias: bool = True,
      attn_drop: float = 0.0,
      proj_drop: float = 0.0,
   ) -> None:
      super().__init__()
      self.num_heads = num_heads
      head_dim = dim // num_heads
      self.scale = head_dim**-0.5

      self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
      self.attn_drop = nn.Dropout(attn_drop)
      self.proj = nn.Linear(dim, dim, bias=proj_bias)
      self.proj_drop = nn.Dropout(proj_drop)

   def forward(self, x: Tensor) -> Tensor:
      B, N, C = x.shape
      qkv = (
         self.qkv(x)
         .reshape(B, N, 3, self.num_heads, C // self.num_heads)
         .permute(2, 0, 3, 1, 4)
      )

      q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
      attn = q @ k.transpose(-2, -1)

      attn = attn.softmax(dim=-1)
      attn = self.attn_drop(attn)

      x = (attn @ v).transpose(1, 2).reshape(B, N, C)
      x = self.proj(x)
      x = self.proj_drop(x)
      return x


class MemEffAttention(Attention):
   def forward(self, x: Tensor, attn_bias=None) -> Tensor:
      if not XFORMERS_AVAILABLE:
         if attn_bias is not None:
               raise AssertionError("xFormers is required for using nested tensors")
         return super().forward(x)

      B, N, C = x.shape
      qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)

      q, k, v = unbind(qkv, 2)

      x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
      x = x.reshape([B, N, C])

      x = self.proj(x)
      x = self.proj_drop(x)
      return x


class AttentionBlock(nn.Module):
   def __init__(self, state_dim):
      super().__init__()

      q = nn.Parameter(torch.randn(1, state_dim, 1))

      init.xavier_uniform_(q)

      self.q = q

      self.state_dim = state_dim

   def forward(self, x):
      x = x.view(tuple(x.shape[:2]) + (-1,))

      r = self.q * x

      unnormalized_attn_map = r.sum(1)

      attn_map = softmax(unnormalized_attn_map, dim=1)

      a = (torch.unsqueeze(attn_map, 1) * x).sum(2)

      return a