# 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
import e3nn.o3 as o3
import pylibcugraphops

from e3nn.o3 import Irreps
from typing import Tuple, Optional, Iterator

from pylibcugraphops.pytorch.utils import maybe_detach, maybe_empty_like
from pylibcugraphops.operators import (
    fused_tp_n2n_fwd,
    fused_tp_n2n_bwd,
    fused_tp_n2n_bwd_bwd,
)


def _transpose_irrep_to_m_last(tensor: torch.Tensor, irreps: Irreps):
    new_tensor = torch.empty_like(tensor)
    offset = 0
    for irr in irreps:
        num_channels_per_m = irr.mul
        m = irr.ir.dim
        tmp = tensor[:, offset : offset + irr.dim].clone()
        if m > 1:
            tmp = tmp.reshape(-1, m, num_channels_per_m)
            tmp = torch.transpose(tmp, 2, 1).reshape(-1, irr.dim)
        new_tensor[:, offset : offset + irr.dim] = tmp
        offset += irr.dim

    return new_tensor


def _transpose_irrep_to_channels_last(tensor: torch.Tensor, irreps: Irreps):
    new_tensor = torch.empty_like(tensor)
    offset = 0
    for irr in irreps:
        num_channels_per_m = irr.mul
        m = irr.ir.dim
        tmp = tensor[:, offset : offset + irr.dim].clone()
        if m > 1:
            tmp = tmp.reshape(-1, num_channels_per_m, m)
            tmp = torch.transpose(tmp, 2, 1).reshape(-1, irr.dim)
        new_tensor[:, offset : offset + irr.dim] = tmp
        offset += irr.dim

    return new_tensor


class _transpose_irrep_to_channels_last_autograd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor: torch.Tensor, irreps: Irreps):
        ctx.irreps = irreps
        return _transpose_irrep_to_channels_last(tensor, irreps)

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        grad_tensor = _transpose_irrep_to_m_last_autograd.apply(grad_out, ctx.irreps)
        return grad_tensor, None


class _transpose_irrep_to_m_last_autograd(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor: torch.Tensor, irreps: Irreps):
        ctx.irreps = irreps
        return _transpose_irrep_to_m_last(tensor, irreps)

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        grad_tensor = _transpose_irrep_to_channels_last_autograd.apply(
            grad_out, ctx.irreps
        )
        return grad_tensor, None


def transpose_irrep_to_channels_last(
    tensor: torch.Tensor, irreps: Irreps
) -> torch.Tensor:
    return _transpose_irrep_to_channels_last_autograd.apply(tensor, irreps)


def transpose_irrep_to_m_last(tensor: torch.Tensor, irreps: Irreps) -> torch.Tensor:
    return _transpose_irrep_to_m_last_autograd.apply(tensor, irreps)


def sort_tp_info_for_operand(
    path_offsets_and_dims: torch.Tensor, path_cg_values: torch.Tensor, op_idx: int
):
    # idx = 0 -> in1
    # idx = 1 -> in2
    # idx = 2 -> weight
    # 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.to(device="cuda")
    path_offsets = path_offsets.to(device="cuda")
    path_csr_offsets = path_csr_offsets.to(device="cuda")

    return (
        path_csr_offsets,
        path_offsets,
        path_cg_values,
        pylibcugraphops.make_tp_info(path_csr_offsets, path_offsets, path_cg_values),
    )


