import math
from functools import partial
from typing import Optional, Tuple, Union

import torch
import torch.nn.functional as F
from torch import Tensor

from continual.logging import getLogger
from continual.module import _callmode

from .mha_base import MultiheadAttentionBase, scaled_dot_prod_attn_flops

logger = getLogger(__name__)
logger_once = getLogger(__name__, log_once=True)

State = Tuple[
    Tensor,  # Q_mem, (B, Nt-1, E)
    Tensor,  # K_T_mem, (B, E, Ns)
    Tensor,  # V_mem, (B, Ns, E)
    Tensor,  # stride_index
]


def _scaled_dot_product_attention_default_state(
    batch_size: int,
    sequence_len: int,
    embed_dim: int,
    num_heads: int,
    query_index=-1,
    init_fn=torch.zeros,
    dtype=None,
    device=None,
):
    init_fn = partial(init_fn, dtype=dtype, device=device)
    E = embed_dim // num_heads
    B = batch_size * num_heads
    N = sequence_len
    Nq = sequence_len - query_index - 1 if query_index >= 0 else -query_index - 1
    Q_mem = init_fn((B, Nq, E))
    K_T_mem = init_fn((B, E, N))
    V_mem = init_fn((B, N, E))
    return (Q_mem, K_T_mem, V_mem)


def _scaled_dot_product_attention_step(
    prev_state: State,
    q_step: Tensor,  # step input (B, E)
    k_step: Tensor,  # step input (B, E)
    v_step: Tensor,  # step input (B, E)
    attn_mask: Optional[Tensor] = None,
    dropout_p: float = 0.0,
) -> Tuple[Tensor, State]:
    """
    Computes the Continual Singe-output Scaled Dot-Product Attention on query, key and value tensors.
    Returns attended values and updated states.

    Args:
        q_step, k_step, v_step: query, key and value tensors for a step. See Shape section for shape details.
        attn_mask: optional tensor containing mask values to be added to calculated
            attention. May be 2D or 3D; see Shape section for details.
        dropout_p: dropout probability. If greater than 0.0, dropout is applied.

    Shape:
        - q_step: :math:`(B, E)` where B is batch size and E is embedding dimension.
        - k_step: :math:`(B, E)` where B is batch size and E is embedding dimension.
        - v_step: :math:`(B, E)` where B is batch size and E is embedding dimension.

        - Output: attention values have shape :math:`(B, Nt, E)`; new state
    """
    if attn_mask is not None:  # pragma: no cover
        logger_once.warning(
            "attn_mask is not supported in scaled_dot_product_attention yet and will be skipped"
        )
    if dropout_p != 0.0:  # pragma: no cover
        logger_once.warning(
            "dropout_p is not supported in continual scaled_dot_product_attention yet and will be skipped"
        )

    (
        Q_mem,  # (B, Nq, E)
        K_T_mem,  # (B, E, Ns)
        V_mem,  # (B, Ns, E)
    ) = prev_state

    B, E = q_step.shape
    q_step = q_step / math.sqrt(E)
    q_sel = (Q_mem[:, 0] if Q_mem.shape[1] > 0 else q_step).unsqueeze(1)

    # Update states
    # Note: We're allowing the K and V mem to have one more entry than
    # strictly necessary to simplify computatations.

    K_T_new = torch.roll(K_T_mem, shifts=-1, dims=(2,))
    K_T_new[:, :, -1] = k_step

    V_new = torch.roll(V_mem, shifts=-1, dims=(1,))
    V_new[:, -1] = v_step

    attn = torch.bmm(q_sel, K_T_new)
    attn_sm = F.softmax(attn, dim=-1)

    if dropout_p > 0.0:  # pragma: no cover
        attn_sm = F.dropout(attn_sm, p=dropout_p)

    # (B, Nt, Ns) x (B, Ns, E) -> (B, Nt, E)
    output = torch.bmm(attn_sm, V_new)

    if Q_mem.shape[1] > 0:
        Q_new = torch.roll(Q_mem, shifts=-1, dims=(1,))
        Q_new[:, -1] = q_step
    else:
        Q_new = Q_mem

    new_states = (Q_new, K_T_new, V_new)

    return output, new_states


