# 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.

import torch
from torch.autograd.function import once_differentiable

from ..utils import maybe_empty, get_tensor_meta_data
from ..graph import StaticCSC, UnsupportedGraphError
from pylibcugraphops.operators import (
    AggOp,
    pool_fg_n2s_fwd,
    pool_fg_n2s_bwd,
)

_aggr_ops = ["max", "mean", "min", "sum"]


def pool_n2s(
    feat: torch.Tensor,
    graph: StaticCSC,
    aggr: str = "sum",
    use_atomics: bool = False,
) -> torch.Tensor:
    r"""PyTorch autograd function which pools node features from a graph
    into a scalar (pool_n2s).

    Parameters
    ----------
    feat : torch.Tensor
        The input node features.
        Shape: (n_in_nodes, dim).

    graph : StaticCSC
        The graph used for the operation.

    aggr : str, default="sum"
        Aggregation operation. Choose from 'max', 'mean', 'min', 'sum'.

    use_atomics : bool, default=False
        Flag indicating whether atomics operations are used or a workspace
        in global memory. It is only valid for 'mean' and 'sum' aggregations
        and might change certain performance characteristics. It defaults
        to ``False``, i.e. a global workspace is used by default.

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: (1, dim)
    """
    if aggr not in _aggr_ops:
        raise ValueError(
            f"Supported aggregation operations: {_aggr_ops}, " f"but got '{aggr}'."
        )
    aggr = getattr(AggOp, aggr.capitalize())
    if isinstance(graph, StaticCSC):
        return _pool_fg_n2s_autograd.apply(feat, graph, aggr, use_atomics)
    else:
        raise UnsupportedGraphError(StaticCSC, graph)


class _pool_fg_n2s_autograd(torch.autograd.Function):
    """Custom autograd function which pools node features from a graph
    into a scalar (pool_n2s) operating on a full graph (fg)."""

    @staticmethod
    def forward(ctx, feat, graph, aggr, use_atomics):
        dim = feat.size(-1)
        fwd_graph = graph._fwd_graph
        agg_output = torch.empty(1, dim, dtype=feat.dtype, device=feat.device)

        if aggr in (AggOp.Max, AggOp.Min):
            idx_type = fwd_graph.__class__.__name__.split("_")[-1]
            out_pos = torch.empty(
                1, dim, dtype=getattr(torch, idx_type), device=feat.device
            )
            # use_atomics only valid for mean/sum
            # silently set this here in case it has been passed
            use_atomics = False
        else:
            out_pos = None

        stream = torch.cuda.current_stream().cuda_stream
        feat = feat.detach().contiguous()
        w_size = pool_fg_n2s_fwd(
            agg_output,
            feat,
            fwd_graph,
            use_atomics,
            None,
            -1,
            aggr,
            out_pos,
            stream_id=stream,
        )

        workspace = maybe_empty(
            w_size, dtype=torch.int64, device=feat.device, create_tensor=(w_size > 0)
        )

        pool_fg_n2s_fwd(
            agg_output,
            feat,
            fwd_graph,
            use_atomics,
            workspace,
            w_size,
            aggr,
            out_pos,
            stream_id=stream,
        )

        ctx.bwd_graph = graph._bwd_graph
        ctx.aggr = aggr
        ctx.feat_meta_data = get_tensor_meta_data(feat)
        ctx.save_for_backward(out_pos)

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        (out_pos,) = ctx.saved_tensors
        needs_feat_grad, _, _, _ = ctx.needs_input_grad
        if not needs_feat_grad:
            return None, None, None

        grad_feat = maybe_empty(**ctx.feat_meta_data, create_tensor=True)
        grad_output = grad_output.detach().contiguous()
        stream = torch.cuda.current_stream().cuda_stream

        pool_fg_n2s_bwd(
            grad_feat,
            grad_output,
            ctx.bwd_graph,
            ctx.aggr,
            out_pos,
            stream_id=stream,
        )

        return grad_feat, None, None, None
