# Copyright Universitat Pompeu Fabra 2020-2023  https://www.compscience.org
# Distributed under the MIT License.
# (See accompanying file README.md file or copy at http://opensource.org/licenses/MIT)

import torch
from typing import Optional, Tuple
from torch import Tensor, nn
from torchmdnet.models.utils import (
    CosineCutoff,
    OptimizedDistance,
    rbf_class_mapping,
    act_class_mapping,
)

__all__ = ["TensorNet"]


def _decompose_tensor(tensor):
    """Full tensor decomposition into irreducible components.

    shapes are [N, 3, 3, F]
    """
    A = 0.5 * (tensor - tensor.transpose(1, 2))
    S = tensor - A
    I = (tensor.diagonal(dim1=1, dim2=2)).mean(-1)
    S = S - I_to_tensor(I)
    return I, A, S


def _compose_tensor(I, A, S):
    """Compose tensor from irreducible components.

    input shapes:
    I: [batch_size, hidden_channels]
    A: [batch_size, 3, 3, hidden_channels]
    S: [batch_size, 3, 3, hidden_channels]
    output shape:
    [batch_size, 3, 3, hidden_channels]"""
    return I_to_tensor(I) + A + S


def _tensor_matmul_o3(Y, msg):
    A = torch.matmul(msg.permute(0, 3, 1, 2), Y.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
    B = torch.matmul(Y.permute(0, 3, 1, 2), msg.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
    return A + B


def _tensor_matmul_so3(Y, msg):
    return torch.matmul(Y.permute(0, 3, 1, 2), msg.permute(0, 3, 1, 2)).permute(
        0, 2, 3, 1
    )


try:
    from torchmdnet.extensions.warp_ops import (
        graph_transform,
        fn_radial_message_passing as fn_tensornet_embedding_message_passing,
        fn_message_passing as fn_tensornet_interaction_message_passing,
        fn_compose_tensor as compose_tensor,
        fn_decompose_tensor as decompose_tensor,
        fn_tensor_matmul_o3_3x3 as tensor_matmul_o3,
        fn_tensor_matmul_so3_3x3 as tensor_matmul_so3,
        fn_tensor_norm3,
    )

    OPT = True
except ImportError:
    import warnings

    warnings.warn(
        "Could not import warp_ops (warp-lang not installed?). "
        "Falling back to pure-Python TensorNet ops. "
        "Install warp-lang for better performance: pip install warp-lang",
        ImportWarning,
        stacklevel=2,
    )
    OPT = False
    compose_tensor = _compose_tensor
    decompose_tensor = _decompose_tensor
    tensor_matmul_o3 = _tensor_matmul_o3
    tensor_matmul_so3 = _tensor_matmul_so3


def vector_to_skewtensor(vector):
    """Creates a skew-symmetric tensor from a vector."""
    B, _, F = vector.shape
    zero = torch.zeros((B, F), device=vector.device, dtype=vector.dtype)
    tensor = torch.stack(
        (
            zero,
            -vector[..., 2, :],
            vector[..., 1, :],
            vector[..., 2, :],
            zero,
            -vector[..., 0, :],
            -vector[..., 1, :],
            vector[..., 0, :],
            zero,
        ),
        dim=1,
    )
    tensor = tensor.view(B, 3, 3, F)
    return tensor


def skewtensor_to_vector(tensor):
    """Converts 3x3 skewtensor A to a vector

    input shape must be [N, 3, 3, F]
    output shape is [N, 3, F]
    """

    tensor = tensor.flatten(1, 2)
    vector = 0.5 * torch.stack(
        (
            tensor[:, 7, :] - tensor[:, 5, :],
            tensor[:, 2, :] - tensor[:, 6, :],
            tensor[:, 3, :] - tensor[:, 1, :],
        ),
        dim=1,
    )
    return vector


def I_to_tensor(I):
    """Converts scalar to 3x3 I tensor with shape [N, 3, 3, F]"""
    return (
        I[:, None, None, :]
        * torch.eye(3, 3, device=I.device, dtype=I.dtype)[None, ..., None]
    )


def outer_to_symtensor(tensor):
    """Creates a symmetric traceless tensor from the outer product tensor

    shapes are [N, 3, 3, F]
    """
    S = 0.5 * (tensor + tensor.transpose(1, 2))
    I = (tensor.diagonal(dim1=1, dim2=2)).mean(-1)
    S = S - I_to_tensor(I)
    return S


def tensor_norm(tensor):
    """Computes Frobenius norm."""
    return (tensor**2).sum((1, 2))


class TensorNet(nn.Module):
    r"""TensorNet's architecture. From
    TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials; G. Simeon and G. de Fabritiis.
    NeurIPS 2023.

    This function optionally supports periodic boundary conditions with arbitrary triclinic boxes.
    For a given cutoff, :math:`r_c`, the box vectors :math:`\vec{a},\vec{b},\vec{c}` must satisfy certain requirements:

    .. math::

      \begin{align*}
      a_y = a_z = b_z &= 0 \\
      a_x, b_y, c_z &\geq 2 r_c \\
      a_x &\geq 2  b_x \\
      a_x &\geq 2  c_x \\
      b_y &\geq 2  c_y
      \end{align*}

    These requirements correspond to a particular rotation of the system and reduced form of the vectors, as well as the requirement that the cutoff be no larger than half the box width.

    Args:
        hidden_channels (int, optional): Hidden embedding size.
            (default: :obj:`128`)
        num_layers (int, optional): The number of interaction layers.
            (default: :obj:`2`)
        num_rbf (int, optional): The number of radial basis functions :math:`\mu`.
            (default: :obj:`32`)
        rbf_type (string, optional): The type of radial basis function to use.
            (default: :obj:`"expnorm"`)
        trainable_rbf (bool, optional): Whether to train RBF parameters with
            backpropagation. (default: :obj:`False`)
        activation (string, optional): The type of activation function to use.
            (default: :obj:`"silu"`)
        cutoff_lower (float, optional): Lower cutoff distance for interatomic interactions.
            (default: :obj:`0.0`)
        cutoff_upper (float, optional): Upper cutoff distance for interatomic interactions.
            (default: :obj:`4.5`)
        max_z (int, optional): Maximum atomic number. Used for initializing embeddings.
            (default: :obj:`128`)
        max_num_neighbors (int, optional): Maximum number of neighbors to return for a
            given node/atom when constructing the molecular graph during forward passes.
            (default: :obj:`64`)
        equivariance_invariance_group (string, optional): Group under whose action on input
            positions internal tensor features will be equivariant and scalar predictions
            will be invariant. O(3) or SO(3).
            (default :obj:`"O(3)"`)
        box_vecs (Tensor, optional):
            The vectors defining the periodic box.  This must have shape `(3, 3)`,
            where `box_vectors[0] = a`, `box_vectors[1] = b`, and `box_vectors[2] = c`.
            If this is omitted, periodic boundary conditions are not applied.
            (default: :obj:`None`)
        static_shapes (bool, optional): Whether to enforce static shapes.
            Makes the model CUDA-graph compatible.
            (default: :obj:`True`)
    """

    def __init__(
        self,
        hidden_channels=128,
        num_layers=2,
        num_rbf=32,
        rbf_type="expnorm",
        trainable_rbf=False,
        activation="silu",
        cutoff_lower=0,
        cutoff_upper=4.5,
        max_num_neighbors=64,
        max_z=128,
        equivariance_invariance_group="O(3)",
        static_shapes=True,
        dtype=torch.float32,
        box_vecs=None,
    ):
        super(TensorNet, self).__init__()

        assert rbf_type in rbf_class_mapping, (
            f'Unknown RBF type "{rbf_type}". '
            f'Choose from {", ".join(rbf_class_mapping.keys())}.'
        )
        assert activation in act_class_mapping, (
            f'Unknown activation function "{activation}". '
            f'Choose from {", ".join(act_class_mapping.keys())}.'
        )

        assert equivariance_invariance_group in ["O(3)", "SO(3)"], (
            f'Unknown group "{equivariance_invariance_group}". '
            f"Choose O(3) or SO(3)."
        )
        self.hidden_channels = hidden_channels
        self.equivariance_invariance_group = equivariance_invariance_group
        self.num_layers = num_layers
        self.num_rbf = num_rbf
        self.rbf_type = rbf_type
        self.activation = activation
        self.cutoff_lower = cutoff_lower
        self.cutoff_upper = cutoff_upper
        act_class = act_class_mapping[activation]
        self.distance_expansion = rbf_class_mapping[rbf_type](
            cutoff_lower, cutoff_upper, num_rbf, trainable_rbf
        )
        self.tensor_embedding = TensorEmbedding(
            hidden_channels,
            num_rbf,
            act_class,
            cutoff_lower,
            cutoff_upper,
            trainable_rbf,
            max_z,
            dtype,
        )

        self.layers = nn.ModuleList()
        if num_layers != 0:
            for _ in range(num_layers):
                self.layers.append(
                    Interaction(
                        num_rbf,
                        hidden_channels,
                        act_class,
                        cutoff_lower,
                        cutoff_upper,
                        equivariance_invariance_group,
                        dtype,
                    )
                )
        self.linear = nn.Linear(3 * hidden_channels, hidden_channels, dtype=dtype)
        self.out_norm = nn.LayerNorm(3 * hidden_channels, dtype=dtype)
        self.act = act_class()
        # Resize to fit set to false ensures Distance returns a statically-shaped tensor of size max_num_pairs=pos.size*max_num_neigbors
        # negative max_num_pairs argument means "per particle"
        # long_edge_index set to False saves memory and spares some kernel launches by keeping neighbor indices as int32.
        self.static_shapes = static_shapes
        self.distance = OptimizedDistance(
            cutoff_lower,
            cutoff_upper,
            max_num_pairs=-max_num_neighbors,
            return_vecs=True,
            loop=True,
            resize_to_fit=not self.static_shapes,
            box=box_vecs,
            long_edge_index=True,
        )

        # for torchscript
        self.opt = OPT

        self.reset_parameters()

    def reset_parameters(self):
        self.tensor_embedding.reset_parameters()
        for layer in self.layers:
            layer.reset_parameters()
        self.linear.reset_parameters()
        self.out_norm.reset_parameters()

    def setup_for_inference(self, z, batch):
        self.tensor_embedding.setup_for_inference(z, batch)
        self.inference_mode = True

    def forward(
        self,
        z: Tensor,
        pos: Tensor,
        batch: Tensor,
        box: Optional[Tensor] = None,
        q: Optional[Tensor] = None,
        s: Optional[Tensor] = None,
    ) -> Tuple[Tensor, Optional[Tensor], Tensor, Tensor, Tensor]:
        # Obtain graph, with distances and relative position vectors
        edge_index, edge_weight, edge_vec = self.distance(pos, batch, box)

        if self.opt:
            # perpare graph indices for message passing
            row_data, row_indices, row_indptr, col_data, col_indices, col_indptr = (
                graph_transform(edge_index.int(), z.shape[0])
            )
        else:
            row_data, row_indices, row_indptr, col_data, col_indices, col_indptr = (
                torch.empty(0, dtype=torch.long),
                torch.empty(0, dtype=torch.long),
                torch.empty(0, dtype=torch.long),
                torch.empty(0, dtype=torch.long),
                torch.empty(0, dtype=torch.long),
                torch.empty(0, dtype=torch.long),
            )

        # This assert convinces TorchScript that edge_vec is a Tensor and not an Optional[Tensor]
        assert (
            edge_vec is not None
        ), "Distance module did not return directional information"
        # Distance module returns -1 for non-existing edges, to avoid having to resize the tensors when we want to ensure static shapes (for CUDA graphs) we make all non-existing edges pertain to a ghost atom
        # Total charge q is a molecule-wise property. We transform it into an atom-wise property, with all atoms belonging to the same molecule being assigned the same charge q
        if q is None:
            q = torch.zeros_like(z, device=z.device, dtype=z.dtype)
        else:
            q = q[batch]
        zp = z
        if self.static_shapes:
            mask = (edge_index[0] < 0).unsqueeze(0).expand_as(edge_index)
            zp = torch.cat((z, torch.zeros(1, device=z.device, dtype=z.dtype)), dim=0)

            if not self.opt:
                q = torch.cat(
                    (q, torch.zeros(1, device=q.device, dtype=q.dtype)), dim=0
                )
            # I trick the model into thinking that the masked edges pertain to the extra atom
            # WARNING: This can hurt performance if max_num_pairs >> actual_num_pairs
            edge_index = edge_index.masked_fill(mask, z.shape[0])
            edge_weight = edge_weight.masked_fill(mask[0], 0)
            edge_vec = edge_vec.masked_fill(
                mask[0].unsqueeze(-1).expand_as(edge_vec), 0
            )

        edge_attr = self.distance_expansion(edge_weight)
        mask = edge_index[0] == edge_index[1]
        # Normalizing edge vectors by their length can result in NaNs, breaking Autograd.
        # I avoid dividing by zero by setting the weight of self edges and self loops to 1
        edge_vec = edge_vec / edge_weight.masked_fill(mask, 1).unsqueeze(1)

        graph_info = {
            "edge_index": edge_index,
            "row_data": row_data,
            "row_indices": row_indices,
            "row_indptr": row_indptr,
            "col_data": col_data,
            "col_indices": col_indices,
            "col_indptr": col_indptr,
        }

        X = self.tensor_embedding(
            zp, graph_info, edge_weight, edge_vec, edge_attr
        )  # [N, 3, 3, F]
        for layer in self.layers:
            X = layer(X, graph_info, edge_weight, edge_attr, q)  # [N, 3, 3, F]

        if not self.opt:
            I, A, S = decompose_tensor(X)
            x = torch.cat((3 * I**2, tensor_norm(A), tensor_norm(S)), dim=-1)
            x = self.out_norm(x)
            x = self.act(self.linear((x)))

            # Remove the extra atom
            if self.static_shapes:
                x = x[:-1]

        else:

            x = fn_tensor_norm3(X)
            x = self.out_norm(x)
            x = self.act(self.linear((x)))

            # with the OPT code and static shapes there is no dummy atom in the outputs

        return x, None, z, pos, batch


def tensornet_embedding_message_passing(
    edge_vec_norm, edge_attr_processed, edge_index, num_atoms: int
):
    # written to mirror the optmized fn_tensornet_embedding_message_passing

    E, _, F = edge_attr_processed.shape

    Iij = edge_attr_processed[:, 0, :]  # [num_edges, hidden_channels]
    Aij = edge_attr_processed[:, 1, None, :] * edge_vec_norm.unsqueeze(
        -1
    )  # [num_edges, 3, hidden_channels]
    _outer = torch.matmul(
        edge_vec_norm.unsqueeze(2), edge_vec_norm.unsqueeze(-2)
    )  # [num_edges, 3, 3]
    Sij = edge_attr_processed[:, 2, None, None, :] * _outer.unsqueeze(
        -1
    )  # [num_edges, 3, 3, hidden_channels]

    source_scalar = torch.zeros(
        num_atoms, F, device=edge_attr_processed.device, dtype=Iij.dtype
    )
    source_vector = torch.zeros(
        num_atoms, 3, F, device=edge_attr_processed.device, dtype=Iij.dtype
    )
    source_tensor = torch.zeros(
        num_atoms, 3, 3, F, device=edge_attr_processed.device, dtype=Iij.dtype
    )
    I = source_scalar.index_add(
        dim=0, index=edge_index[0], source=Iij
    )  # [num_atoms, hidden_channels]
    A_vec = source_vector.index_add(
        dim=0, index=edge_index[0], source=Aij
    )  # [num_atoms, 3, hidden_channels]
    S = source_tensor.index_add(
        dim=0, index=edge_index[0], source=Sij
    )  # [num_atoms, 3, 3, hidden_channels]

    A = vector_to_skewtensor(A_vec)  # [num_atoms, 3, 3, hidden_channels]
    S = outer_to_symtensor(S)  # [num_atoms, 3, 3, hidden_channels]

    return I, A, S


class TensorEmbedding(nn.Module):
    """Tensor embedding layer.

    :meta private:
    """

    def __init__(
        self,
        hidden_channels,
        num_rbf,
        activation,
        cutoff_lower,
        cutoff_upper,
        trainable_rbf=False,
        max_z=128,
        dtype=torch.float32,
    ):
        super(TensorEmbedding, self).__init__()

        self.hidden_channels = hidden_channels
        self.distance_proj1 = nn.Linear(num_rbf, hidden_channels, dtype=dtype)
        self.distance_proj2 = nn.Linear(num_rbf, hidden_channels, dtype=dtype)
        self.distance_proj3 = nn.Linear(num_rbf, hidden_channels, dtype=dtype)
        self.cutoff = CosineCutoff(cutoff_lower, cutoff_upper)
        self.max_z = max_z
        self.emb = nn.Embedding(max_z, hidden_channels, dtype=dtype)
        self.emb2 = nn.Linear(2 * hidden_channels, hidden_channels, dtype=dtype)
        self.act = activation()
        self.linears_tensor = nn.ModuleList()
        for _ in range(3):
            self.linears_tensor.append(
                nn.Linear(hidden_channels, hidden_channels, bias=False)
            )
        self.linears_scalar = nn.ModuleList()
        self.linears_scalar.append(
            nn.Linear(hidden_channels, 2 * hidden_channels, bias=True, dtype=dtype)
        )
        self.linears_scalar.append(
            nn.Linear(2 * hidden_channels, 3 * hidden_channels, bias=True, dtype=dtype)
        )
        self.init_norm = nn.LayerNorm(hidden_channels, dtype=dtype)

        # for torchscript compatibility
        self.Zij_map = torch.empty(0, 0, 0)
        self.opt = OPT

        self.reset_parameters()

    def reset_parameters(self):
        self.distance_proj1.reset_parameters()
        self.distance_proj2.reset_parameters()
        self.distance_proj3.reset_parameters()
        self.emb.reset_parameters()
        self.emb2.reset_parameters()
        for linear in self.linears_tensor:
            linear.reset_parameters()
        for linear in self.linears_scalar:
            linear.reset_parameters()
        self.init_norm.reset_parameters()
        self.inference_mode = False

    def setup_for_inference(self, z, batch):

        # we can precompute the edge-wise atomic embeddings
        with torch.no_grad():
            max_z = z.max().item()
            _zij_map = torch.zeros((max_z + 1, max_z + 1, self.hidden_channels))

            for zi in torch.unique(z):
                Zi = self.emb(zi)
                for zj in torch.unique(z):
                    Zj = self.emb(zj)
                    _zij_map[zi, zj] = self.emb2(torch.cat([Zi, Zj], dim=-1))

        del self.Zij_map
        self.register_buffer("Zij_map", _zij_map)
        self.inference_mode = True

    def _get_atomic_number_message(self, z: Tensor, edge_index: Tensor) -> Tensor:

        if self.inference_mode:
            if torch.jit.is_scripting():
                # torchscript does not like the buffer
                Zij = torch.as_tensor(self.Zij_map)[z[edge_index[0]], z[edge_index[1]]]
            else:
                Zij = self.Zij_map[z[edge_index[0]], z[edge_index[1]]]
        else:
            Z = self.emb(z)
            Zij = self.emb2(
                Z.index_select(0, edge_index.t().reshape(-1)).view(
                    -1, self.hidden_channels * 2
                )
            )
        return Zij

    def forward(
        self,
        z: Tensor,
        graph_info: dict[str, Tensor],
        edge_weight: Tensor,
        edge_vec_norm: Tensor,
        edge_attr: Tensor,
    ) -> Tensor:

        edge_index = graph_info["edge_index"]

        Zij = self._get_atomic_number_message(
            z, edge_index
        )  # [num_edges, hidden_channels]

        dp1 = self.distance_proj1(edge_attr)  # [num_edges, hidden_channels]
        dp2 = self.distance_proj2(edge_attr)  # [num_edges, hidden_channels]
        dp3 = self.distance_proj3(edge_attr)  # [num_edges, hidden_channels]
        CZij = (
            self.cutoff(edge_weight).unsqueeze(-1) * Zij
        )  # [num_edges, hidden_channels]

        edge_attr_processed = CZij.unsqueeze(1) * torch.stack(
            [dp1, dp2, dp3], dim=1
        )  # [num_edges, 3, hidden_channels]

        if not self.opt:
            I, A, S = tensornet_embedding_message_passing(
                edge_vec_norm, edge_attr_processed, edge_index, z.shape[0]
            )  # [num_atoms, hidden_channels], [num_atoms, 3, 3, hidden_channels], [num_atoms, 3, 3, hidden_channels]

        else:
            row_indptr, row_indices, row_data = (
                graph_info["row_indptr"],
                graph_info["row_indices"],
                graph_info["row_data"],
            )
            I, A, S = fn_tensornet_embedding_message_passing(
                edge_vec_norm,
                edge_attr_processed,
                row_data,
                row_indptr,
            )  # [num_atoms, 1, hidden_channels], [num_atoms, 3, hidden_channels], [num_atoms, 5, hidden_channels]

        X = compose_tensor(I, A, S)  # [num_atoms, 3, 3, hidden_channels]

        norm = self.init_norm(tensor_norm(X))
        for linear_scalar in self.linears_scalar:
            norm = self.act(linear_scalar(norm))

        norm = norm.reshape(-1, 3, self.hidden_channels)

        if not self.opt:
            I = (
                self.linears_tensor[0](I) * norm[:, 0, :]
            )  # [num_atoms, hidden_channels]
            A = (
                self.linears_tensor[1](A) * norm[:, 1, None, None, :]
            )  # [num_atoms, 3, 3, hidden_channels]
            S = (
                self.linears_tensor[2](S) * norm[:, 2, None, None, :]
            )  # [num_atoms, 3, 3, hidden_channels]

        else:
            I = (
                self.linears_tensor[0](I) * norm[:, 0, None, :]
            )  # [num_atoms, 1, hidden_channels]
            A = (
                self.linears_tensor[1](A) * norm[:, 1, None, :]
            )  # [num_atoms, 3, hidden_channels]
            S = (
                self.linears_tensor[2](S) * norm[:, 2, None, :]
            )  # [num_atoms, 5, hidden_channels]

        X = compose_tensor(I, A, S)  # [num_atoms, 3, 3, hidden_channels]

        return X


def tensornet_interaction_message_passing(
    I, A, S, edge_attr_processed, edge_index, natoms: int
):

    # written to mirror the optimized fn_tensornet_interaction_message_passing

    A_vec = skewtensor_to_vector(A)  # [num_atoms, 3, hidden_channels]

    factor_scalar = edge_attr_processed[..., 0, :]
    factor_vector = edge_attr_processed[..., 1, None, :]
    factor_tensor = edge_attr_processed[..., 2, None, None, :]

    I = scalar_message_passing(
        edge_index, factor_scalar, I, natoms
    )  # [num_atoms, hidden_channels]
    A_vec = vector_message_passing(
        edge_index, factor_vector, A_vec, natoms
    )  # [num_atoms, 3, hidden_channels]
    S = tensor_message_passing(
        edge_index, factor_tensor, S, natoms
    )  # [num_atoms, 3, 3, hidden_channels]

    A = vector_to_skewtensor(A_vec)  # [num_atoms, 3, 3, hidden_channels]
    return I, A, S


def scalar_message_passing(
    edge_index: Tensor, factor: Tensor, tensor: Tensor, natoms: int
) -> Tensor:
    msg = factor * tensor.index_select(0, edge_index[1])  # [num_edges, hidden_channels]
    shape = (natoms, tensor.shape[1])
    tensor_m = torch.zeros(*shape, device=tensor.device, dtype=tensor.dtype)
    tensor_m = tensor_m.index_add(0, edge_index[0], msg)
    return tensor_m


def vector_message_passing(
    edge_index: Tensor, factor: Tensor, tensor: Tensor, natoms: int
) -> Tensor:
    msg = factor * tensor.index_select(
        0, edge_index[1]
    )  # [num_edges, 3, hidden_channels]
    shape = (natoms, tensor.shape[1], tensor.shape[2])
    tensor_m = torch.zeros(*shape, device=tensor.device, dtype=tensor.dtype)
    tensor_m = tensor_m.index_add(0, edge_index[0], msg)
    return tensor_m


def tensor_message_passing(
    edge_index: Tensor, factor: Tensor, tensor: Tensor, natoms: int
) -> Tensor:
    msg = factor * tensor.index_select(
        0, edge_index[1]
    )  # [num_edges, 3, 3, hidden_channels]
    shape = (natoms, tensor.shape[1], tensor.shape[2], tensor.shape[3])
    tensor_m = torch.zeros(*shape, device=tensor.device, dtype=tensor.dtype)
    tensor_m = tensor_m.index_add(0, edge_index[0], msg)
    return tensor_m


class Interaction(nn.Module):
    """Interaction layer.

    :meta private:
    """

    def __init__(
        self,
        num_rbf,
        hidden_channels,
        activation,
        cutoff_lower,
        cutoff_upper,
        equivariance_invariance_group,
        dtype=torch.float32,
    ):
        super(Interaction, self).__init__()

        self.num_rbf = num_rbf
        self.hidden_channels = hidden_channels
        self.cutoff = CosineCutoff(cutoff_lower, cutoff_upper)
        self.linears_scalar = nn.ModuleList()
        self.linears_scalar.append(
            nn.Linear(num_rbf, hidden_channels, bias=True, dtype=dtype)
        )
        self.linears_scalar.append(
            nn.Linear(hidden_channels, 2 * hidden_channels, bias=True, dtype=dtype)
        )
        self.linears_scalar.append(
            nn.Linear(2 * hidden_channels, 3 * hidden_channels, bias=True, dtype=dtype)
        )
        self.linears_tensor = nn.ModuleList()
        for _ in range(6):
            self.linears_tensor.append(
                nn.Linear(hidden_channels, hidden_channels, bias=False)
            )
        self.act = activation()
        self.equivariance_invariance_group = equivariance_invariance_group
        self.opt = OPT
        self.reset_parameters()

    def reset_parameters(self):
        for linear in self.linears_scalar:
            linear.reset_parameters()
        for linear in self.linears_tensor:
            linear.reset_parameters()

    def forward(
        self,
        X: Tensor,
        graph_info: dict[str, Tensor],
        edge_weight: Tensor,
        edge_attr: Tensor,
        q: Tensor,
    ) -> Tensor:

        C = self.cutoff(edge_weight)
        for linear_scalar in self.linears_scalar:
            edge_attr = self.act(linear_scalar(edge_attr))
        edge_attr = (edge_attr * C.view(-1, 1)).reshape(
            edge_attr.shape[0], 3, self.hidden_channels
        )

        X = X / (tensor_norm(X) + 1)[:, None, None, :]

        I, A, S = decompose_tensor(
            X
        )  # ([num_atoms, hidden_channels], [num_atoms, 3, 3, hidden_channels], [num_atoms, 3, 3, hidden_channels])
        # or ([num_atoms, 1, hidden_channels], [num_atoms, 3, hidden_channels], [num_atoms, 5, hidden_channels])

        I = self.linears_tensor[0](I)
        A = self.linears_tensor[1](A)
        S = self.linears_tensor[2](S)
        Y = compose_tensor(I, A, S)  # [num_atoms, 3, 3, hidden_channels]

        if not self.opt:
            edge_index = graph_info["edge_index"]

            Im, Am, Sm = tensornet_interaction_message_passing(
                I, A, S, edge_attr, edge_index, X.shape[0]
            )  # ([num_atoms, hidden_channels], [num_atoms, 3, 3, hidden_channels], [num_atoms, 3, 3, hidden_channels])

        else:

            row_data = graph_info["row_data"]
            row_indices = graph_info["row_indices"]
            row_indptr = graph_info["row_indptr"]
            col_data = graph_info["col_data"]
            col_indices = graph_info["col_indices"]
            col_indptr = graph_info["col_indptr"]

            Im, Am, Sm = fn_tensornet_interaction_message_passing(
                I,
                A,
                S,
                edge_attr,
                col_data,
                col_indices,
                col_indptr,
                row_data,
                row_indices,
                row_indptr,
            )  # ([num_atoms, 1, hidden_channels], [num_atoms, 3, hidden_channels], [num_atoms, 5, hidden_channels])

        msg = compose_tensor(Im, Am, Sm)  # [num_atoms, 3, 3, hidden_channels]

        if self.equivariance_invariance_group == "O(3)":
            C = (1 + 0.1 * q[..., None, None, None]) * tensor_matmul_o3(Y, msg)
            I, A, S = decompose_tensor(C)
        if self.equivariance_invariance_group == "SO(3)":
            C = 2 * tensor_matmul_so3(Y, msg)
            I, A, S = decompose_tensor(C)

        normp1 = tensor_norm(C) + 1

        if not self.opt:
            I, A, S = (
                I / normp1,
                A / normp1[..., None, None, :],
                S / normp1[..., None, None, :],
            )
        else:
            I = I / normp1.unsqueeze(1)
            A = A / normp1.unsqueeze(1)
            S = S / normp1.unsqueeze(1)

        I = self.linears_tensor[3](I)
        A = self.linears_tensor[4](A)
        S = self.linears_tensor[5](S)
        dX = compose_tensor(I, A, S)
        X = X + dX + (1 + 0.1 * q[..., None, None, None]) * tensor_matmul_so3(dX, dX)

        return X
