# Copyright (c) 2023, 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 CuGraphAnyCsr
from ...binding.pylibcugraphops_ext import ActivationOp, mha_params
from ...binding.pylibcugraphops_ext.operators import (
    activation_params,
    mha_gat_v2_fg_n2n_fwd as fg_n2n_fwd,
    mha_gat_v2_fg_n2n_bwd as fg_n2n_bwd,
)

_activation_ops = [x for x in dir(ActivationOp) if not x.startswith("__")]


def mha_gat_v2_n2n(
    feat: torch.Tensor,
    attn_weights: torch.Tensor,
    graph: CuGraphAnyCsr,
    num_heads: int = 1,
    activation: str = "LeakyReLU",
    negative_slope: float = 0.2,
    add_own_node: bool = False,
    concat_heads: bool = True,
) -> torch.Tensor:
    r"""PyTorch autograd function for a multi-head attention layer (GAT-like)
    without using cudnn (mha_gat_v2) with an activation prior to the dot
    product but none afterwards in a node-to-node reduction (n2n).

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

    attn_weights : torch.Tensor
        The attention weights.
        Shape: (dim,).

    graph : fg_csr_int[32|64] | mfg_csr_int[32|64]
        The graph used in for the operation.

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

    activation : str, default="LeakyReLU"
        The activation function used in the attention mechanism. Choose from
        'ELU', 'LeakyReLU', 'Linear', 'ReLU', 'Scalar', 'Sigmoid' and 'Tanh'.

    negative_slope : float, default=0.2
        LeakyReLU angle of negative slope. Only effective when `activation` is
        `"LeakyReLU"`.

    add_own_node : bool, default=False
        Add self-loop for each node if set.

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

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

    Notes
    -----
    Replace n_in_nodes and n_out_nodes with n_nodes when given a full graph.
    """
    if activation not in _activation_ops:
        raise ValueError(
            f"Supported activation functions: {_activation_ops}, "
            f"but got '{activation}'."
        )
    act_params = activation_params(getattr(ActivationOp, activation), negative_slope)
    params = mha_params(num_heads, act_params, add_own_node, concat_heads)

    use_mfg = hasattr(graph, "sample_size")
    if use_mfg:
        raise NotImplementedError
    return _mha_gat_v2_fg_n2n_autograd.apply(feat, attn_weights, graph, params)


class _mha_gat_v2_fg_n2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for a multi-head attention layer (GAT-like)
    without using cudnn (mha_gat_v2) with an activation prior to the dot
    product but none afterwards operating on full graphs (fg) in a node-to-node
    reduction (n2n)."""

    @staticmethod
    def forward(ctx, feat, attn_weights, fg, params):
        feat = feat.contiguous()
        attn_weights = attn_weights.contiguous()
        dim = feat.shape[-1]
        dim_out = dim if params.concat_heads else (dim // params.num_heads)
        output = torch.zeros(
            fg.n_nodes, dim_out, dtype=feat.dtype, device=feat.device
        )
        sm_scores = torch.zeros(
            2, params.num_heads, fg.n_indices, dtype=feat.dtype, device=feat.device
        )
        act_scores = torch.zeros(
            fg.n_indices, dim, dtype=feat.dtype, device=feat.device
        )

        stream = torch.cuda.current_stream().cuda_stream
        fg_n2n_fwd(
            output,
            sm_scores,
            act_scores,
            feat.detach(),
            attn_weights.detach(),
            fg,
            params,
            stream_id=stream,
        )

        ctx.fg = fg
        ctx.params = params
        ctx.save_for_backward(feat, attn_weights, sm_scores, act_scores)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        grad_output = grad_output.contiguous()
        feat, attn_weights, sm_scores, act_scores = ctx.saved_tensors

        grad_feat = torch.zeros_like(feat)
        grad_attn_weights = torch.zeros_like(attn_weights)
        grad_sm_scores = torch.zeros_like(sm_scores)

        stream = torch.cuda.current_stream().cuda_stream
        fg_n2n_bwd(
            grad_feat,
            grad_attn_weights,
            grad_sm_scores,
            grad_output,
            feat.detach(),
            attn_weights.detach(),
            sm_scores,
            act_scores,
            ctx.fg,
            ctx.params,
            stream_id=stream,
        )

        return grad_feat, grad_attn_weights, None, None
