# 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 torch.nn as nn
from functools import cmp_to_key

import pylibcugraphops

from pylibcugraphops.pytorch.utils import (
    maybe_detach,
    get_operator_from_module,
)

import pylibcugraphops.operators as ops


def _cmp_indices_tuples(a, b):
    if a[1][0] > b[1][0]:
        return 1
    elif a[1][0] == b[1][0]:
        if len(a[1]) < len(b[1]):
            return -1
        elif len(a[1]) == len(b[1]):
            return 0
        else:
            return 1
    else:
        return -1


def _matricization(path_vals, path_indices, num_out_segments, correlation):
    path_padded_vals = torch.tensor([], dtype=torch.float64)
    path_padded_indices = torch.tensor([], dtype=torch.int32)

    path_offsets = torch.zeros(2 * num_out_segments, dtype=torch.int32)

    sorted_vals_indices = sorted(
        zip(path_vals, path_indices), key=cmp_to_key(_cmp_indices_tuples)
    )

    curr_output = -1
    path_start, path_end, path_num = 0, 0, 0

    for path_val, path_index in sorted_vals_indices:

        if len(path_index) > correlation + 2:
            raise AssertionError("Invalid length of index tuple")

        if curr_output == -1:

            curr_output = path_index[0]
            path_num = 1

        elif curr_output == path_index[0]:

            path_num = path_num + 1

        elif curr_output != path_index[0]:

            path_end = path_start + path_num
            path_offsets[2 * curr_output + 0] = path_start
            path_offsets[2 * curr_output + 1] = path_end

            num_padded = ((path_num + 3) // 4) * 4 - path_num

            for i in range(0, num_padded):
                path_padded_vals = torch.cat(
                    (path_padded_vals, torch.tensor([0], dtype=torch.float64))
                )
                path_padded_indices = torch.cat(
                    (
                        path_padded_indices,
                        torch.tensor([0] * (correlation + 1), dtype=torch.int32),
                    )
                )

            path_num = 1
            path_start = path_end + num_padded
            path_end = path_start

            curr_output = path_index[0]

        else:
            raise AssertionError(
                "Unexpected value of curr_output and path index",
                curr_output,
                path_index,
            )

        path_padded_vals = torch.cat(
            (path_padded_vals, torch.tensor([path_val], dtype=torch.float64))
        )

        index = (
            path_index[1:-1]
            + [-1] * (correlation - len(path_index[1:-1]))
            + [path_index[-1]]
        )

        path_padded_indices = torch.cat(
            (path_padded_indices, torch.tensor(index, dtype=torch.int32))
        )

    if curr_output != -1:
        num_padded = ((path_num + 3) // 4) * 4 - path_num

        for i in range(0, num_padded):
            path_padded_vals = torch.cat(
                (path_padded_vals, torch.tensor([0], dtype=torch.float64))
            )
            path_padded_indices = torch.cat(
                (
                    path_padded_indices,
                    torch.tensor([0] * (correlation + 1), dtype=torch.int32),
                )
            )

        path_end = path_start + path_num
        path_offsets[2 * curr_output + 0] = path_start
        path_offsets[2 * curr_output + 1] = path_end

    return path_padded_vals, path_padded_indices, path_offsets


class SymmetricContractionFwdInfo(nn.Module):
    def __init__(self, path_vals, path_indices, num_out_segments, correlation, dtype):
        super().__init__()

        self.num_out_segments = num_out_segments

        path_vals, path_indices, path_offsets = _matricization(
            path_vals, path_indices, num_out_segments, correlation
        )

        path_vals = path_vals.to(dtype).contiguous()

        self.register_buffer("path_vals", path_vals)

        path_indices = path_indices.to(torch.int16).contiguous()
        self.register_buffer("path_indices", path_indices)

        path_offsets = path_offsets.to(torch.int32).flatten().contiguous()
        self.register_buffer("path_offsets", path_offsets)

        self.make_clebsch_gordan_tensor = get_operator_from_module(
            pylibcugraphops,
            "make_clebsch_gordan_tensor",
            (dtype,),
        )

    def forward(self):
        raise NotImplementedError(
            "SymmetricContractionFwdInfo module is not intended to be called."
        )

    def c_type(self):
        return self.make_clebsch_gordan_tensor(
            self.path_vals, self.path_indices, self.path_offsets, self.num_out_segments
        )


def _diff_index_tuple_list_tensor(path_vals, path_indices, num_in_segments):
    new_path_vals, new_path_indices = [], []

    for path_val, path_index in zip(path_vals, path_indices):

        for i_segment in range(0, num_in_segments):

            tmp = path_index[1:-1]
            num = tmp.count(i_segment)

            if num == 0:
                continue

            tmp.pop(tmp.index(i_segment))
            new_path_index = [i_segment] + [path_index[0]] + tmp + [path_index[-1]]
            new_path_val = num * path_val

            new_path_vals.append(new_path_val)
            new_path_indices.append(new_path_index)

    return new_path_vals, new_path_indices


def _diff_index_tuple_list_weight(path_vals, path_indices):

    new_path_indices = []
    for path_index in path_indices:
        index = [path_index[-1]] + path_index[1:-1] + [path_index[0]]
        new_path_indices.append(index)

    return path_vals, new_path_indices


class SymmetricContractionBwdInfo(nn.Module):
    def __init__(
        self,
        path_vals,
        path_indices,
        num_in_segments,
        num_coupling_paths,
        correlation,
        dtype=torch.float32,
    ):
        super().__init__()

        self.num_in_segments = num_in_segments
        self.num_coupling_paths = num_coupling_paths

        # for weight derivatives
        new_path_vals, new_path_indices = _diff_index_tuple_list_weight(
            path_vals, path_indices
        )

        path_padded_vals, path_padded_indices, path_offsets = _matricization(
            new_path_vals, new_path_indices, num_coupling_paths, correlation
        )

        path_vals_w = path_padded_vals.to(dtype).contiguous()
        self.register_buffer("path_vals_w", path_vals_w)

        path_indices_w = path_padded_indices.to(torch.int16).contiguous()
        self.register_buffer("path_indices_w", path_indices_w)

        path_offsets_w = path_offsets.to(torch.int32).contiguous()
        self.register_buffer("path_offsets_w", path_offsets_w)

        # for input tensor derivatives
        new_path_vals, new_path_indices = _diff_index_tuple_list_tensor(
            path_vals, path_indices, num_in_segments
        )

        path_padded_vals, path_padded_indices, path_offsets = _matricization(
            new_path_vals, new_path_indices, num_in_segments, correlation
        )

        path_vals = path_padded_vals.to(dtype).contiguous()
        self.register_buffer("path_vals", path_vals)

        path_indices = path_padded_indices.to(torch.int16).contiguous()
        self.register_buffer("path_indices", path_indices)

        path_offsets = path_offsets.to(torch.int32).flatten().contiguous()
        self.register_buffer("path_offsets", path_offsets)

        self.make_clebsch_gordan_tensor = get_operator_from_module(
            pylibcugraphops,
            "make_clebsch_gordan_tensor",
            (dtype,),
        )

    def forward(self):
        raise NotImplementedError(
            "SymmetricContractionBwdInfo module is not intended to be called."
        )

    def c_type(self):
        cg0 = self.make_clebsch_gordan_tensor(
            self.path_vals_w,
            self.path_indices_w,
            self.path_offsets_w,
            self.num_coupling_paths,
        )
        cg1 = self.make_clebsch_gordan_tensor(
            self.path_vals, self.path_indices, self.path_offsets, self.num_in_segments
        )
        return cg0, cg1


class _symmetric_tensor_contraction_autograd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, in1_tensor, in2_tensor, in2_tensor_id, args, kernels):

        (num_out_segments, correlation, cg_fwd, cg0_bwd, cg1_bwd) = args
        (fwd_kernel, bwd_kernel) = kernels

        out_tensor = torch.zeros(
            [in1_tensor.shape[0], num_out_segments, in1_tensor.shape[-1]],
            dtype=in1_tensor.dtype,
            device=in1_tensor.device,
        )

        stream = torch.cuda.current_stream().cuda_stream

        fwd_kernel(
            maybe_detach(out_tensor),
            maybe_detach(in1_tensor),
            maybe_detach(in2_tensor),
            maybe_detach(in2_tensor_id),
            cg_fwd,
            correlation,
            stream_id=stream,
        )

        ctx.save_for_backward(in1_tensor, in2_tensor, in2_tensor_id)

        ctx.args = args
        ctx.kernels = kernels

        return out_tensor

    @staticmethod
    def backward(ctx, grad_out_tensor):

        if grad_out_tensor is None:
            return None, None, None, None, None

        (in1_tensor, in2_tensor, in2_tensor_id) = ctx.saved_tensors

        (_, correlation, _, cg0_bwd, cg1_bwd) = ctx.args
        (_, bwd_kernel) = ctx.kernels

        stream = torch.cuda.current_stream().cuda_stream

        d_in1_tensor = torch.zeros(
            in1_tensor.shape, dtype=in1_tensor.dtype, device=in1_tensor.device
        )

        accum_type = torch.float32 if in1_tensor.element_size() <= 4 else torch.float64

        d_in2_tensor = torch.zeros(
            in2_tensor.shape, dtype=accum_type, device=in2_tensor.device
        )

        bwd_kernel(
            maybe_detach(d_in1_tensor),
            maybe_detach(d_in2_tensor),
            maybe_detach(grad_out_tensor),
            maybe_detach(in1_tensor),
            maybe_detach(in2_tensor),
            maybe_detach(in2_tensor_id),
            cg0_bwd,
            cg1_bwd,
            correlation,
            stream,
        )

        if d_in2_tensor.dtype != in2_tensor.dtype:
            d_in2_tensor = d_in2_tensor.to(in2_tensor.dtype)

        return d_in1_tensor, d_in2_tensor, None, None, None


def _map_to_path_indices(path_segment_indices):
    path_indices = [[index[-1]] + index[0:-1] for index in path_segment_indices]
    return path_indices


class SymmetricTensorContraction(nn.Module):
    """
    A class to perform the symmetric tensor contraction in the Mace model

    Attributes
    ---------------------------
    num_in_segments : int
        number of segments in the fist input tensor
    num_couplings : int
        number of segments in the second input tensor
    num_out_segments : int
        number of segments in the output tensor
    correlation: int
        correlation length
    dtype: torch.dtype
        data type for computation
    fwd_info: SymmetricContractionFwdInfo
        object to store precompute data for the forward
    bwd_info: SymmetricContractionBwdInfo
        object to store precompute data for the backward

    Methods
    ---------------------------
    forward(in1_tensor, in2_tensor, in2_tensor_id):
        forward function for the symmetric tensor contraction in the Mace model
    get_segments_dims()
        get the numbers of the segments for the input and the output tensors
    """

    def __init__(
        self,
        path_segment_indices: list[list[int]],
        path_coefficients: list[float],
        num_in_segments: int,
        num_couplings: int,
        num_out_segments: int,
        correlation: int,
        dtype: torch.dtype,
    ):
        """
        Construct all the necessary attributes

        Parameters
        ---------------------------
        path_segment_indices: list[list[int]]
            list of integer list to represent each path
        path_coefficients: list[float]
            list of floating number for each path
        num_in_segments: int
            number of segments in the first input tensor,
            the repeated operand
        num_couplings: int
            number of segments in the second input tensor
        num_out_segments: int
            number of segments in the output tensor
        correlation: int
            correlation length for the symmetric tensor contraction
        dtype: torch.dtype
            data type to perform computation
        """
        super().__init__()

        if not torch.cuda.is_available():
            raise AssertionError("No Nvidia GPU is detected")

        if len(path_segment_indices) != len(path_coefficients):
            raise AssertionError(
                "Number of the path coefficients and of the path segment indices \
                have are different"
            )

        self.num_in_segments = num_in_segments
        self.num_couplings = num_couplings
        self.num_out_segments = num_out_segments

        self.correlation = correlation
        self.dtype = dtype

        path_indices = _map_to_path_indices(path_segment_indices)

        self.fwd_info = SymmetricContractionFwdInfo(
            path_coefficients, path_indices, self.num_out_segments, correlation, dtype
        )
        self.bwd_info = SymmetricContractionBwdInfo(
            path_coefficients,
            path_indices,
            self.num_in_segments,
            self.num_couplings,
            correlation,
            dtype,
        )

    def forward(self, in1_tensor, in2_tensor, in2_tensor_id):
        """
        Forward funciton call

        Parameters
        ---------------------------
        in1_tensor: torch.Tensor
            The input tensor with shape [num_batches, num_in_segments, num_features]
        in2_tensor: torch.Tensor
            The input tensor with shape [num_batches, num_couplings, num_features]
        in2_tensor_id:
            The ID with shape [num_batches]
        """
        cg_fwd = self.fwd_info.c_type()
        cg0_bwd, cg1_bwd = self.bwd_info.c_type()
        args = (
            self.num_out_segments,
            self.correlation,
            cg_fwd,
            cg0_bwd,
            cg1_bwd,
        )

        fwd_call = get_operator_from_module(
            ops,
            "symmetric_tensor_contraction_n2n_fwd",
            (self.dtype,),
        )
        bwd_call = get_operator_from_module(
            ops,
            "symmetric_tensor_contraction_n2n_bwd",
            (self.dtype,),
        )

        kernels = (fwd_call, bwd_call)

        return _symmetric_tensor_contraction_autograd.apply(
            in1_tensor, in2_tensor, in2_tensor_id, args, kernels
        )

    def get_segments_dims(self):
        """
        Get the numbers of the segments for the input and the output tensors
        """
        return (
            self.num_in_segments,
            self.num_out_segments,
            self.num_coupling_paths,
        )