class SingleOutputMultiheadAttention(MultiheadAttentionBase):
    """
    MultiHeadAttention which only computes the attention output for the a single query during continual inference.

    Continual MHAs were proposed by Hedegaard et al. in
    "Continual Transformers: Redundancy-Free Attention for Online Inference"
    https://arxiv.org/abs/2201.06268 (paper) https://www.youtube.com/watch?v=gy802Tlp-eQ (video).

    This module augments the MultiHeadAttention in PyTorch with
    `forward_step` / `forward_steps` functions, in which one / more
    query, key, and value tokens are passed to yield the multihead attentions, and
    updated outputs are computed for each token input.

    Args:
        embed_dim: total dimension of the model.
        num_heads: parallel attention heads.
        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
        bias: add bias as module parameter. Default: True.
        add_bias_kv: add bias to the key and value sequences at dim=0.
        add_zero_attn: add a new batch of zeros to the key and
                       value sequences at dim=1.
        kdim: total number of features in key. Default: None.
        vdim: total number of features in value. Default: None.
        batch_first: If ``True``, then the input and output tensors are provided
            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
        device: torch device to initialize layer on. Defaults to None.
        dtype: datatype of layer parameters. Defaults to None.
        sequence_len: Length of token sequence.
        query_index: The index of the query to compute the attention. Here, -1 denotes the latest query.
        forward_returns_attn_mask: Whether forward should return attention mask.
        embed_dim_second: Whether the embed dimension should be second.
        single_output_forward: Whether forward should be restricted to compute attention for only one query.

    .. note::
        If :attr:`kdim` and :attr:`vdim` are None, they will be set
        to :attr:`embed_dim` such that query, key, and value have the same
        number of features.

    Examples::

        mha = co.SingleOutputMultiheadAttention(
            embed_dim=512,
            num_heads=8,
            sequence_len=32,
            dropout=0.0,
            batch_first=True,
            embed_dim_second=True,
        )
        x = torch.rand(10, 512, 32)

        out, attn_mask = mha.forward(x)

        # continual inference API
        firsts = mha.forward_steps(x[:,:,:-1])
        last = mha.forward_step(x[:,:,-1])

        assert firsts is None  # The module first needs to observe ``sequence_len`` values
        assert torch.allclose(out[:,:,-1], last, atol=1e-6)
    """

    _state_shape = 4
    #                      Q_mem, K_T_mem, V_mem, stride_index
    _dynamic_state_inds = [True, True, True, False]

    def __init__(
        self,
        embed_dim,
        num_heads,
        dropout=0.0,
        bias=True,
        add_bias_kv=False,
        add_zero_attn=False,
        kdim=None,
        vdim=None,
        batch_first=True,
        device=None,
        dtype=None,
        sequence_len=None,
        query_index=-1,
        forward_returns_attn_mask=True,
        embed_dim_second=False,
        single_output_forward=True,
    ) -> None:
        MultiheadAttentionBase.__init__(
            self,
            embed_dim,
            num_heads,
            dropout,
            bias,
            add_bias_kv,
            add_zero_attn,
            kdim,
            vdim,
            batch_first,
            device,
            dtype,
            sequence_len,
            partial(
                _scaled_dot_product_attention_default_state,
                sequence_len=sequence_len,
                embed_dim=embed_dim,
                num_heads=num_heads,
                query_index=query_index,
            ),
            _scaled_dot_product_attention_step,
            forward_returns_attn_mask,
            embed_dim_second,
        )
        assert query_index < sequence_len
        self.query_index = query_index
        self.single_output_forward = single_output_forward
        self.register_buffer("Q_mem", torch.tensor([]), persistent=False)
        self.register_buffer("K_T_mem", torch.tensor([]), persistent=False)
        self.register_buffer("V_mem", torch.tensor([]), persistent=False)
        self.register_buffer("stride_index", torch.tensor(0), persistent=False)

    def get_state(self) -> Optional[State]:
        if len(self.V_mem) > 0:
            return (
                self.Q_mem,
                self.K_T_mem,
                self.V_mem,
                self.stride_index,
            )
        return None

    def set_state(self, state: State):
        (
            self.Q_mem,
            self.K_T_mem,
            self.V_mem,
            self.stride_index,
        ) = state

    def clean_state(self):
        self.Q_mem = torch.tensor([], device=self.Q_mem.device)
        self.K_T_mem = torch.tensor([], device=self.K_T_mem.device)
        self.V_mem = torch.tensor([], device=self.V_mem.device)
        self.stride_index = torch.tensor(0)

    @property
    def delay(self) -> int:
        return (
            self.sequence_len - self.query_index - 1
            if self.query_index >= 0
            else -self.query_index - 1
        )

    def forward(
        self,
        query: Tensor,
        key: Tensor = None,
        value: Tensor = None,
        key_padding_mask: Optional[Tensor] = None,
        need_weights: bool = True,
        attn_mask: Optional[Tensor] = None,
    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
        if not self.single_output_forward:
            return MultiheadAttentionBase.forward(
                self, query, key, value, key_padding_mask, need_weights, attn_mask
            )

        if key is None:
            key = query
        if value is None:
            value = query

        # Select a single query entry
        if self.batch_first:
            if self.embed_dim_second:
                query = query[:, :, self.query_index].unsqueeze(2)
            else:
                query = query[:, self.query_index].unsqueeze(1)
        else:
            query = query[self.query_index].unsqueeze(0)

        o = MultiheadAttentionBase.forward(
            self, query, key, value, key_padding_mask, need_weights, attn_mask
        )

        return o

    def _forward_step(
        self,
        query: Tensor,
        key: Tensor = None,
        value: Tensor = None,
        prev_state: Optional[State] = None,
    ) -> Tuple[Optional[Tensor], State]:
        """
        Args:
            query, key, value: step_inputs for mapping a query and a set of key-value pairs to an output.
                See "Attention Is All You Need" for more details.

        Shapes for inputs:
            - query: :math:`(N, E)` N is the batch size, E is
              the embedding dimension.
            - key: :math:`(N, E)`, N is the batch size, E is
              the embedding dimension.
            - value: :math:`(N, E)` N is the batch size, E is
              the embedding dimension.

        Shapes for outputs:
            - attn_output: :math:`(N, E)` where N is the batch size,
              E is the embedding dimension. :math:`(N, L, E)` if ``batch_first`` is ``True``.
        """
        if key is None:
            key = query
        if value is None:
            value = query

        o, next_state = MultiheadAttentionBase._forward_step(
            self, query, key, value, prev_state
        )

        if o is not None:
            o = o.squeeze(0)

        return o, next_state

    def forward_step(
        self,
        query: Tensor,
        key: Tensor = None,
        value: Tensor = None,
        update_state=True,
        *args,
        **kwargs,
    ) -> Optional[Tensor]:
        """
        Args:
            query, key, value: step_inputs for mapping a query and a set of key-value pairs to an output.
                See "Attention Is All You Need" for more details.

        Shapes for inputs:
            - query: :math:`(N, E)` N is the batch size, E is
              the embedding dimension.
            - key: :math:`(N, E)`, N is the batch size, E is
              the embedding dimension.
            - value: :math:`(N, E)` N is the batch size, E is
              the embedding dimension.

        Shapes for outputs:
            - attn_output: :math:`(N, E)` where N is the batch size,
              E is the embedding dimension. :math:`(N, L, E)` if ``batch_first`` is ``True``.
        """
        o = MultiheadAttentionBase.forward_step(
            self, query, key, value, update_state, *args, **kwargs
        )
        return o.squeeze(1 if self.batch_first else 0) if isinstance(o, Tensor) else o

    def forward_steps(
        self,
        query: Tensor,
        key: Tensor = None,
        value: Tensor = None,
        update_state=True,
        *args,
        **kwargs,
    ) -> Optional[Tensor]:
        """Forward computation for multiple steps with state initialisation

        Args:
            query (Tensor): query.
            key (Tensor): key.
            value (Tensor): value.
            update_state (bool): Whether internal state should be updated during this operation.

        Returns:
            Tensor: Stepwise layer outputs
        """
        o = MultiheadAttentionBase.forward_steps(
            self, query, key, value, update_state, *args, **kwargs
        )

        if isinstance(o, Tensor):
            o = o.squeeze(2)
            if self.embed_dim_second:
                o = o.transpose(1, 2)  # N T E -> N E T

        return o

    def flops(self, include_muls=True, include_adds=False, include_exps=False):
        f = 0

        # Linear projection
        steps_taken = {
            _callmode("forward"): 1
            if self.single_output_forward
            else self.sequence_len,
            _callmode("forward_step"): 1,
        }[self.call_mode]

        f += (
            steps_taken
            * self.embed_dim
            * self.embed_dim
            * 3  # Assuming equal len for Q, K, and V
        )
        if include_adds:
            f += 3 * steps_taken * self.embed_dim * (self.embed_dim - 1)

        if self.in_proj_bias is not None:
            f += 3 * steps_taken * self.embed_dim

            if include_adds:
                f += 3 * steps_taken * self.embed_dim

        # Multi-head Scaled Dot-Product Attention
        f += self.num_heads * {
            _callmode("forward"): single_output_scaled_dot_prod_attn_flops
            if self.single_output_forward
            else scaled_dot_prod_attn_flops,
            _callmode("forward_step"): single_output_scaled_dot_prod_attn_flops,
        }[self.call_mode](
            self.sequence_len,
            self.embed_dim // self.num_heads,
            include_muls,
            include_adds,
            include_exps,
        )

        # Linear projection
        f += 1 * self.embed_dim * (self.embed_dim + 1)

        return f


def single_output_scaled_dot_prod_attn_flops(
    sequence_len, embed_dim, include_muls=True, include_adds=False, include_exps=False
):
    n = sequence_len
    d = embed_dim

    flops = 0

    if include_muls:
        flops += 2 * n * d + 2 * d
    if include_adds:
        flops += 2 * n * d - d - 1
    if include_exps:
        flops += n

    return flops
