#################################################################################################
# Copyright (c) 2022-2024 Ali Hassani.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
#################################################################################################
from typing import Optional

import torch
from torch import nn, Tensor
from torch.nn.functional import pad
from torch.nn.init import trunc_normal_

from .functional import na3d_av, na3d_qk_with_bias


class NeighborhoodAttention3D(nn.Module):
    """
    Neighborhood Attention 3D Module
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        kernel_size: int,
        kernel_size_d: Optional[int] = None,
        dilation: int = 1,
        dilation_d: Optional[int] = None,
        bias: bool = True,
        qkv_bias: bool = True,
        qk_scale: Optional[float] = None,
        attn_drop: float = 0.0,
        proj_drop: float = 0.0,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // self.num_heads
        self.scale = qk_scale or self.head_dim**-0.5
        kernel_size_d = kernel_size_d or kernel_size
        dilation = dilation or 1
        dilation_d = dilation_d or dilation
        assert (
            kernel_size > 1 and kernel_size % 2 == 1
        ), f"Kernel size must be an odd number greater than 1, got {kernel_size}."
        assert (
            dilation >= 1
        ), f"Dilation must be greater than or equal to 1, got {dilation}."
        assert (
            dilation_d >= 1
        ), f"Dilation (depth) must be greater than or equal to 1, got {dilation_d}."
        self.kernel_size = kernel_size
        self.kernel_size_d = kernel_size_d
        self.dilation = dilation
        self.dilation_d = dilation_d
        self.window_size = self.kernel_size * self.dilation
        self.window_size_d = self.kernel_size_d * self.dilation_d

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        if bias:
            self.rpb = nn.Parameter(
                torch.zeros(
                    num_heads,
                    (2 * kernel_size_d - 1),
                    (2 * kernel_size - 1),
                    (2 * kernel_size - 1),
                )
            )
            trunc_normal_(self.rpb, std=0.02, mean=0.0, a=-2.0, b=2.0)
        else:
            self.register_parameter("rpb", None)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x: Tensor) -> Tensor:
        if x.dim() != 5:
            raise ValueError(
                f"NeighborhoodAttention2D expected a rank-5 input tensor; got {x.dim()=}."
            )

        B, D, H, W, C = x.shape
        # Pad if the input is small than the minimum supported size
        D_padded, H_padded, W_padded = D, H, W
        padding_d = padding_h = padding_w = 0
        if D < self.window_size_d or H < self.window_size or W < self.window_size:
            padding_d = max(0, self.window_size_d - D_padded)
            padding_h = max(0, self.window_size - H_padded)
            padding_w = max(0, self.window_size - W_padded)
            x = pad(x, (0, 0, 0, padding_w, 0, padding_h))
            _, D_padded, H_padded, W_padded, _ = x.shape

        qkv = (
            self.qkv(x)
            .reshape(B, D_padded, H_padded, W_padded, 3, self.num_heads, self.head_dim)
            .permute(4, 0, 5, 1, 2, 3, 6)
        )
        q, k, v = qkv[0], qkv[1], qkv[2]
        q = q * self.scale
        attn = na3d_qk_with_bias(
            q,
            k,
            self.rpb,
            self.kernel_size_d,
            self.kernel_size,
            self.dilation_d,
            self.dilation,
        )
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = na3d_av(
            attn,
            v,
            self.kernel_size_d,
            self.kernel_size,
            self.dilation_d,
            self.dilation,
        )
        x = x.permute(0, 2, 3, 4, 1, 5).reshape(B, D_padded, H_padded, W_padded, C)

        # Remove padding, if added any
        if padding_d or padding_h or padding_w:
            x = x[:, :D, :H, :W, :]

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

    def extra_repr(self) -> str:
        return (
            f"head_dim={self.head_dim}, num_heads={self.num_heads}, "
            + f"kernel_size_d={self.kernel_size_d}, dilation_d={self.dilation_d}, "
            + f"kernel_size={self.kernel_size}, dilation={self.dilation}, "
            + f"has_bias={self.rpb is not None}"
        )
