# Copyright (c) 2022-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
import pylibcugraphops
import warnings

from typing import Optional, Union, Tuple


class CSC:
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        num_src_nodes: int,
        map_csc_to_coo: Optional[torch.Tensor] = None,
        map_dst_to_src: Optional[torch.Tensor] = None,
        dst_max_in_degree: int = -1,
        reverse_graph_bwd: bool = False,
        is_bipartite: bool = False,
    ):
        self.offsets = offsets
        self.indices = indices
        self.num_src_nodes = num_src_nodes
        self.num_dst_nodes = offsets.size(0) - 1
        self.num_indices = indices.size(0)

        self.rev_offsets = (
            torch.empty(
                num_src_nodes + 1,
                device=offsets.device,
                dtype=offsets.dtype,
            )
            if reverse_graph_bwd
            else None
        )
        self.rev_indices = torch.empty_like(indices) if reverse_graph_bwd else None

        self.map_dst_to_src = map_dst_to_src
        self.map_csc_to_coo = map_csc_to_coo
        self.map_rev_to_coo = torch.empty_like(indices) if reverse_graph_bwd else None

        self.dst_max_in_degree = dst_max_in_degree
        # TODO remove warning once we again have support for these
        if dst_max_in_degree > 0:
            warnings.warn("dst_max_in_degree currently has no effect")
        if map_dst_to_src is not None:
            raise NotImplementedError("map_dst_to_src is currently not supported")

        self.is_bipartite = is_bipartite

        if reverse_graph_bwd:
            self.add_reverse_graph()

        else:
            if self.is_bipartite:
                self.graph_csc = pylibcugraphops.make_bipartite_csc(
                    offsets,
                    indices,
                    num_src_nodes,
                    map_csc_to_coo,
                )

            else:
                self.graph_csc = pylibcugraphops.make_csc(
                    offsets,
                    indices,
                    num_src_nodes,
                    map_csc_to_coo,
                )

        self.reverse_graph_bwd = reverse_graph_bwd

    def __to__(self, *args, **kwargs):
        self.offsets = self.offsets.to(*args, **kwargs)
        self.indices = self.indices.to(*args, **kwargs)
        if self.map_dst_to_src is not None:
            self.map_dst_to_src = self.map_dst_to_src.to(*args, **kwargs)
        if self.map_csc_to_coo is not None:
            self.map_csc_to_coo = self.map_csc_to_coo.to(*args, **kwargs)
        if self.map_rev_to_coo is not None:
            self.map_rev_to_coo = self.map_rev_to_coo.to(*args, **kwargs)
        if self.reverse_graph_bwd:
            self.rev_offsets = self.rev_offsets.to(*args, **kwargs)
            self.rev_indices = self.rev_indices.to(*args, **kwargs)

    def add_reverse_graph(self):
        if self.reverse_graph_bwd:
            assert self.rev_offsets is not None
            assert self.rev_indices is not None
            assert self.map_rev_to_coo is not None
            return

        self.rev_offsets = torch.empty(
            self.num_src_nodes + 1,
            device=self.offsets.device,
            dtype=self.offsets.dtype,
        )

        self.rev_indices = torch.empty_like(self.indices)
        self.map_rev_to_coo = torch.empty_like(self.indices)

        if self.is_bipartite:
            self.graph_csc = pylibcugraphops.make_bipartite_csc(
                self.offsets,
                self.indices,
                self.num_src_nodes,
                self.map_csc_to_coo,
                self.rev_offsets,
                self.rev_indices,
                self.map_rev_to_coo,
            )

        else:
            self.graph_csc = pylibcugraphops.make_csc(
                self.offsets,
                self.indices,
                self.num_src_nodes,
                self.map_csc_to_coo,
                self.rev_offsets,
                self.rev_indices,
                self.map_rev_to_coo,
            )

        stream = torch.cuda.current_stream().cuda_stream
        # call reverse to get necessary workspace size and
        # create necessary workspace buffers
        node_counts = torch.empty(
            self.num_src_nodes,
            dtype=self.offsets.dtype,
            device=self.offsets.device,
        )
        cub_workspace = None
        cub_workspace_size = pylibcugraphops.reverse_graph(
            self.graph_csc, node_counts, cub_workspace, 0, stream
        )
        cub_workspace = torch.empty(
            cub_workspace_size, dtype=torch.int8, device=self.offsets.device
        )
        # actually reverse graph (in-place)
        pylibcugraphops.reverse_graph(
            self.graph_csc, node_counts, cub_workspace, cub_workspace_size, stream
        )