def get_tp_infos_from_e3nn_instructions(
    instructions: Iterator[o3.Instruction],
    irreps_in1: o3.Irreps,
    irreps_in2: o3.Irreps,
    irreps_out: o3.Irreps,
):
    connection_mode = instructions[0].connection_mode

    path_offsets_and_dims = []
    path_cg_values = []

    slices_a = irreps_in1.slices()
    slices_b = irreps_in2.slices()
    slices_o = irreps_out.slices()

    off_weight = 0

    for inst in instructions:
        error_msg = "Only one connection mode at a time supported"
        error_msg += f"got {connection_mode} and {inst.connection_mode} respectively"
        if not inst.connection_mode == connection_mode:
            raise ValueError(error_msg)
        ir_idx_a = inst.i_in1
        ir_idx_b = inst.i_in2
        ir_idx_o = inst.i_out

        off_a = slices_a[ir_idx_a].start
        off_b = slices_b[ir_idx_b].start
        off_o = slices_o[ir_idx_o].start
        path_weight = inst.path_weight

        dim_u = inst.path_shape[0]
        dim_v = inst.path_shape[1] if len(inst.path_shape) > 1 else 0
        dim_w = inst.path_shape[2] if len(inst.path_shape) > 2 else 0

        if connection_mode == "uvw":
            dim_weight = dim_u * dim_v * dim_w
        elif connection_mode in {"uvu", "uvv"}:
            dim_weight = dim_u * dim_v
        elif connection_mode == "uuw":
            dim_w = dim_v
            dim_v = dim_u
            dim_weight = dim_u * dim_w
        else:
            dim_weight = 0

        # dim_a, dim_b, dim_o here actually are the number of channels
        dim_a = irreps_in1[ir_idx_a].mul
        dim_b = irreps_in2[ir_idx_b].mul
        dim_o = irreps_out[ir_idx_o].mul

        l_a = irreps_in1[ir_idx_a].ir.l
        l_b = irreps_in2[ir_idx_b].ir.l
        l_o = irreps_out[ir_idx_o].ir.l

        cg = o3.wigner_3j(l_a, l_b, l_o)

        for m_a in range(2 * l_a + 1):
            for m_b in range(2 * l_b + 1):
                for m_o in range(2 * l_o + 1):
                    cg_val = cg[m_a, m_b, m_o]
                    if cg_val != 0:
                        path_cg_values.append(cg_val * path_weight)
                        path_offsets_and_dims.append(
                            (
                                off_a + m_a * dim_a,
                                off_b + m_b * dim_b,
                                off_weight,
                                off_o + m_o * dim_o,
                                dim_u,
                                dim_v,
                                dim_w,
                                -1,  # place-holder / padding
                            )
                        )

        if inst.has_weight:
            off_weight += dim_weight

    strides = (
        irreps_in1.dim,
        irreps_in2.dim,
        off_weight,
        irreps_out.dim,
    )

    if connection_mode == "uvw" and all([ir.mul == 1 for ir in irreps_in1]):
        connection_mode = "1vw"
        pass
    elif connection_mode == "uvw" and all([ir.mul == 1 for ir in irreps_in2]):
        connection_mode = "u1w"
        pass
    connection_mode = getattr(pylibcugraphops.ConnectionMode, connection_mode.upper())

    path_cg_values = torch.FloatTensor(path_cg_values)
    path_offsets_and_dims = torch.IntTensor(path_offsets_and_dims)

    tp_info_fwd = make_tp_info(3, strides[3], path_offsets_and_dims, path_cg_values)
    tp_info_bwd_dgrad_in1 = make_tp_info(
        0, strides[0], path_offsets_and_dims, path_cg_values
    )
    tp_info_bwd_dgrad_in2 = make_tp_info(
        1, strides[1], path_offsets_and_dims, path_cg_values
    )
    tp_info_bwd_dgrad_w = make_tp_info(
        2, strides[2], path_offsets_and_dims, path_cg_values
    )

    return (
        tp_info_fwd,
        tp_info_bwd_dgrad_in1,
        tp_info_bwd_dgrad_in2,
        tp_info_bwd_dgrad_w,
        connection_mode,
        strides,
    )


class _fused_tensor_product_autograd(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        in1: torch.Tensor,
        in2: torch.Tensor,
        weight: Optional[torch.Tensor],
        args: Tuple,
    ) -> torch.Tensor:
        (
            tp_info_fwd,
            tp_info_bwd_dgrad_in1,
            tp_info_bwd_dgrad_in2,
            tp_info_bwd_dgrad_w,
            connection_mode,
            stride_out,
        ) = args

        out = torch.empty((in1.size(0), stride_out), device=in1.device, dtype=in1.dtype)
        stream = torch.cuda.current_stream().cuda_stream

        fused_tp_n2n_fwd(
            maybe_detach(out),
            maybe_detach(in1),
            maybe_detach(in2),
            maybe_detach(weight),
            connection_mode,
            tp_info_fwd,
            stream_id=stream,
        )

        ctx.save_for_backward(
            in1,
            in2,
            weight,
        )
        ctx.args = args

        return out

    @staticmethod
    def backward(
        ctx, grad_out: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]:
        # args should not require_grad in any case
        # assert not ctx.needs_input_grad[3]

        (in1, in2, weight) = ctx.saved_tensors

        return _fused_tensor_product_backward_autograd.apply(
            grad_out,
            in1,
            in2,
            weight,
            (*ctx.args, *ctx.needs_input_grad),
        )


