from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from fastai.torch_core import Module
from torch import Tensor

from fedot_ind.core.architecture.settings.computational import backend_methods as np


class ScaledDotProductAttention(Module):
    """Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual
    attention from previous layer (Realformer: Transformer likes residual attention by He et al, 2020) and locality
    self attention (Vision Transformer for Small-Size Datasets by Lee et al, 2021)

    """

    def __init__(self, d_model, n_heads, attn_dropout=0., res_attention=False, lsa=False):
        self.attn_dropout = nn.Dropout(attn_dropout)
        self.res_attention = res_attention
        head_dim = d_model // n_heads
        self.scale = nn.Parameter(torch.tensor(
            head_dim ** -0.5), requires_grad=lsa)
        self.lsa = lsa

    def forward(self, q: Tensor, k: Tensor, v: Tensor, prev: Optional[Tensor] = None,
                key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):
        """
        Method for forward pass of scaled dot-product attention.

        Args:
            q: [bs x n_heads x max_q_len x d_k]
            k: [bs x n_heads x d_k x seq_len]
            v: [bs x n_heads x seq_len x d_v]
            prev: [bs x n_heads x q_len x seq_len]
            key_padding_mask: [bs x seq_len]
            attn_mask: [1 x seq_len x seq_len]

        Returns:
            output: [bs x n_heads x max_q_len x d_v]
            attn: [bs x n_heads x max_q_len x seq_len]
            scores: [bs x n_heads x max_q_len x seq_len]

        """

        # Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence
        # attn_scores : [bs x n_heads x max_q_len x q_len]
        attn_scores = torch.matmul(q, k) * self.scale

        # Add pre-softmax attention scores from the previous layer (optional)
        if prev is not None:
            attn_scores = attn_scores + prev

        # Attention mask (optional)
        # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len
        if attn_mask is not None:
            if attn_mask.dtype == torch.bool:
                attn_scores.masked_fill_(attn_mask, -np.inf)
            else:
                attn_scores += attn_mask

        # Key padding mask (optional)
        # mask with shape [bs x q_len] (only when max_w_len == q_len)
        if key_padding_mask is not None:
            attn_scores.masked_fill_(
                key_padding_mask.unsqueeze(1).unsqueeze(2), -np.inf)

        # normalize the attention weights
        # attn_weights   : [bs x n_heads x max_q_len x q_len]
        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.attn_dropout(attn_weights)

        # compute the new values given the attention weights
        # output: [bs x n_heads x max_q_len x d_v]
        output = torch.matmul(attn_weights, v)

        if self.res_attention:
            return output, attn_weights, attn_scores
        else:
            return output, attn_weights


class MultiHeadAttention(Module):
    def __init__(self, d_model,
                 n_heads,
                 d_k=None,
                 d_v=None,
                 res_attention=False,
                 attn_dropout=0.,
                 proj_dropout=0.,
                 qkv_bias=True,
                 lsa=False):
        """Multi Head Attention Layer

        Args:
            d_model: model dimensionality
            n_heads: number of heads
            d_k: dimensionality of K and Q
            d_v: dimensionality of V
            res_attention: whether to use residual attention from previous layer
            attn_dropout: dropout for attention weights
            proj_dropout: dropout for output
            qkv_bias: whether to use bias for q, k, v projections
            lsa: whether to use learnable scale for attention scores

        """

        if d_k is None:
            d_k = d_model // n_heads
        if d_v is None:
            d_v = d_model // n_heads

        self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v

        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=qkv_bias)

        # Scaled Dot-Product Attention (multiple heads)
        self.res_attention = res_attention
        self.sdp_attn = ScaledDotProductAttention(d_model, n_heads, attn_dropout=attn_dropout,
                                                  res_attention=self.res_attention, lsa=lsa)

        # Poject output
        self.to_out = nn.Sequential(
            nn.Linear(n_heads * d_v, d_model), nn.Dropout(proj_dropout))

    def forward(self, Q: Tensor, K: Optional[Tensor] = None, V: Optional[Tensor] = None, prev: Optional[Tensor] = None,
                key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):

        bs = Q.size(0)
        if K is None:
            K = Q
        if V is None:
            V = Q

        # Linear (+ split in multiple heads)
        q_s = self.W_Q(Q).view(bs, -1, self.n_heads, self.d_k).transpose(1,
                                                                         2)
        # q_s    : [bs x n_heads x max_q_len x d_k]
        k_s = self.W_K(K).view(bs, -1, self.n_heads, self.d_k).permute(0, 2, 3,
                                                                       1)
        # k_s    : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)
        # v_s    : [bs x n_heads x q_len x d_v]
        v_s = self.W_V(V).view(bs, -1, self.n_heads, self.d_v).transpose(1, 2)

        # Apply Scaled Dot-Product Attention (multiple heads)
        if self.res_attention:
            output, attn_weights, attn_scores = self.sdp_attn(q_s, k_s, v_s, prev=prev,
                                                              key_padding_mask=key_padding_mask, attn_mask=attn_mask)
        else:
            output, attn_weights = self.sdp_attn(
                q_s, k_s, v_s, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
        # output: [bs x n_heads x q_len x d_v],
        # attn: [bs x n_heads x q_len x q_len],
        # scores: [bs x n_heads x max_q_len x q_len]

        # back to the original inputs dimensions
        output = output.transpose(1, 2).contiguous().view(bs, -1,
                                                          self.n_heads * self.d_v)
        # output: [bs x q_len x n_heads * d_v]
        output = self.to_out(output)

        if self.res_attention:
            return output, attn_weights, attn_scores
        else:
            return output, attn_weights
