# 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 math
import torch
import torch.nn as nn
import pylibcugraphops

from dataclasses import dataclass
from typing import Tuple, Optional

from pylibcugraphops.pytorch.utils import (
    maybe_detach,
    maybe_empty_like,
    get_operator_from_module,
)
import pylibcugraphops.operators as ops


def sort_tp_info_for_operand(
    path_offsets_and_dims: torch.Tensor, path_cg_values: torch.Tensor, op_idx: int
):
    # idx = 0 -> in0
    # idx = 1 -> in1
    # idx = 2 -> in2
    # idx = 3 -> out
    key = path_offsets_and_dims[:, op_idx]
    perm = key.sort(descending=False, stable=True)[1]
    return path_offsets_and_dims[perm, :], path_cg_values[perm]


def path_offsets_to_csr(
    path_offsets_and_dims: torch.Tensor, op_idx: int, op_stride: int
):
    diff = torch.diff(
        path_offsets_and_dims[:, op_idx],
        n=1,
        append=path_offsets_and_dims.new_tensor([op_stride]),
    )
    path_csr_offsets = diff.nonzero().to(dtype=torch.int32).squeeze(dim=1)
    path_csr_offsets = torch.cat(
        [
            path_csr_offsets.new_tensor([0]),
            path_csr_offsets + 1,
        ]
    )
    return path_csr_offsets


def make_tp_info(
    op_idx: int,
    op_stride: int,
    path_offsets: torch.Tensor,
    path_cg_values: torch.Tensor,
):
    path_offsets = path_offsets.clone().detach()
    path_cg_values = path_cg_values.clone().detach()
    path_offsets, path_cg_values = sort_tp_info_for_operand(
        path_offsets, path_cg_values, op_idx
    )
    path_csr_offsets = path_offsets_to_csr(path_offsets, op_idx, op_stride)

    path_cg_values = path_cg_values
    path_offsets = path_offsets
    path_csr_offsets = path_csr_offsets

    return (
        path_csr_offsets,
        path_offsets,
        path_cg_values,
    )


def make_tp_infos(
    strides: list[int], path_offsets_and_dims: list[int], path_cg_values: list[float]
):
    path_cg_values = torch.FloatTensor(path_cg_values)
    path_offsets_and_dims = torch.IntTensor(path_offsets_and_dims)

    tp_infos = None

    if len(strides) == 4:
        tp_infos = (
            make_tp_info(3, strides[3], path_offsets_and_dims, path_cg_values),
            make_tp_info(0, strides[0], path_offsets_and_dims, path_cg_values),
            make_tp_info(1, strides[1], path_offsets_and_dims, path_cg_values),
            make_tp_info(2, strides[2], path_offsets_and_dims, path_cg_values),
        )

    elif len(strides) == 3:
        tp_infos = (
            make_tp_info(2, strides[2], path_offsets_and_dims, path_cg_values),
            make_tp_info(0, strides[0], path_offsets_and_dims, path_cg_values),
            make_tp_info(1, strides[1], path_offsets_and_dims, path_cg_values),
        )

    else:
        raise NotImplementedError

    return tp_infos


class TensorProductPathData(nn.Module):
    def __init__(
        self,
        path_csr_offsets: torch.Tensor,
        path_offsets: torch.Tensor,
        path_cg_values: torch.Tensor,
    ):
        super().__init__()
        # for now, don't have these buffers as part of the state dict to avoid
        # breaking changes w.r.t. loading state-dicts for now
        self.register_buffer("path_csr_offsets", path_csr_offsets, persistent=False)
        self.register_buffer("path_offsets", path_offsets, persistent=False)
        # Register float cg_values as int via bitcasting, so that model.to(dtype)
        # won't affect the precision of cg_values buffer
        self.register_buffer(
            "path_cg_values", path_cg_values.view(dtype=torch.int32), persistent=False
        )

    def forward(self, *args, **kwargs):
        raise NotImplementedError(
            "TensorProductPathData module is not intended to be called."
        )

    def c_type(self):
        return pylibcugraphops.make_tp_info(
            self.path_csr_offsets,
            self.path_offsets,
            self.path_cg_values.view(dtype=torch.float32),
        )