class _fused_tensor_product_backward_autograd(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        grad_out: torch.Tensor,
        in1: torch.Tensor,
        in2: torch.Tensor,
        weight: Optional[torch.Tensor],
        args,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, None, None]:
        (
            tp_info_fwd,
            tp_info_bwd_dgrad_in1,
            tp_info_bwd_dgrad_in2,
            tp_info_bwd_dgrad_w,
            connection_mode,
            stride_out,
            needs_input_grad_in1,
            needs_input_grad_in2,
            needs_input_grad_w,
            needs_input_grad_args,
        ) = args

        # calling function takes care of ensuring contiguous and detach
        grad_in1 = maybe_empty_like(in1, needs_input_grad_in1)
        grad_in2 = maybe_empty_like(in2, needs_input_grad_in2)
        grad_weight = maybe_empty_like(weight, needs_input_grad_w)

        stream = torch.cuda.current_stream().cuda_stream

        fused_tp_n2n_bwd(
            grad_in1,
            grad_in2,
            grad_weight,
            maybe_detach(grad_out),
            maybe_detach(in1),
            maybe_detach(in2),
            maybe_detach(weight),
            connection_mode,
            tp_info_bwd_dgrad_in1,
            tp_info_bwd_dgrad_in2,
            tp_info_bwd_dgrad_w,
            stream_id=stream,
        )

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

        return grad_in1, grad_in2, grad_weight, None

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

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

        stream = torch.cuda.current_stream().cuda_stream

        tp_info_fwd = ctx.args[0]
        tp_info_bwd_dgrad_in1 = ctx.args[1]
        tp_info_bwd_dgrad_in2 = ctx.args[2]
        tp_info_bwd_dgrad_w = ctx.args[3]
        connection_mode = ctx.args[4]

        fused_tp_n2n_bwd_bwd(
            grad_in1,
            grad_in2,
            grad_weight,
            grad_grad_out,
            maybe_detach(grad_grad_in1),
            maybe_detach(grad_grad_in2),
            maybe_detach(grad_grad_weight),
            maybe_detach(grad_out),
            maybe_detach(in1),
            maybe_detach(in2),
            maybe_detach(weight),
            connection_mode,
            tp_info_fwd,
            tp_info_bwd_dgrad_in1,
            tp_info_bwd_dgrad_in2,
            tp_info_bwd_dgrad_w,
            stream_id=stream,
        )

        return grad_grad_out, grad_in1, grad_in2, grad_weight, None


