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

from typing import Tuple, Optional, Iterator, List, Union

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

from pylibcugraphops.pytorch.operators.fused_tensor_product_utils import (
    _irreps_to_segment_info,
    transpose_irrep_to_m_last,
    transpose_irrep_to_channels_last,
    get_tp_infos_from_e3nn_instructions,
    get_linear_infos_from_e3nn_instructions,
    get_sym_cont_u_matrices,
    init_sym_cont_weights,
    get_sym_contraction_path_offsets,
    contract_sym_cont_weights,
)


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)
        self.register_buffer("path_cg_values", path_cg_values, 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,
        )


class _fused_tensor_product_autograd(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        in1: torch.Tensor,
        in2: torch.Tensor,
        weight: Optional[torch.Tensor],
        bias: 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

        fwd_fun = get_operator_from_module(
            ops,
            "fused_tp_n2n_fwd",
            (
                in1.dtype,
                in2.dtype,
                weight.dtype if weight is not None else out.dtype,
                out.dtype,
            ),
        )
        fwd_fun(
            maybe_detach(out),
            maybe_detach(in1),
            maybe_detach(in2),
            maybe_detach(weight),
            maybe_detach(bias),
            connection_mode,
            tp_info_fwd,
            stream_id=stream,
        )

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

        return out

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

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

        return _fused_tensor_product_backward_autograd.apply(
            grad_out,
            in1,
            in2,
            weight,
            bias,
            (*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],
        bias: 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_bias,
            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)
        grad_bias = maybe_empty_like(bias, needs_input_grad_bias)

        stream = torch.cuda.current_stream().cuda_stream

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

        bwd_fun(
            grad_in1,
            grad_in2,
            grad_weight,
            grad_bias,
            maybe_detach(grad_out),
            maybe_detach(in1),
            maybe_detach(in2),
            maybe_detach(weight),
            maybe_detach(bias),
            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, grad_bias, None

    @staticmethod
    def backward(
        ctx, grad_grad_in1, grad_grad_in2, grad_grad_weight, grad_grad_bias, 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]

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

        bwd_bwd_fun(
            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,
        )

        if grad_grad_bias is not None:
            grad_grad_out += grad_grad_bias

        return grad_grad_out, grad_in1, grad_in2, grad_weight, None, 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,
        internal_weights: Optional[torch.Tensor] = None,
        shared_weights: Optional[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.internal_weights = False
        self.shared_weights = shared_weights
        if internal_weights is not None:
            self.internal_weights = True
            self.weight = nn.Parameter(internal_weights)

        if self.e3nn_compat_mode:
            _irreps_in1_segment_info = _irreps_to_segment_info(self.irreps_in1)
            _irreps_in2_segment_info = _irreps_to_segment_info(self.irreps_in2)
            _irreps_out_segment_info = _irreps_to_segment_info(self.irreps_out)

            self.register_buffer("_irreps_in1_segment_info", _irreps_in1_segment_info)
            self.register_buffer("_irreps_in2_segment_info", _irreps_in2_segment_info)
            self.register_buffer("_irreps_out_segment_info", _irreps_out_segment_info)

        self._tp_info_fwd = TensorProductPathData(*tp_infos[0])
        self._tp_info_bwd_dgrad_in1 = TensorProductPathData(*tp_infos[1])
        self._tp_info_bwd_dgrad_in2 = TensorProductPathData(*tp_infos[2])
        self._tp_info_bwd_dgrad_w = TensorProductPathData(*tp_infos[3])

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

        # TODO: better checks against shape matches and other things
        if self.internal_weights:
            weight = self.weight
        out = _fused_tensor_product_autograd.apply(in1, in2, weight, None, args)

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

        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,
        internal_weights: Optional[bool] = None,
        shared_weights: Optional[bool] = None,
    ):
        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
        ]

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            _tp = o3.TensorProduct(
                irreps_in1,
                irreps_in2,
                irreps_out,
                instr,
                internal_weights=internal_weights,
                shared_weights=shared_weights,
                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,
            internal_weights=_tp.weight.data if internal_weights else None,
            shared_weights=shared_weights,
        )


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,
    ):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            _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,
    ):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            _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,
        )