@dataclass
class TensorProductArgs:
    output_shape: Tuple[int, int]
    tp_info_fwd: pylibcugraphops.tp_info_fp32
    tp_info_bwd_dgrad_in0: pylibcugraphops.tp_info_fp32
    tp_info_bwd_dgrad_in1: pylibcugraphops.tp_info_fp32
    tp_info_bwd_dgrad_in2: pylibcugraphops.tp_info_fp32
    connection_mode: pylibcugraphops.ConnectionMode
    needs_input_grad_in0: bool
    needs_input_grad_in1: bool
    needs_input_grad_in2: bool
    needs_input_grad_bias: bool


class _fused_tensor_product_autograd(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        in0: torch.Tensor,
        in1: torch.Tensor,
        in2: Optional[torch.Tensor],
        bias: Optional[torch.Tensor],
        args: TensorProductArgs,
    ) -> torch.Tensor:

        out = torch.empty(args.output_shape, device=in0.device, dtype=in0.dtype)
        stream = torch.cuda.current_stream().cuda_stream

        fwd_fun = get_operator_from_module(
            ops,
            "fused_tp_n2n_fwd",
            (
                in0.dtype,
                in1.dtype,
                in2.dtype if in2 is not None else out.dtype,
                out.dtype,
            ),
        )

        fwd_fun(
            out,
            maybe_detach(in0),
            maybe_detach(in1),
            maybe_detach(in2),
            maybe_detach(bias),
            args.connection_mode,
            args.tp_info_fwd,
            stream_id=stream,
        )

        ctx.save_for_backward(
            in0,
            in1,
            in2,
            bias,
        )
        ctx.args = args

        return out

    @staticmethod
    def backward(
        ctx, grad_out: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]:

        (in0, in1, in2, bias) = ctx.saved_tensors

        return _fused_tensor_product_backward_autograd.apply(
            grad_out,
            in0,
            in1,
            in2,
            bias,
            ctx.args,
        )


class _fused_tensor_product_backward_autograd(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        grad_out: torch.Tensor,
        in0: torch.Tensor,
        in1: torch.Tensor,
        in2: Optional[torch.Tensor],
        bias: Optional[torch.Tensor],
        args: TensorProductArgs,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, None, None]:

        # calling function takes care of ensuring contiguous and detach
        grad_in0 = maybe_empty_like(in0, args.needs_input_grad_in0)
        grad_in1 = maybe_empty_like(in1, args.needs_input_grad_in1)
        grad_in2 = maybe_empty_like(in2, args.needs_input_grad_in2)
        grad_bias = maybe_empty_like(bias, args.needs_input_grad_bias)

        stream = torch.cuda.current_stream().cuda_stream

        bwd_fun = get_operator_from_module(
            ops,
            "fused_tp_n2n_bwd",
            (
                in0.dtype,
                in1.dtype,
                in2.dtype if in2 is not None else grad_out.dtype,
                grad_out.dtype,
            ),
        )

        in0 = maybe_detach(in0)
        in1 = maybe_detach(in1)
        in2 = maybe_detach(in2)
        bias = maybe_detach(bias)
        grad_out = maybe_detach(grad_out)

        bwd_fun(
            grad_in0,
            grad_in1,
            grad_in2,
            grad_bias,
            grad_out,
            in0,
            in1,
            in2,
            bias,
            args.connection_mode,
            args.tp_info_bwd_dgrad_in0,
            args.tp_info_bwd_dgrad_in1,
            args.tp_info_bwd_dgrad_in2,
            stream_id=stream,
        )

        ctx.args = args
        ctx.save_for_backward(
            in0,
            in1,
            in2,
            grad_out,
        )

        return grad_in0, grad_in1, grad_in2, grad_bias, None

    @staticmethod
    def backward(
        ctx, grad_grad_in0, grad_grad_in1, grad_grad_in2, grad_grad_bias, grad_args
    ):
        # args should not require_grad in any case
        # assert not ctx.needs_input_grad[4]
        in0, in1, in2, grad_out = ctx.saved_tensors

        grad_grad_out = maybe_empty_like(grad_out, ctx.needs_input_grad[0])
        grad_in0 = maybe_empty_like(in0, ctx.needs_input_grad[1])
        grad_in1 = maybe_empty_like(in1, ctx.needs_input_grad[2])
        grad_in2 = maybe_empty_like(in2, ctx.needs_input_grad[3])

        stream = torch.cuda.current_stream().cuda_stream

        bwd_bwd_fun = get_operator_from_module(
            ops,
            "fused_tp_n2n_bwd_bwd",
            (
                in0.dtype,
                in1.dtype,
                in2.dtype if in2 is not None else grad_out.dtype,
                grad_out.dtype,
            ),
        )

        bwd_bwd_fun(
            grad_in0,
            grad_in1,
            grad_in2,
            grad_grad_out,
            maybe_detach(grad_grad_in0),
            maybe_detach(grad_grad_in1),
            maybe_detach(grad_grad_in2),
            grad_out,
            in0,
            in1,
            in2,
            ctx.args.connection_mode,
            ctx.args.tp_info_fwd,
            ctx.args.tp_info_bwd_dgrad_in0,
            ctx.args.tp_info_bwd_dgrad_in1,
            ctx.args.tp_info_bwd_dgrad_in2,
            stream_id=stream,
        )

        if grad_grad_bias is not None:
            grad_grad_out += grad_grad_bias

        return grad_grad_out, grad_in0, grad_in1, grad_in2, None, None


