# Copyright (c) 2023-2024, NVIDIA CORPORATION.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
from typing import Optional
from torch.autograd.function import once_differentiable

from ..utils import (
    maybe_empty_like,
    maybe_detach,
    maybe_empty,
    get_tensor_meta_data,
    maybe_zeros,
    maybe_zeros_like,
)
from ..graph import CSC
from pylibcugraphops import mha_params
from pylibcugraphops.operators import (
    activation_params,
    mha_simple_n2n_fwd,
    mha_simple_n2n_bwd,
    mha_simple_n2n_efeat_fwd,
    mha_simple_n2n_efeat_bwd,
)


def mha_simple_n2n(
    key_emb: torch.Tensor,
    query_emb: torch.Tensor,
    value_emb: torch.Tensor,
    graph: CSC,
    num_heads: int = 1,
    concat_heads: bool = True,
    edge_emb: Optional[torch.Tensor] = None,
    norm_by_dim: bool = True,
    score_bias: Optional[torch.Tensor] = None,
    score_weight: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    r"""PyTorch autograd function for a multi-head attention layer (mha_simple)
    without using cudnn with an activation prior to the dot
    product but none afterwards in a node-to-node reduction (n2n).

    Parameters
    ----------
    key_emb : torch.Tensor
        The key embeddings of input nodes. Each row consists of concatenated features
        from different heads after the linear transformation.
        Shape: (n_src_nodes, dim_in) where dim_in = dim_head * num_heads, with
                dim_head being the feature dimension per head.

    query_emb : torch.Tensor
        The query embeddings of input nodes. Each row consists of concatenated features
        from different heads after the linear transformation.
        Shape: (n_dst_nodes, dim_in) where dim_in = dim_head * num_heads, with
                dim_head being the feature dimension per head.

    value_emb : torch.Tensor
        The value embeddings of input nodes. Each row consists of concatenated features
        from different heads after the linear transformation.
        Shape: (n_src_nodes, dim_in) where dim_in = dim_head * num_heads, with
                dim_head being the feature dimension per head.

    graph : CSC
        The graph used in for the operation.

    num_heads : int, default=1
        Number of heads in multi-head attention.

    concat_heads : bool, default=True
        Aggregated embeddings from each head are concatenated if `True` or
        averaged if `False`.

    edge_emb : Optional[torch.Tensor], default=None
        Optional input edge embeddings. Each row consists of concatenated
        features from different heads after the linear transformation.
        Shape: (n_edges, dim_in).

    norm_by_dim : bool, default=True
        flag indicating whether or not to normalize the result of
        the key-query dot-product with the square root of dim_in

    score_bias : Optional[torch.Tensor], default=None
        Optional additive bias which is added onto the result of the
        key-query dot-product before the softmax.

    score_weight : Optional[torch.Tensor], default=None
        Optional multplicative bias applied onto the result of the key-query
        dot-product before the softmax (prior to score_bias if also applied).

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: (n_dst_nodes, dim_in) if `concat_heads` is `True`;
               (n_dst_nodes, dim_head) otherwise.
    """

    params = mha_params(activation_params(), num_heads, concat_heads)
    if not graph.is_bipartite:
        raise ValueError(
            "Expected bipartite with key_emb, query_emb,"
            " and value_emb being distinct buffers."
        )
    return _mha_simple_n2n_autograd.apply(
        key_emb,
        query_emb,
        value_emb,
        edge_emb,
        score_bias,
        score_weight,
        (graph, params, norm_by_dim),
    )


class _mha_simple_n2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for a multi-head attention layer (mha_simple)
    without using cudnn (mha_gat_v2) operating on a graph
    in a node-to-node reduction (n2n)."""

    @staticmethod
    def forward(
        ctx, key_emb, query_emb, value_emb, edge_emb, score_bias, score_weight, params
    ):
        graph, params, norm_by_dim = params
        dim = value_emb.size(-1)
        dim_out = dim if params.concat_heads else (dim // params.num_heads)
        output = torch.empty(
            graph.num_dst_nodes, dim_out, dtype=value_emb.dtype, device=value_emb.device
        )
        sm_scores = torch.empty(
            2 if score_weight is None else 3,
            params.num_heads,
            graph.num_indices,
            dtype=value_emb.dtype,
            device=value_emb.device,
        )

        stream = torch.cuda.current_stream().cuda_stream

        key_emb = key_emb.detach().contiguous()
        query_emb = query_emb.detach().contiguous()
        value_emb = value_emb.detach().contiguous()
        edge_emb = maybe_detach(edge_emb)
        score_weight = maybe_detach(score_weight)
        score_bias = maybe_detach(score_bias)

        if graph.num_dst_nodes > 0:
            if edge_emb is None:
                mha_simple_n2n_fwd(
                    output,
                    sm_scores,
                    key_emb,
                    query_emb,
                    value_emb,
                    graph.graph_csc,
                    params,
                    norm_by_dim,
                    score_bias,
                    score_weight,
                    stream_id=stream,
                )
            else:
                mha_simple_n2n_efeat_fwd(
                    output,
                    sm_scores,
                    key_emb,
                    query_emb,
                    value_emb,
                    edge_emb,
                    graph.graph_csc,
                    params,
                    norm_by_dim,
                    score_bias,
                    score_weight,
                    stream_id=stream,
                )

        ctx.graph = graph
        ctx.params = params
        ctx.norm_by_dim = norm_by_dim
        ctx.score_bias_meta_data = get_tensor_meta_data(score_bias)
        ctx.save_for_backward(
            key_emb, query_emb, value_emb, edge_emb, sm_scores, score_weight
        )

        return output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        (
            key_emb,
            query_emb,
            value_emb,
            edge_emb,
            sm_scores,
            score_weight,
        ) = ctx.saved_tensors
        (
            needs_grad_key_emb,
            needs_grad_query_emb,
            needs_grad_value_emb,
            needs_grad_edge_emb,
            needs_grad_score_bias,
            needs_grad_score_weight,
            _,
        ) = ctx.needs_input_grad
        if not (
            needs_grad_key_emb
            or needs_grad_query_emb
            or needs_grad_value_emb
            or needs_grad_edge_emb
            or needs_grad_score_bias
            or needs_grad_score_weight
        ):
            return None, None, None, None, None, None

        grad_output = grad_output.detach().contiguous()

        if ctx.graph.num_dst_nodes > 0:
            grad_key_emb = maybe_empty_like(key_emb, needs_grad_key_emb)
            grad_query_emb = maybe_empty_like(query_emb, needs_grad_query_emb)
            grad_value_emb = maybe_empty_like(value_emb, needs_grad_value_emb)
            grad_sm_scores = torch.empty_like(sm_scores)

            # These two should be created if edge_emb and score_bias exist in
            # in the forward pass However, if they are passed as None, they also
            # don't require a gradient. Thus needs_grad_x should capture both
            # them existing in the forward pass but not requiring a gradient
            # and not existing in the first place.
            grad_edge_emb = maybe_empty_like(edge_emb, needs_grad_edge_emb)
            grad_score_bias = maybe_empty(
                **ctx.score_bias_meta_data, create_tensor=needs_grad_score_bias
            )
            grad_score_weight = maybe_empty_like(score_weight, needs_grad_score_weight)

            stream = torch.cuda.current_stream().cuda_stream

            if edge_emb is None:
                mha_simple_n2n_bwd(
                    grad_key_emb,
                    grad_query_emb,
                    grad_value_emb,
                    grad_sm_scores,
                    grad_output,
                    key_emb,
                    query_emb,
                    value_emb,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    ctx.norm_by_dim,
                    score_weight,
                    grad_score_bias,
                    grad_score_weight,
                    stream_id=stream,
                )

            else:
                mha_simple_n2n_efeat_bwd(
                    grad_key_emb,
                    grad_query_emb,
                    grad_value_emb,
                    grad_edge_emb,
                    grad_sm_scores,
                    grad_output,
                    key_emb,
                    query_emb,
                    value_emb,
                    edge_emb,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    ctx.norm_by_dim,
                    score_weight,
                    grad_score_bias,
                    grad_score_weight,
                    stream_id=stream,
                )
        else:
            grad_key_emb = maybe_zeros_like(key_emb, needs_grad_key_emb)
            grad_query_emb = maybe_zeros_like(query_emb, needs_grad_query_emb)
            grad_value_emb = maybe_zeros_like(value_emb, needs_grad_value_emb)
            grad_edge_emb = maybe_zeros_like(edge_emb, needs_grad_edge_emb)
            grad_score_bias = maybe_zeros(
                **ctx.score_bias_meta_data, create_tensor=needs_grad_score_bias
            )
            grad_score_weight = maybe_zeros_like(score_weight, needs_grad_score_weight)

        return (
            grad_key_emb,
            grad_query_emb,
            grad_value_emb,
            grad_edge_emb,
            grad_score_bias,
            grad_score_weight,
            None,
        )