class FusedLinear(nn.Module):
    def __init__(
        self,
        irreps_in: o3.Irreps,
        irreps_out: o3.Irreps,
        e3nn_compat_mode: bool = False,
        internal_weights: Optional[bool] = None,
        shared_weights: Optional[bool] = None,
    ):
        super().__init__()

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            _tp = o3.Linear(
                irreps_in,
                irreps_out,
                internal_weights=internal_weights,
                shared_weights=shared_weights,
            )

        self.internal_weights = internal_weights
        self.shared_weights = shared_weights
        if internal_weights:
            self.weight = nn.Parameter(_tp.weight.data)

        tp_infos = get_linear_infos_from_e3nn_instructions(
            _tp.instructions,
            irreps_in,
            irreps_out,
        )

        self.e3nn_compat_mode = e3nn_compat_mode
        self.connection_mode = tp_infos[3]
        self.stride_in = tp_infos[4][0]
        self.stride_w = tp_infos[4][1]
        self.stride_out = tp_infos[4][3]
        self.irreps_in = irreps_in
        self.irreps_out = irreps_out
        self.weight_numel = self.stride_w

        if self.e3nn_compat_mode:
            _irreps_in_segment_info = _irreps_to_segment_info(self.irreps_in)
            _irreps_out_segment_info = _irreps_to_segment_info(self.irreps_out)

            self.register_buffer("_irreps_in_segment_info", _irreps_in_segment_info)
            self.register_buffer("_irreps_out_segment_info", _irreps_out_segment_info)

        self._tp_info_fwd = TensorProductPathData(*tp_infos[0])
        self._tp_info_bwd_dgrad_in = TensorProductPathData(*tp_infos[1])
        self._tp_info_bwd_dgrad_w = TensorProductPathData(*tp_infos[2])

    def forward(self, x: torch.Tensor, weight: Optional[torch.Tensor] = None):
        args = (
            self._tp_info_fwd.c_type(),
            self._tp_info_bwd_dgrad_in.c_type(),
            self._tp_info_bwd_dgrad_w.c_type(),
            self._tp_info_bwd_dgrad_w.c_type(),
            self.connection_mode,
            self.stride_out,
        )

        if self.e3nn_compat_mode:
            x = transpose_irrep_to_channels_last(x, self._irreps_in_segment_info)

        # TODO: better checks
        if self.internal_weights:
            weight = self.weight

        out = _fused_tensor_product_autograd.apply(x, weight, None, None, args)

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

        return out


class FusedMainContraction(nn.Module):
    def __init__(
        self,
        tp_infos: List[Tuple[torch.Tensor]],
        strides: List[int],
    ):
        super().__init__()

        self._tp_info_fwd = TensorProductPathData(*tp_infos[0])
        self._tp_info_bwd_dgrad_x = TensorProductPathData(*tp_infos[1])
        self._tp_info_bwd_dgrad_y = TensorProductPathData(*tp_infos[2])
        self._tp_info_bwd_dgrad_w = TensorProductPathData(*tp_infos[3])

        self.connection_mode = getattr(pylibcugraphops.ConnectionMode, "UVU")

        self.stride_x = strides[0]
        self.stride_y = strides[1]
        self.stride_w = strides[2]
        self.stride_o = strides[3]

    def forward(
        self,
        x: torch.Tensor,
        y: torch.Tensor,
        contracted_weight: torch.Tensor,
        next_contracted_weight: Optional[torch.Tensor],
    ):
        args = (
            self._tp_info_fwd.c_type(),
            self._tp_info_bwd_dgrad_x.c_type(),
            self._tp_info_bwd_dgrad_y.c_type(),
            self._tp_info_bwd_dgrad_w.c_type(),
            self.connection_mode,
            self.stride_o,
        )

        out = _fused_tensor_product_autograd.apply(
            x, y, contracted_weight, next_contracted_weight, args
        )

        return out