class FusedTensorProductOp4(nn.Module):
    def __init__(
        self,
        operand_segment_modes: list[str],
        operand_segment_offsets: list[list[int]],
        operand_segment_shapes: list[list[list[int]]],
        path_indices: list[list[int]],
        path_coefficients: list[float],
    ):
        super().__init__()

        mode = "_".join(operand_segment_modes)

        mode_mappings = {
            "u__u_u": "UVU",
            "u__uw_w": "U1W",
            "u__uv_v": "U1W",
            "_v_vw_w": "1VW",
            "u_v_uv_u": "UVU",
            "u_v_uv_v": "UVV",
            "u_u_uw_w": "UUW",
            "u_v_uvw_w": "UVW",
        }

        # for know, we can always extract
        # u,v,w from individual operands
        dims_dict = {
            "u__u_u": (0, 1, -1),
            "u__uw_w": (0, 1, 3),
            "u__uv_v": (0, 1, 3),
            "_v_vw_w": (0, 1, 3),
            "u_v_uv_u": (0, 1, -1),
            "u_v_uv_v": (0, 1, -1),
            "u_u_uw_w": (0, -1, 3),
            "u_v_uvw_w": (0, 1, 3),
        }

        assert (
            mode in mode_mappings
        ), f"{mode} not found, expected {list(mode_mappings.keys())}"

        u_op_idx, v_op_idx, w_op_idx = dims_dict[mode]

        if (
            len(operand_segment_modes) != 4
            or len(operand_segment_offsets) != 4
            or len(operand_segment_shapes) != 4
        ):
            raise ValueError("Expected 4 operands (3 inputs, 1 output)")

        if len(path_indices) != len(path_coefficients):
            msg = (
                "path_coefficients and path_indices both have to be of length num_paths"
            )
            msg += f", got {len(path_coefficients)} & {len(path_indices)} respectively."

            raise ValueError(msg)

        operand_segment_strides = [
            [math.prod(segment) for segment in op] for op in operand_segment_shapes
        ]

        strides = [sum(op) for op in operand_segment_strides]

        path_offsets_and_dims = []
        self.connection_mode = mode_mappings[mode]

        for indices in path_indices:
            u = (
                operand_segment_strides[u_op_idx][indices[u_op_idx]]
                if u_op_idx != -1
                else -1
            )
            v = (
                operand_segment_strides[v_op_idx][indices[v_op_idx]]
                if v_op_idx != -1
                else -1
            )
            w = (
                operand_segment_strides[w_op_idx][indices[w_op_idx]]
                if w_op_idx != -1
                else -1
            )

            path_offsets_and_dims.append(
                (
                    operand_segment_offsets[0][indices[0]],
                    operand_segment_offsets[1][indices[1]],
                    operand_segment_offsets[2][indices[2]],
                    operand_segment_offsets[3][indices[3]],
                    u,
                    v,
                    w,
                    -1,
                )
            )

        tp_infos = make_tp_infos(strides, path_offsets_and_dims, path_coefficients)

        self.stride_out = strides[-1]
        self._tp_info_fwd = TensorProductPathData(*tp_infos[0])
        self._tp_info_bwd_dgrad_in0 = TensorProductPathData(*tp_infos[1])
        self._tp_info_bwd_dgrad_in1 = TensorProductPathData(*tp_infos[2])
        self._tp_info_bwd_dgrad_in2 = TensorProductPathData(*tp_infos[3])

    def forward(
        self,
        in0: torch.Tensor,
        in1: torch.Tensor,
        in2: torch.Tensor,
    ) -> torch.Tensor:
        if in0.dim() == 1 and in1.dim() == 1 and in2.dim() != 1:
            batch_size = in2.size(0)
        elif in0.dim() == 1 and in1.dim() != 1:
            batch_size = in1.size(0)
        else:
            batch_size = in0.size(0)

        args = TensorProductArgs(
            (batch_size, self.stride_out),
            self._tp_info_fwd.c_type(),
            self._tp_info_bwd_dgrad_in0.c_type(),
            self._tp_info_bwd_dgrad_in1.c_type(),
            self._tp_info_bwd_dgrad_in2.c_type(),
            getattr(pylibcugraphops.ConnectionMode, self.connection_mode),
            in0.requires_grad,
            in1.requires_grad,
            in2.requires_grad,
            False,
        )

        out = _fused_tensor_product_autograd.apply(in0, in1, in2, None, args)

        return out