class FusedTensorProduct(nn.Module):
    def __init__(
        self,
        instructions: Iterator[o3.Instruction],
        irreps_in1: o3.Irreps,
        irreps_in2: o3.Irreps,
        irreps_out: o3.Irreps,
        e3nn_compat_mode: bool = False,
    ):
        super().__init__()

        tp_infos = get_tp_infos_from_e3nn_instructions(
            instructions,
            irreps_in1,
            irreps_in2,
            irreps_out,
        )

        self.e3nn_compat_mode = e3nn_compat_mode
        self.connection_mode = tp_infos[4]
        self.stride_in1 = tp_infos[5][0]
        self.stride_in2 = tp_infos[5][1]
        self.stride_w = tp_infos[5][2]
        self.stride_out = tp_infos[5][3]
        self.irreps_in1 = irreps_in1
        self.irreps_in2 = irreps_in2
        self.irreps_out = irreps_out
        self.weight_numel = self.stride_w

        self._tp_info_fwd = tp_infos[0][3]
        self._tp_info_bwd_dgrad_in1 = tp_infos[1][3]
        self._tp_info_bwd_dgrad_in2 = tp_infos[2][3]
        self._tp_info_bwd_dgrad_w = tp_infos[3][3]

        self.register_buffer("_path_csr_offsets_fwd", tp_infos[0][0])
        self.register_buffer("_path_offsets_and_dims_fwd", tp_infos[0][1])
        self.register_buffer("_path_cg_values_fwd", tp_infos[0][2])

        self.register_buffer("_path_csr_offsets_bwd_dgrad_in1", tp_infos[1][0])
        self.register_buffer("_path_offsets_and_dims_bwd_dgrad_in1", tp_infos[1][1])
        self.register_buffer("_path_cg_values_bwd_dgrad_in1", tp_infos[1][2])

        self.register_buffer("_path_csr_offsets_bwd_dgrad_in2", tp_infos[2][0])
        self.register_buffer("_path_offsets_and_dims_bwd_dgrad_in2", tp_infos[2][1])
        self.register_buffer("_path_cg_values_bwd_dgrad_in2", tp_infos[2][2])

        self.register_buffer("_path_csr_offsets_bwd_dgrad_w", tp_infos[3][0])
        self.register_buffer("_path_offsets_and_dims_bwd_dgrad_w", tp_infos[3][1])
        self.register_buffer("_path_cg_values_bwd_dgrad_w", tp_infos[3][2])

    def forward(
        self,
        in1: torch.Tensor,
        in2: torch.Tensor,
        weight: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        args = (
            self._tp_info_fwd,
            self._tp_info_bwd_dgrad_in1,
            self._tp_info_bwd_dgrad_in2,
            self._tp_info_bwd_dgrad_w,
            self.connection_mode,
            self.stride_out,
        )
        if self.e3nn_compat_mode:
            in1 = transpose_irrep_to_channels_last(in1, self.irreps_in1)
            in2 = transpose_irrep_to_channels_last(in2, self.irreps_in2)

        out = _fused_tensor_product_autograd.apply(in1, in2, weight, args)

        if self.e3nn_compat_mode:
            out = transpose_irrep_to_m_last(out, self.irreps_out)

        return out


class FusedFullyConnectedTensorProduct(FusedTensorProduct):
    def __init__(
        self,
        irreps_in1: o3.Irreps,
        irreps_in2: o3.Irreps,
        irreps_out: o3.Irreps,
        connection_mode: str = "uvw",
        irrep_normalization: Optional[str] = None,
        path_normalization: Optional[str] = None,
        e3nn_compat_mode: bool = False,
    ):
        connection_mode = connection_mode.lower()
        supported_modes = {"uvw", "uvu", "uvv", "uuw"}
        if connection_mode not in supported_modes:
            msg = f"Invalid connection_mode, received {connection_mode}"
            msg += f", supported modes are {supported_modes}."
            raise ValueError(msg)

        instr = [
            (i_1, i_2, i_out, connection_mode, True, 1.0)
            for i_1, (_, ir_1) in enumerate(irreps_in1)
            for i_2, (_, ir_2) in enumerate(irreps_in2)
            for i_out, (_, ir_out) in enumerate(irreps_out)
            if ir_out in ir_1 * ir_2
        ]

        _tp = o3.TensorProduct(
            irreps_in1,
            irreps_in2,
            irreps_out,
            instr,
            internal_weights=False,
            shared_weights=False,
            irrep_normalization=irrep_normalization,
            path_normalization=path_normalization,
        )

        super().__init__(
            _tp.instructions,
            _tp.irreps_in1,
            _tp.irreps_in2,
            _tp.irreps_out,
            e3nn_compat_mode=e3nn_compat_mode,
        )


class FusedFullTensorProduct(FusedTensorProduct):
    def __init__(
        self,
        irreps_in1: o3.Irreps,
        irreps_in2: o3.Irreps,
        filter_ir_out: Iterator[o3.Irrep] = None,
        irrep_normalization: str = None,
        e3nn_compat_mode: bool = False,
    ):
        _tp = o3.FullTensorProduct(
            irreps_in1,
            irreps_in2,
            filter_ir_out=filter_ir_out,
            irrep_normalization=irrep_normalization,
        )

        super().__init__(
            _tp.instructions,
            _tp.irreps_in1,
            _tp.irreps_in2,
            _tp.irreps_out,
            e3nn_compat_mode=e3nn_compat_mode,
        )


class FusedElementWiseTensorProduct(FusedTensorProduct):
    def __init__(
        self,
        irreps_in1: o3.Irreps,
        irreps_in2: o3.Irrep,
        filter_ir_out=None,
        irrep_normalization: str = None,
        e3nn_compat_mode: bool = False,
    ):
        _tp = o3.ElementwiseTensorProduct(
            irreps_in1,
            irreps_in2,
            filter_ir_out=filter_ir_out,
            irrep_normalization=irrep_normalization,
        )

        super().__init__(
            _tp.instructions,
            _tp.irreps_in1,
            _tp.irreps_in2,
            _tp.irreps_out,
            e3nn_compat_mode=e3nn_compat_mode,
        )