class FusedFeatureContraction(nn.Module):
    def __init__(
        self,
        tp_infos: List[Tuple[torch.Tensor]],
        strides: List[int],
    ):
        super().__init__()

        self._tp_info_fwd = TensorProductPathData(*tp_infos[0])
        self._tp_info_bwd_dgrad_x = TensorProductPathData(*tp_infos[1])
        self._tp_info_bwd_dgrad_w = TensorProductPathData(*tp_infos[2])

        self.connection_mode = getattr(pylibcugraphops.ConnectionMode, "UUU")

        self.stride_x = strides[0]
        self.stride_w = strides[1]
        self.stride_o = strides[3]

    def forward(
        self,
        x: torch.Tensor,
        contracted_weight: torch.Tensor,
        next_contracted_weight: Optional[torch.Tensor],
    ):
        args = (
            self._tp_info_fwd.c_type(),
            self._tp_info_bwd_dgrad_x.c_type(),
            self._tp_info_bwd_dgrad_w.c_type(),
            # duplication only for API compatibility until fixed, not accessed
            self._tp_info_bwd_dgrad_w.c_type(),
            self.connection_mode,
            self.stride_o,
        )

        out = _fused_tensor_product_autograd.apply(
            x, contracted_weight, None, next_contracted_weight, args
        )

        return out