class FusedTensorProductOp3(nn.Module):
    def __init__(
        self,
        operand_segment_modes: list[str],
        operand_segment_offsets: list[list[int]],
        operand_segment_shapes: list[list[list[int]]],
        path_indices: list[list[int]],
        path_coefficients: list[float],
    ):
        super().__init__()

        mode = "_".join(operand_segment_modes)

        mode_mappings = {
            "u_u_u": "UUU",
            "u_v_uv": "UVUV",
            "u_uv_v": "UUVV",
            "u__u": "U1U1",
            "_v_v": "1V1V",
        }

        # for know, we can always extract
        # u,v,w from individual operands
        dims_dict = {
            "u_u_u": (0, -1),
            "u_v_uv": (0, 1),
            "u_uv_v": (0, 2),
            "u__u": (0, 1),
            "_v_v": (0, 1),
        }

        assert mode in mode_mappings

        u_op_idx, v_op_idx = dims_dict[mode]

        if (
            len(operand_segment_modes) != 3
            or len(operand_segment_offsets) != 3
            or len(operand_segment_shapes) != 3
        ):
            raise ValueError("Expected 3 operands (3 inputs, 1 output)")

        if len(path_indices) != len(path_coefficients):
            msg = ""
            msg += (
                "path_coefficients and path_indices both have to be of length num_paths"
            )
            msg += f", got {len(path_coefficients)} & {len(path_indices)} respectively."
            raise ValueError(msg)

        operand_segment_strides = [
            [math.prod(segment) for segment in op] for op in operand_segment_shapes
        ]
        strides = [sum(op) for op in operand_segment_strides]

        path_offsets_and_dims = []
        self.connection_mode = mode_mappings[mode]

        for indices in path_indices:
            u = (
                operand_segment_strides[u_op_idx][indices[u_op_idx]]
                if u_op_idx != -1
                else -1
            )
            v = (
                operand_segment_strides[v_op_idx][indices[v_op_idx]]
                if v_op_idx != -1
                else -1
            )
            path_offsets_and_dims.append(
                (
                    operand_segment_offsets[0][indices[0]],
                    operand_segment_offsets[1][indices[1]],
                    -1,
                    operand_segment_offsets[2][indices[2]],
                    u,
                    v,
                    -1,
                    -1,
                )
            )

        tp_infos = make_tp_infos(strides, path_offsets_and_dims, path_coefficients)

        self.stride_out = strides[-1]
        self._tp_info_fwd = TensorProductPathData(*tp_infos[0])
        self._tp_info_bwd_dgrad_in0 = TensorProductPathData(*tp_infos[1])
        self._tp_info_bwd_dgrad_in1 = TensorProductPathData(*tp_infos[2])

    def forward(
        self,
        in0: torch.Tensor,
        in1: torch.Tensor,
    ) -> torch.Tensor:
        if in0.dim() == 1 and in1.dim() != 1:
            batch_size = in1.size(0)
        else:
            batch_size = in0.size(0)

        args = TensorProductArgs(
            (batch_size, self.stride_out),
            self._tp_info_fwd.c_type(),
            self._tp_info_bwd_dgrad_in0.c_type(),
            self._tp_info_bwd_dgrad_in1.c_type(),
            self._tp_info_bwd_dgrad_in1.c_type(),
            getattr(pylibcugraphops.ConnectionMode, self.connection_mode),
            in0.requires_grad,
            in1.requires_grad,
            False,
            False,
        )

        out = _fused_tensor_product_autograd.apply(in0, in1, None, None, args)

        return out