class HeteroCSC(CSC):
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        edge_types: torch.Tensor,
        num_src_nodes: int,
        num_edge_types: int,
        num_node_types: int = 0,
        node_types: Optional[Union[torch.Tensor, Tuple[torch.Tensor]]] = None,
        map_dst_to_src: Optional[torch.Tensor] = None,
        map_csc_to_coo: Optional[torch.Tensor] = None,
        dst_max_in_degree: int = -1,
        reverse_graph_bwd: bool = False,
        is_bipartite: bool = False,
    ):
        if is_bipartite:
            if num_node_types > 0:
                assert isinstance(node_types, Tuple), (
                    "Expected tuple of node types for a bipartite graph"
                    f", but got {type(node_types)} instead."
                )

        else:
            if num_node_types > 0:
                assert isinstance(node_types, torch.Tensor), (
                    "Expected node_types to be of type torch.Tensor"
                    f", but got {type(node_types)} instead."
                )

        super(HeteroCSC, self).__init__(
            offsets,
            indices,
            num_src_nodes,
            map_dst_to_src=map_dst_to_src,
            map_csc_to_coo=map_csc_to_coo,
            dst_max_in_degree=dst_max_in_degree,
            reverse_graph_bwd=reverse_graph_bwd,
            is_bipartite=is_bipartite,
        )
        self.edge_types = edge_types
        self.num_edge_types = num_edge_types

        if self.is_bipartite:
            src_node_types, dst_node_types = node_types
            self.graph_csc = pylibcugraphops.make_bipartite_csc_hg(
                offsets,
                indices,
                num_src_nodes,
                num_node_types,
                num_edge_types,
                src_node_types,
                dst_node_types,
                edge_types,
                map_csc_to_coo,
                self.rev_offsets,
                self.rev_indices,
                self.map_rev_to_coo,
            )

        else:
            self.graph_csc = pylibcugraphops.make_csc_hg(
                offsets,
                indices,
                num_src_nodes,
                num_node_types,
                num_edge_types,
                node_types,
                edge_types,
                map_csc_to_coo,
                self.rev_offsets,
                self.rev_indices,
                self.map_rev_to_coo,
            )

    def __to__(self, *args, **kwargs):
        super(HeteroCSC, self).__to__(*args, **kwargs)
        self.edge_types = self.edge_types.to(*args, **kwargs)


#################################################################################
# Deprecated Graph Types
# for now, they inherit from CSC which mostly makes all previous types
# (in 23.04) compatible with the refactored operators. Deprecation in 23.08,
# potential deletion in 23.10.
#################################################################################


class BipartiteCSC(CSC):
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        num_src_nodes: int,
        ef_indices: Optional[torch.Tensor] = None,
        reverse_graph_bwd: bool = False,
    ):
        super(BipartiteCSC, self).__init__(
            offsets,
            indices,
            num_src_nodes=num_src_nodes,
            map_csc_to_coo=ef_indices,
            reverse_graph_bwd=reverse_graph_bwd,
            is_bipartite=True,
        )

        warnings.warn(
            "BipartiteCSC is deprecated with the 23.08 release"
            " and will be removed in 23.10, use CSC instead.",
            DeprecationWarning,
        )


class SampledCSC(CSC):
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        max_num_neighbors: int,
        num_src_nodes: int,
    ):
        super(SampledCSC, self).__init__(
            offsets,
            indices,
            num_src_nodes=num_src_nodes,
            dst_max_in_degree=max_num_neighbors,
            is_bipartite=False,
        )

        warnings.warn(
            "SampledCSC is deprecated with the 23.08 release"
            " and will be removed in 23.10, use CSC instead.",
            DeprecationWarning,
        )


class StaticCSC(CSC):
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        ef_indices: Optional[torch.Tensor] = None,
        reverse_graph_bwd: bool = False,
    ):
        num_src_nodes = offsets.size(0) - 1
        super(StaticCSC, self).__init__(
            offsets,
            indices,
            num_src_nodes=num_src_nodes,
            map_csc_to_coo=ef_indices,
            reverse_graph_bwd=reverse_graph_bwd,
            is_bipartite=False,
        )

        warnings.warn(
            "StaticCSC is deprecated with the 23.08 release"
            " and will be removed in 23.10, use CSC instead.",
            DeprecationWarning,
        )


class SampledHeteroCSC(HeteroCSC):
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        edge_types: torch.Tensor,
        max_num_neighbors: int,
        num_src_nodes: int,
        num_edge_types: int,
    ):
        super(SampledHeteroCSC, self).__init__(
            offsets,
            indices,
            edge_types=edge_types,
            num_src_nodes=num_src_nodes,
            num_edge_types=num_edge_types,
            dst_max_in_degree=max_num_neighbors,
            is_bipartite=False,
        )

        warnings.warn(
            "SampledHeteroCSC is deprecated with the 23.08 release"
            " and will be removed in 23.10, use HeteroCSC instead.",
            DeprecationWarning,
        )


class StaticHeteroCSC(HeteroCSC):
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        edge_types: torch.Tensor,
        num_edge_types: int,
        ef_indices: Optional[torch.Tensor] = None,
    ):
        num_src_nodes = offsets.size(0) - 1

        super(StaticHeteroCSC, self).__init__(
            offsets,
            indices,
            edge_types=edge_types,
            num_src_nodes=num_src_nodes,
            num_edge_types=num_edge_types,
            map_csc_to_coo=ef_indices,
            is_bipartite=False,
        )

        warnings.warn(
            "StaticHeteroCSC is deprecated with the 23.08 release"
            " and will be removed in 23.10, use HeteroCSC instead.",
            DeprecationWarning,
        )