class FusedSymmetricContraction(nn.Module):
    def __init__(
        self,
        irreps_in: o3.Irreps,
        irreps_out: o3.Irreps,
        correlation: int,
        num_elements: int,
        e3nn_compat_mode: bool = False,
    ):
        super().__init__()

        self.irreps_in = irreps_in
        self.irreps_out = irreps_out
        self.correlation = correlation
        self.num_irreps_out = len(irreps_out)
        self.num_elements = num_elements
        self.num_features = irreps_in.count((0, 1))
        self.e3nn_compat_mode = e3nn_compat_mode

        if self.e3nn_compat_mode:
            _irreps_out_segment_info = _irreps_to_segment_info(self.irreps_out)
            self.register_buffer("_irreps_out_segment_info", _irreps_out_segment_info)

        u_matrices = get_sym_cont_u_matrices(irreps_in, irreps_out, correlation)
        weights = init_sym_cont_weights(
            u_matrices, self.num_elements, self.num_features, correlation
        )

        weight_dict = {}
        for idx in range(self.num_irreps_out):
            for nu in range(0, correlation):
                self.register_buffer(f"u_matrix_{idx}_{nu}", u_matrices[idx][nu])
                weight_dict[f"weight_{idx}_{nu}"] = weights[idx][nu]
        self.weights = nn.ParameterDict(weight_dict)

        tp_infos, weight_indices, strides = get_sym_contraction_path_offsets(
            weights, u_matrices, correlation, self.num_features, self.num_elements
        )

        for nu in range(correlation):
            for i, idx in enumerate(weight_indices[nu]):
                self.register_buffer(f"weight_selection_index_{nu}_{i}", idx)
        self._idx_len = [len(idx) for idx in weight_indices]

        self.main_contraction = FusedMainContraction(
            tp_infos[correlation - 1], strides[correlation - 1]
        )
        feature_contractions = []
        for nu in range(0, correlation - 1):
            feature_contractions.append(
                FusedFeatureContraction(tp_infos[nu], strides[nu])
            )
        self.feature_contractions = nn.ModuleList(feature_contractions)

    def forward(self, x: torch.Tensor, y: torch.Tensor):
        if self.e3nn_compat_mode:
            # if compat mode, x is of shape [BS, NUM_FEAT, NUM_ELL]
            # we need [BS, NUM_ELL*NUM_FEAT]
            x = x.permute(0, 2, 1).contiguous()
            x = x.view(x.size(0), -1)

        u_matrices = self._get_u_matrices_from_buffers()
        weights = self._get_weights_from_buffers()
        contracted_weights = contract_sym_cont_weights(
            weights, u_matrices, self.correlation
        )

        out = None

        if self.correlation == 1:
            cw = self._index_contracted_weights(
                self.correlation - 1, contracted_weights
            )
            cw = cw.view(-1)
            out = self.main_contraction(x, y, cw, None)

        else:
            cw = self._index_contracted_weights(
                self.correlation - 1, contracted_weights
            )
            cw = cw.view(-1)
            cw_next = self._index_contracted_weights(
                self.correlation - 2, contracted_weights
            )
            cw_next = self._contract_weight(cw_next, y)
            out = self.main_contraction(x, y, cw, cw_next)

        for nu in range(self.correlation - 2, -1, -1):
            cw_next = None
            if nu - 1 >= 0:
                cw_next = self._index_contracted_weights(nu - 1, contracted_weights)
                cw_next = self._contract_weight(cw_next, y)
            out = self.feature_contractions[nu](x, out, cw_next)

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

        return out

    def _contract_weight(self, cw: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        cw = cw.view(-1, cw.size(-1))
        c_tensor = torch.einsum("be,ae->ba", y, cw)
        return c_tensor

    def _index_contracted_weights(
        self, nu: int, contracted_weights: List[torch.Tensor]
    ) -> torch.Tensor:
        idx = self._get_weight_selection_indices(nu)
        cw = contracted_weights[nu][idx]
        return cw

    def _get_u_matrices_from_buffers(self):
        u_matrices = []
        for idx in range(self.num_irreps_out):
            tmp = []
            for nu in range(0, self.correlation):
                tmp.append(getattr(self, f"u_matrix_{idx}_{nu}"))
            u_matrices.append(tmp)
        return u_matrices

    def _get_weights_from_buffers(self):
        weights = []
        for idx in range(self.num_irreps_out):
            tmp = []
            for nu in range(0, self.correlation):
                tmp.append(self.weights[f"weight_{idx}_{nu}"])
            weights.append(tmp)
        return weights

    def _load_weights_from_reference(self, module: nn.Module):
        for i in range(len(self.irreps_out)):
            ref_layer = module.contractions[i]
            for nu in range(self.correlation):
                if nu == (self.correlation - 1):
                    weight = ref_layer.weights_max
                else:
                    weight = ref_layer.weights[self.correlation - 2 - nu]

                self.weights[f"weight_{i}_{nu}"].data = weight.data.clone()

    def _get_weight_selection_indices(self, nu: int) -> Tuple:
        return tuple(
            [
                getattr(self, f"weight_selection_index_{nu}_{i}")
                for i in range(self._idx_len[nu])
            ]
        )


class LoopedFusedSymmetricContraction(nn.Module):
    def __init__(
        self,
        irreps_in: o3.Irreps,
        irreps_out: o3.Irreps,
        correlation: Union[int, Iterator[int]],
        num_elements: int,
        e3nn_compat_mode: bool = False,
    ):
        # unlike the original implementation, FusedContraction
        # already would work for many different irreps at once
        # however, the current implementation for now requires rather
        # a lot of shared memory and to be compatible with having
        # different correlations for different irreps, this can be
        # used instead

        super().__init__()
        self.irreps_in = irreps_in
        self.irreps_out = irreps_out
        self.correlation = correlation

        if isinstance(correlation, int):
            correlation = [correlation] * len(irreps_out)
        else:
            if len(correlation) != len(irreps_out):
                raise ValueError("List of iteration values does not match irreps_out.")

        self.contractions = torch.nn.ModuleList()
        for i, irrep_out in enumerate(self.irreps_out):
            self.contractions.append(
                FusedSymmetricContraction(
                    irreps_in,
                    o3.Irreps(str(irrep_out)),
                    correlation=correlation[i],
                    num_elements=num_elements,
                    e3nn_compat_mode=e3nn_compat_mode,
                )
            )

    def forward(self, x: torch.Tensor, y: torch.Tensor):
        outs = [contraction(x, y) for contraction in self.contractions]
        return torch.cat(outs, dim=1)

    def _load_weights_from_reference(self, module: nn.Module):
        for i in range(len(self.irreps_out)):
            ref_layer = module.contractions[i]
            for nu in range(self.correlation):
                if nu == (self.correlation - 1):
                    weight = ref_layer.weights_max
                else:
                    weight = ref_layer.weights[self.correlation - 2 - nu]

                self.contractions[i].weights[
                    f"weight_{0}_{nu}"
                ].data = weight.data.clone()
