# Copyright (c) 2022-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.

from typing import Optional

import torch
from torch.autograd.function import once_differentiable

from ..utils import maybe_empty_like, maybe_detach
from ..graph import HeteroCSC
from pylibcugraphops.operators import (
    agg_hg_basis_n2n_post_fwd,
    agg_hg_basis_n2n_post_bwd,
)


def agg_hg_basis_n2n_post(
    feat: torch.Tensor,
    weights_comb: Optional[torch.Tensor],
    graph: HeteroCSC,
    concat_own: bool = False,
    norm_by_out_degree: bool = False,
) -> torch.Tensor:
    """
    PyTorch autograd function for node-to-node RGCN-like basis regularized
    aggregation, with features being transformed after (post) this aggregation.

    Parameters
    ----------
    feat : torch.Tensor
        The input node features.
        Shape: ``(n_src_nodes, dim_in)``.

    weights_comb : Optional[torch.Tensor], default=None
        The combination weights. If `None`, no weights are used and features of
        different edge types are aggregated into separate output dimensions.
        Shape: ``(n_edge_types, n_bases)`` if provided.

    graph : HeteroCSC
        The graph used for the operation.

    concat_own : bool, default=False
        Concatenate output node embeddings in the aggregation.

    norm_by_out_degree : bool, default=False
        If set, output embeddings are normed by the degree of the output node
        per edge type.

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape for ``concat_own=False``: ``(n_dst_nodes, dim_out)``.
        Shape for ``concat_own=True``: ``(n_dst_nodes, dim_out + dim_in)``,
        where ``dim_out = dim_in * n_bases`` if ``weights_comb`` is provided,
        and ``dim_out = dim_in * n_edge_types`` if not provided.
    """
    options = (concat_own, norm_by_out_degree)

    if graph.is_bipartite:
        raise NotImplementedError(
            "This operator currently does not implement support for bipartite graphs."
        )

    return _agg_hg_basis_n2n_post_autograd.apply(feat, weights_comb, graph, options)


class _agg_hg_basis_n2n_post_autograd(torch.autograd.Function):
    """
    Custom autograd function for node-to-node full-graph RGCN-like basis
    regularized aggregation, with features being transformed after (post)
    this aggregation.
    """

    @staticmethod
    def forward(ctx, feat, weights_comb, graph, options):
        concat_own, norm_by_out_degree = options
        dim = feat.size(-1)
        n_bases = (
            graph.num_edge_types if weights_comb is None else weights_comb.size(-1)
        )
        dim_out = dim * (n_bases + (1 if concat_own else 0))

        agg_output = torch.empty(
            graph.num_dst_nodes, dim_out, dtype=feat.dtype, device=feat.device
        )

        stream = torch.cuda.current_stream().cuda_stream

        feat = feat.detach().contiguous()
        weights_comb = maybe_detach(weights_comb)
        agg_hg_basis_n2n_post_fwd(
            agg_output,
            feat,
            weights_comb,
            graph.graph_csc,
            concat_own,
            norm_by_out_degree,
            stream_id=stream,
        )

        ctx.graph = graph
        ctx.concat_own = concat_own
        ctx.norm_by_out_degree = norm_by_out_degree
        ctx.save_for_backward(feat, weights_comb)

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        feat, weights_comb = ctx.saved_tensors
        needs_grad_feat, needs_grad_w, _, _ = ctx.needs_input_grad
        if not (needs_grad_feat or needs_grad_w):
            return None, None, None, None

        grad_feat = maybe_empty_like(feat, needs_grad_feat)

        stream = torch.cuda.current_stream().cuda_stream
        if weights_comb is None:
            grad_weights_comb = None
        else:
            grad_weights_comb = maybe_empty_like(weights_comb, needs_grad_w)

        agg_hg_basis_n2n_post_bwd(
            grad_feat,
            grad_weights_comb,
            grad_output,
            feat,
            weights_comb,
            ctx.graph.graph_csc,
            ctx.concat_own,
            ctx.norm_by_out_degree,
            stream_id=stream,
        )

        return grad_feat, grad_weights_comb, None, None
