# Copyright (c) 2023-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.operators as ops


class transpose_segment_to_contiguous_autograd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor: torch.Tensor, segment_info: torch.Tensor):
        segment_info = segment_info.detach().contiguous()
        ctx.save_for_backward(segment_info)
        tensor_transpose = torch.empty_like(tensor)
        stream = torch.cuda.current_stream().cuda_stream

        if tensor.numel() > 0:
            ops.segmented_transpose(
                tensor_transpose.detach().contiguous(),
                tensor.detach().contiguous(),
                segment_info,
                True,
                stream,
            )
        return tensor_transpose

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        (segment_info,) = ctx.saved_tensors
        grad_tensor = transpose_segment_to_strided_autograd.apply(
            grad_out, segment_info
        )
        return grad_tensor, None


class transpose_segment_to_strided_autograd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor: torch.Tensor, segment_info: torch.Tensor):
        segment_info = segment_info.detach().contiguous()
        ctx.save_for_backward(segment_info)
        tensor_transpose = torch.empty_like(tensor)
        stream = torch.cuda.current_stream().cuda_stream

        if tensor.numel() > 0:
            ops.segmented_transpose(
                tensor_transpose.detach().contiguous(),
                tensor.detach().contiguous(),
                segment_info,
                False,
                stream,
            )
        return tensor_transpose

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        (segment_info,) = ctx.saved_tensors
        grad_tensor = transpose_segment_to_contiguous_autograd.apply(
            grad_out, segment_info
        )
        return grad_tensor, None


def transpose_segment_to_contiguous(
    tensor: torch.Tensor, segment_info: torch.Tensor
) -> torch.Tensor:
    if segment_info is None:
        return tensor
    return transpose_segment_to_contiguous_autograd.apply(tensor, segment_info)


def transpose_segment_to_strided(
    tensor: torch.Tensor, segment_info: torch.Tensor
) -> torch.Tensor:
    if segment_info is None:
        return tensor
    return transpose_segment_to_strided_autograd.apply(tensor, segment_info)
