# 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 Optional, Union
from torch.autograd.function import once_differentiable

from ..utils import maybe_empty_like, maybe_empty, get_tensor_meta_data
from ..graph import SampledCSC, StaticCSC, UnsupportedGraphError
from pylibcugraphops.operators import (
    AggOp,
    agg_concat_fg_n2n_fwd as fg_n2n_fwd,
    agg_concat_fg_n2n_bwd as fg_n2n_bwd,
    agg_concat_fg_e2n_fwd as fg_e2n_fwd,
    agg_concat_fg_e2n_bwd as fg_e2n_bwd,
    agg_concat_fg_n2n_e2n_fwd as fg_n2n_e2n_fwd,
    agg_concat_fg_n2n_e2n_bwd as fg_n2n_e2n_bwd,
    agg_concat_weighted_fg_n2n_fwd as fg_n2n_weighted_fwd,
    agg_concat_weighted_fg_n2n_bwd as fg_n2n_weighted_bwd,
    agg_concat_mfg_n2n_fwd as mfg_n2n_fwd,
    agg_concat_mfg_n2n_bwd as mfg_n2n_bwd,
)

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


def agg_concat_n2n(
    feat: torch.Tensor,
    graph: Union[SampledCSC, StaticCSC],
    aggr: str = "sum",
    edge_weight: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    PyTorch autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) while concatenating the original features of
    output at the end (agg_concat).

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

    graph : SampledCSC | StaticCSC
        The graph used for the operation.

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

    edge_weight : Optional[torch.Tensor], default=None
        When passing additional edge weights, the aggregation is done in
        a weighted fashion. Note that for min/max the result after
        weighting features with those edge weights is returned.
        Edge weights are indexed in the same fashion as the indices in ``graph``.
        Shape: ``(n_edges, )``.

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: ``(n_out_nodes, 2 * dim_in)``.
    """

    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, SampledCSC):
        if edge_weight is not None:
            raise NotImplementedError(
                "agg_concat_n2n with edge weights not implemented for type SampledCSC."
            )
        return _agg_concat_mfg_n2n_autograd.apply(feat, graph, aggr)
    elif isinstance(graph, StaticCSC):
        return _agg_concat_fg_n2n_autograd.apply(feat, edge_weight, graph, aggr)

    else:
        raise UnsupportedGraphError([SampledCSC, StaticCSC], graph)


def agg_concat_e2n(
    node_feat: torch.Tensor,
    edge_feat: torch.Tensor,
    graph: StaticCSC,
    aggr: str = "sum",
) -> torch.Tensor:
    r"""PyTorch autograd function for simple aggregation using edge features in
    while concatenating the original features of output at the end (agg_concat).

    Parameters
    ----------
    node_feat : torch.Tensor
        The input node features.
        Shape: ``(n_out_nodes, dim_node)``.

    edge_feat : torch.Tensor
        The input edge features.
        Shape: ``(n_edges, dim_edge)``.

    graph : StaticCSC
        The graph used for the operation.

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

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: ``(n_out_nodes, dim_node + dim_edge)``.
    """
    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 _agg_concat_fg_e2n_autograd.apply(edge_feat, node_feat, graph, aggr)

    else:
        raise UnsupportedGraphError(StaticCSC, graph)


def agg_concat_n2n_e2n(
    node_feat: torch.Tensor,
    edge_feat: torch.Tensor,
    graph: StaticCSC,
    aggr: str = "sum",
) -> torch.Tensor:
    r"""PyTorch autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) and edge features in and edge-to-node
    reduction (e2n) while concatenating the original features of output at the
    end (agg_concat).

    Parameters
    ----------
    node_feat : torch.Tensor
        The input node features.
        Shape: ``(n_out_nodes, dim_node)``.

    edge_feat : torch.Tensor
        The input edge features.
        Shape: ``(n_edges, dim_edge)``.

    graph : StaticCSC
        The graph used for the operation.

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

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: ``(n_out_nodes, dim_node + dim_edge + dim_node)``.
    """
    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 _agg_concat_fg_n2n_e2n_autograd.apply(node_feat, edge_feat, graph, aggr)

    else:
        raise UnsupportedGraphError(StaticCSC, graph)


class _agg_concat_fg_n2n_autograd(torch.autograd.Function):
    """
    Custom autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) on a full graph (fg) while concatenating the
    original features of output at the end (agg_concat).
    """

    @staticmethod
    def forward(ctx, feat, edge_weight, graph, aggr):
        fwd_graph = graph._fwd_graph
        dim_in = feat.size(-1)
        agg_output = torch.empty(
            fwd_graph.n_nodes, 2 * dim_in, dtype=feat.dtype, device=feat.device
        )
        node_deg = None

        if aggr in (AggOp.Max, AggOp.Min):
            idx_type = fwd_graph.__class__.__name__.split("_")[-1]
            out_pos = torch.empty(
                fwd_graph.n_nodes,
                dim_in,
                dtype=getattr(torch, idx_type),
                device=feat.device,
            )
        else:
            out_pos = None

        stream = torch.cuda.current_stream().cuda_stream
        feat = feat.detach().contiguous()
        if edge_weight is None:
            fg_n2n_fwd(agg_output, feat, fwd_graph, aggr, out_pos, stream_id=stream)
        else:
            edge_weight = edge_weight.detach().contiguous()
            node_deg = torch.empty_like(edge_weight)
            fg_n2n_weighted_fwd(
                agg_output,
                feat,
                edge_weight,
                fwd_graph,
                aggr,
                out_pos,
                node_deg,
                stream_id=stream,
            )

        ctx.bwd_graph = graph._bwd_graph
        ctx.aggr = aggr
        ctx.dim_node = dim_in
        ctx.feat_meta_data = get_tensor_meta_data(feat)

        saved_agg_output = agg_output if edge_weight is not None else None
        saved_feat = feat if edge_weight is not None else None
        ctx.save_for_backward(
            saved_feat, edge_weight, node_deg, out_pos, saved_agg_output
        )

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        feat, edge_weight, node_deg, out_pos, agg_output = ctx.saved_tensors
        needs_grad_feat, needs_grad_weight, _, _ = ctx.needs_input_grad
        if not (needs_grad_feat or needs_grad_weight):
            return None, None, None, None

        grad_output = grad_output.detach().contiguous()
        grad_feat = maybe_empty(
            **ctx.feat_meta_data,
            create_tensor=needs_grad_feat,
        )
        # captures both having edge_weight=None and
        # needs_grad_weight=False with edge_weight!=None
        grad_edge_weight = maybe_empty_like(edge_weight, needs_grad_weight)
        stream = torch.cuda.current_stream().cuda_stream

        if edge_weight is None:
            fg_n2n_bwd(
                grad_feat,
                grad_output,
                ctx.bwd_graph,
                ctx.aggr,
                out_pos,
                stream_id=stream,
            )

        else:
            fg_n2n_weighted_bwd(
                grad_feat,
                grad_edge_weight,
                grad_output,
                edge_weight,
                feat,
                ctx.bwd_graph,
                ctx.aggr,
                agg_output,
                out_pos,
                node_deg,
                stream_id=stream,
            )

        return grad_feat, grad_edge_weight, None, None


class _agg_concat_fg_e2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation using edge features in
    an edge-to-node reduction (e2n) on a full graph (fg) while concatenating the
    original features of output at the end (agg_concat)."""

    @staticmethod
    def forward(ctx, edge_feat, node_feat, graph, aggr):
        fwd_graph = graph._fwd_graph
        dim_node = node_feat.size(-1)
        dim_edge = edge_feat.size(-1)
        agg_output = torch.empty(
            fwd_graph.n_nodes,
            dim_edge + dim_node,
            dtype=edge_feat.dtype,
            device=edge_feat.device,
        )

        if aggr in (AggOp.Max, AggOp.Min):
            idx_type = fwd_graph.__class__.__name__.split("_")[-1]
            out_pos = torch.empty(
                fwd_graph.n_nodes,
                dim_edge,
                dtype=getattr(torch, idx_type),
                device=edge_feat.device,
            )
        else:
            out_pos = None

        stream = torch.cuda.current_stream().cuda_stream
        node_feat = node_feat.detach().contiguous()
        edge_feat = edge_feat.detach().contiguous()
        fg_e2n_fwd(
            agg_output,
            edge_feat,
            node_feat,
            fwd_graph,
            aggr,
            out_pos,
            stream_id=stream,
        )

        ctx.bwd_graph = graph._bwd_graph
        ctx.aggr = aggr
        ctx.dim_node = dim_node
        ctx.dim_edge = dim_edge

        ctx.edge_feat_meta_data = get_tensor_meta_data(edge_feat)
        ctx.node_feat_meta_data = get_tensor_meta_data(node_feat)

        ctx.save_for_backward(out_pos)

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        (out_pos,) = ctx.saved_tensors
        needs_grad_node, needs_grad_edge, _, _ = ctx.needs_input_grad
        if not (needs_grad_node or needs_grad_edge):
            return None, None, None, None

        grad_output = grad_output.detach().contiguous()
        grad_node_feat = maybe_empty(
            **ctx.node_feat_meta_data,
            create_tensor=needs_grad_node,
        )
        grad_edge_feat = maybe_empty(
            **ctx.edge_feat_meta_data,
            create_tensor=needs_grad_edge,
        )

        stream = torch.cuda.current_stream().cuda_stream
        fg_e2n_bwd(
            grad_edge_feat,
            grad_node_feat,
            grad_output,
            ctx.bwd_graph,
            ctx.dim_edge,
            ctx.dim_node,
            ctx.aggr,
            out_pos,
            stream_id=stream,
        )

        return grad_edge_feat, grad_node_feat, None, None


class _agg_concat_fg_n2n_e2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) and edge features in
    an edge-to-node reduction (e2n) on a full graph (fg) while concatenating the
    original features of output at the end (agg_concat)."""

    @staticmethod
    def forward(ctx, node_feat, edge_feat, graph, aggr):
        fwd_graph = graph._fwd_graph
        dim_node = node_feat.size(-1)
        dim_edge = edge_feat.size(-1)
        agg_output = torch.empty(
            fwd_graph.n_nodes,
            dim_node + dim_edge + dim_node,
            dtype=edge_feat.dtype,
            device=edge_feat.device,
        )

        if aggr in (AggOp.Max, AggOp.Min):
            idx_type = fwd_graph.__class__.__name__.split("_")[-1]
            out_pos = torch.empty(
                fwd_graph.n_nodes,
                dim_node + dim_edge,
                dtype=getattr(torch, idx_type),
                device=edge_feat.device,
            )
        else:
            out_pos = None

        stream = torch.cuda.current_stream().cuda_stream

        node_feat = node_feat.detach().contiguous()
        edge_feat = edge_feat.detach().contiguous()
        fg_n2n_e2n_fwd(
            agg_output,
            node_feat,
            edge_feat,
            fwd_graph,
            aggr,
            out_pos,
            stream_id=stream,
        )

        ctx.bwd_graph = graph._bwd_graph
        ctx.dim_node = dim_node
        ctx.dim_edge = dim_edge
        ctx.aggr = aggr
        ctx.node_feat_meta_data = get_tensor_meta_data(node_feat)
        ctx.edge_feat_meta_data = get_tensor_meta_data(edge_feat)

        ctx.save_for_backward(out_pos)

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        (out_pos,) = ctx.saved_tensors
        needs_grad_node, needs_grad_edge, _, _ = ctx.needs_input_grad
        if not (needs_grad_node or needs_grad_edge):
            return None, None, None, None

        grad_output = grad_output.detach().contiguous()
        grad_node_feat = maybe_empty(
            **ctx.node_feat_meta_data, create_tensor=needs_grad_node
        )
        grad_edge_feat = maybe_empty(
            **ctx.edge_feat_meta_data, create_tensor=needs_grad_edge
        )

        stream = torch.cuda.current_stream().cuda_stream

        fg_n2n_e2n_bwd(
            grad_node_feat,
            grad_edge_feat,
            grad_output,
            ctx.bwd_graph,
            ctx.dim_node,
            ctx.dim_edge,
            ctx.aggr,
            out_pos,
            stream_id=stream,
        )

        return grad_node_feat, grad_edge_feat, None, None


class _agg_concat_mfg_n2n_autograd(torch.autograd.Function):
    """
    Custom autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) on a message-flow-graph (mfg) while
    concatenating the original features of output at the end (agg_concat).
    """

    @staticmethod
    def forward(ctx, feat, graph, aggr):
        fwd_graph = graph._fwd_graph
        dim_in = feat.size(-1)
        agg_output = torch.empty(
            fwd_graph.n_out_nodes, 2 * dim_in, 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(
                fwd_graph.n_out_nodes,
                dim_in,
                dtype=getattr(torch, idx_type),
                device=feat.device,
            )
        else:
            out_pos = None

        stream = torch.cuda.current_stream().cuda_stream
        feat = feat.detach().contiguous()
        mfg_n2n_fwd(
            agg_output,
            feat,
            fwd_graph,
            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):
        needs_grad_node = ctx.needs_input_grad
        if not needs_grad_node:
            return None, None, None

        grad_output = grad_output.detach().contiguous()
        (out_pos,) = ctx.saved_tensors
        grad_feat = maybe_empty(**ctx.feat_meta_data, create_tensor=True)

        stream = torch.cuda.current_stream().cuda_stream

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

        return grad_feat, None, None
