import math
from typing import Optional, List, Dict, Tuple
import torch
from deepchem.models.torch_models.layers import Fiber
from deepchem.utils.cache_utils import cached_dirpklgz


def fiber2head(F: Dict[str, torch.Tensor],
               h: int,
               structure: Fiber,
               squeeze: bool = False) -> torch.Tensor:
    """
    Converts SE(3)-equivariant features into multi-head format for attention.
    This function reshapes and concatenates fiber-based features into a format
    suitable for multi-head attention.

    - Input Fiber Representation:
      $$ F[d] \in \mathbb{R}^{B \times N \times M_d \times (2d+1)} $$
    - After Multi-Head Transformation:
      $$ F[d] \in \mathbb{R}^{B \times N \times H \times (M_d / H) \times (2d+1)} $$

    where:
    - `B`: Batch size.
    - `N`: Number of nodes.
    - `M_d`: Multiplicity (number of channels for degree `d`).
    - `2d+1`: Number of components in degree `d`.
    - `H`: Number of attention heads.
    - `M_d / H`: Channels per head.

    Parameters
    ----------
    - F (Dict[str, torch.Tensor]):
      Dictionary of SE(3)-equivariant features, where keys are **degrees (`d`).
    - h (int):
      Number of attention heads.
    - structure (Fiber):
      Fiber representation defining feature multiplicities and degrees.
    - squeeze (bool, optional, default=`False`):
      - If `True`: Concatenates along last feature dimension.
      - If `False`: Concatenates along second-last dimension.

    Returns
    -------
    - torch.Tensor:
      - Reshaped tensor with multi-head representation.
      - Shape: `(batch, num_nodes, num_heads, channels_per_head, 2d+1)`

    Example
    -------
    >>> import torch
    >>> from deepchem.models.torch_models.layers import Fiber
    >>> from deepchem.utils.equivariance_utils import fiber2head

    # Create Fiber Structure
    >>> fiber_structure = Fiber(dictionary={0: 16, 1: 32})  # Scalars & Vectors

    # Create Input Feature Dictionary
    >>> F = {
    ...    '0': torch.randn(10, 16, 1),   # Degree 0 features (scalars)
    ...    '1': torch.randn(10, 32, 3)    # Degree 1 features (vectors)
    ... }

    #  Define Number of Attention Heads
    >>> num_heads = 4
    >>> output = fiber2head(F, num_heads, fiber_structure, squeeze=True)
    >>> print(output.shape)
    torch.Size([10, 4, 28])

    References
    ----------
    .. [1] SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
           Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, Max Welling
           NeurIPS 2020, https://arxiv.org/abs/2006.10503
    """
    if h <= 0:
        raise ValueError("Number of attention heads (h) must be > 0")
    if not isinstance(F, dict):
        raise TypeError("F must be a dictionary of tensors")
    if not isinstance(structure, Fiber):
        raise TypeError("%s must be an instance of Fiber" % structure)

    fibers_list = []

    for degree in structure.degrees:
        key = str(degree)
        if key not in F:
            raise KeyError(f"Degree {key} missing in input dictionary F")

        tensor = F[key]

        if squeeze:
            reshaped = tensor.view(
                *tensor.shape[:-2], h,
                -1)  # [batch, nodes, heads, channels_per_head]
            fibers_list.append(reshaped)
        else:
            reshaped = tensor.view(
                *tensor.shape[:-2], h, -1,
                1)  # [batch, nodes, heads, channels_per_head, 1]
            fibers_list.append(reshaped)

    if not fibers_list:
        raise ValueError("No valid tensors found in F to concatenate")

    if squeeze:
        return torch.cat(fibers_list,
                         dim=-1)  # Concatenate along last feature dimension
    else:
        return torch.cat(fibers_list, dim=-2)


def get_basis(G,
              max_degree: int,
              compute_gradients: bool = False) -> Dict[str, torch.Tensor]:
    """
    Precompute the SE(3)-equivariant weight basis, \( W_J^{lk}(x) \).

    This function computes the equivariant weight basis used in SE(3)-equivariant
    convolutions, enabling the model to learn **rotation-equivariant filters.

    It is **called internally** by `get_basis_and_r()`, which computes **both**:
    - **Equivariant weight basis** \( W_J^{lk}(x) \).
    - **Inter-nodal distances** \( r_{ij} \) (used in attention & convolution).

    ---
    ### **Mathematical Background**
    This function follows Equation (8) from [SE(3)-Transformer paper](https://arxiv.org/pdf/2006.10503.pdf):

    \[
    W_J^{lk}(x) = Q_J \cdot Y_J(x)
    \]

    - \( Y_J(x) \): **Spherical Harmonic Basis**
    - \( Q_J \): **Basis transformation matrix**
    - \( J \): **Angular momentum index**
    - \( d_{in}, d_{out} \): **Feature degrees**

    This function precomputes basis functions for SE(3)-equivariant convolutions by calcutating spherical harmonic projections.

    Parameters
    ----------
    G:  dgl.DGLGraph
        DGL graph where `G.edata['edge_attr']` stores edge displacement vectors.
    max_degree: int
        Maximum degree (`l_max`) of equivariant tensors.
    compute_gradients: `bool`, optional, default=`False`
        If `True`: Enables gradient tracking for backpropagation.
        If `False`: Uses `torch.no_grad()` for efficiency.

    Returns
    -------
    basis: Dict[str, torch.Tensor]
        Dictionary mapping `'<d_in>,<d_out>'`: precomputed basis tensor.
        Shape: `(batch_size, 1, 2*d_out+1, 1, 2*d_in+1, num_bases)`

    Example
    -------

    >>> import torch
    >>> import dgl
    >>> import shutil
    >>> import os
    >>> from deepchem.utils.equivariance_utils import get_basis
    >>> from rdkit import Chem
    >>> import deepchem as dc
    >>> mol = Chem.MolFromSmiles('CCO')
    >>> featurizer = dc.feat.EquivariantGraphFeaturizer(fully_connected=False, embeded=True)
    >>> features = featurizer.featurize([mol])[0]
    >>> G = features.to_dgl_graph()
    >>> G.edata['w'] = torch.tensor(features.edge_weights, dtype=torch.float32)
    >>> basis = get_basis(G, max_degree=2)
    >>> print(basis.keys())
    dict_keys(['0,0', '0,1', '0,2', '1,0', '1,1', '1,2', '2,0', '2,1', '2,2'])
    >>> dir_path = "cache"
    >>> if os.path.exists(dir_path) and os.path.isdir(dir_path):
    ...     shutil.rmtree(dir_path)
    """
    # Disable gradients unless explicitly enabled
    context = torch.enable_grad() if compute_gradients else torch.no_grad()

    with context:
        cloned_d = torch.clone(G.edata['edge_attr'])
        if G.edata['edge_attr'].requires_grad:
            cloned_d.requires_grad_()

        # Compute relative positional encodings in spherical coordinates
        r_ij = get_spherical_from_cartesian(cloned_d)

        # Compute Spherical Harmonics (Y_J) for all degrees up to 2*max_degree
        Y = precompute_sh(r_ij, 2 * max_degree)
        device = Y[0].device

        basis = {}
        for d_in in range(max_degree + 1):
            for d_out in range(max_degree + 1):
                K_Js = []
                for J in range(abs(d_in - d_out), d_in + d_out + 1):
                    # Compute basis transformation matrix Q_J
                    Q_J = basis_transformation_Q_J(J, d_in,
                                                   d_out).float().T.to(device)

                    # Apply projection using spherical harmonics
                    K_J = torch.matmul(Y[J], Q_J)
                    K_Js.append(K_J)

                # Reshape for broadcasting
                size = (-1, 1, 2 * d_out + 1, 1, 2 * d_in + 1,
                        2 * min(d_in, d_out) + 1)
                basis[f"{d_in},{d_out}"] = torch.stack(K_Js, -1).view(*size)

        return basis


def get_r(G) -> torch.Tensor:
    """
    Compute inter-nodal distances for a given DGL graph.

    This function computes the Euclidean distance between connected nodes
    based on edge feature `d`, which represents relative displacements.

    Parameters
    ----------
    G : dgl.DGLGraph
        The input graph where `G.edata['edge_attr']` contains edge displacement vectors.

    Returns
    -------
    torch.Tensor
        A tensor containing the computed inter-nodal distances for each edge,
        with shape `(num_edges, 1)`.

    Example
    -------
    >>> import torch
    >>> import dgl
    >>> from deepchem.utils.equivariance_utils import get_r
    >>> from rdkit import Chem
    >>> import deepchem as dc
    >>> mol = Chem.MolFromSmiles('CCO')
    >>> featurizer = dc.feat.EquivariantGraphFeaturizer(fully_connected=False, embeded=True)
    >>> features = featurizer.featurize([mol])[0]
    >>> G = features.to_dgl_graph()
    >>> G.edata['w'] = torch.tensor(features.edge_weights, dtype=torch.float32)
    >>> # Compute internodal distances
    >>> r = get_r(G)
    >>> print(r.shape)  # (num_edges, 1)
    torch.Size([4, 1])
    """
    cloned_d = torch.clone(G.edata['edge_attr'])
    if G.edata['edge_attr'].requires_grad:
        cloned_d.requires_grad_()
    return torch.sqrt(torch.sum(cloned_d**2, -1, keepdim=True))


def get_equivariant_basis_and_r(
    G,
    max_degree: int,
    compute_gradients: bool = False
) -> Tuple[Dict[str, torch.Tensor], torch.Tensor]:
    """
    Compute equivariant weight basis and inter-nodal distances for SE(3)-Transformers.

    This function computes:
    1. **Equivariant weight basis**: Required for SE(3)-equivariant convolutions.
    2. **Inter-nodal distances**: Used in attention mechanisms and basis function computation.

    This should be **called once per forward pass** to compute shared equivariant representations.

    Parameters
    ----------
    G : dgl.DGLGraph
        A DGL graph containing edge features.
    max_degree : int
        The maximum degree of the SE(3) tensor representation.
    compute_gradients : bool, optional (default=False)
        If True, enables gradient computation for the basis.

    Returns
    -------
    Tuple[Dict[str, torch.Tensor], torch.Tensor]
        - **basis (dict[str, torch.Tensor])**: Dictionary of equivariant bases,
          indexed by `'<d_in><d_out>'`, where `d_in` and `d_out` are feature degrees.
        - **r (torch.Tensor)**: Inter-nodal distance tensor with shape `(num_edges, 1)`.

    Example
    -------
    >>> import torch
    >>> import dgl
    >>> from deepchem.utils.equivariance_utils import get_equivariant_basis_and_r
    >>> from rdkit import Chem
    >>> import deepchem as dc
    >>> mol = Chem.MolFromSmiles('CCO')
    >>> featurizer = dc.feat.EquivariantGraphFeaturizer(fully_connected=False, embeded=True)
    >>> features = featurizer.featurize([mol])[0]
    >>> G = features.to_dgl_graph()
    >>> G.edata['w'] = torch.tensor(features.edge_weights, dtype=torch.float32)
    >>> # Compute basis and distances
    >>> basis, r = get_equivariant_basis_and_r(G, max_degree=2)
    >>> print(r.shape)  # Expected: (num_edges, 1)
    torch.Size([4, 1])
    >>> print(basis.keys())  # Expected: dict
    dict_keys(['0,0', '0,1', '0,2', '1,0', '1,1', '1,2', '2,0', '2,1', '2,2'])
    """
    basis = get_basis(G, max_degree, compute_gradients)
    r = get_r(G)
    return basis, r


def semifactorial(x: int) -> float:
    """
    Compute the semifactorial function x!!, defined as:
    x!! = x * (x-2) * (x-4) * ...
    Parameters
    ----------
    x : int
        A positive integer.
    Returns
    -------
    float
        The value of x!!, computed iteratively.
    Examples
    --------
    >>> semifactorial(5)
    15.0
    >>> semifactorial(6)
    48.0
    """
    # edge cases
    if x == 0 or x == 1:
        return 1.0

    y = 1.0
    for n in range(x, 1, -2):
        y *= n
    return y


def pochhammer(x: int, k: int) -> float:
    """
    Compute the Pochhammer symbol (x)_k , defined as:
    (x)_k = x * (x+1) * (x+2) * ... * (x+k-1).
    Parameters
    ----------
    x : int
        The starting integer of the sequence.
    k : int
        The number of terms in the product.
    Returns
    -------
    float
        The Pochhammer symbol value.
    Examples
    --------
    >>> pochhammer(3, 4)
    360.0
    >>> pochhammer(5, 2)
    30.0
    """
    # handle edge case (k=0)
    if k == 0:
        return 1.0

    xf = float(x)
    for n in range(x + 1, x + k):
        xf *= n
    return xf


def lpmv(d: int, m: int, x: torch.Tensor) -> torch.Tensor:
    """
    Compute the associated Legendre function P_l^m(x), including
    the Condon-Shortley phase. The implementation includes the base case
    for P_m^m(x), P_{m+1}^m(x), and uses a recurrence relation
    for higher degrees.
    Parameters
    ----------
    l : int
        Degree of the Legendre polynomial.
    m : int
        Order of the Legendre polynomial.
    x : torch.Tensor
        A tensor of shape (N,) representing the input values for x,
        where ( -1 <= x <= 1).
    Returns
    -------
    torch.Tensor
        A tensor of shape (N,) containing the values of P_l^m(x).
    Examples
    --------
    >>> lpmv(2, 1, torch.tensor([0.5]))
    tensor([-1.2990])
    """
    m_abs = abs(m)
    if m_abs > d:
        return torch.zeros_like(x)

    # base case: P_m^m(x)
    p_mm = ((-1)**m_abs * semifactorial(2 * m_abs - 1)) * torch.pow(
        1 - x * x, m_abs / 2)

    if d == m_abs:
        return p_mm

    # Compute P_{m+1}^m(x)
    p_m1m = x * (2 * m_abs + 1) * p_mm

    if d == m_abs + 1:
        return p_m1m

    # Recurrence relation for higher degrees (l)
    p_lm = p_m1m.clone()
    for i in range(m_abs + 2, d + 1):
        p_lm = ((2 * i - 1) * x * p_m1m - (i + m_abs - 1) * p_mm) / (i - m_abs)
        p_mm, p_m1m = p_m1m, p_lm

    # Negative values of m
    if m < 0:
        p_lm *= ((-1)**m) * pochhammer(d - m + 1, -2 * m)

    return p_lm


class SphericalHarmonics:
    """
    A class for computing real tesseral spherical harmonics, including
    the Condon-Shortley phase.
    This implementation is inspired by the SE(3)-Transformer library by Fabian Fuchs,
    which provides efficient computation of spherical harmonics and related transformations.
    For more details, see the SE(3)-Transformer repository:
    https://github.com/FabianFuchsML/se3-transformer-public

    Methods
    -------
    get_element(l, m, theta, phi)
        Compute the tesseral spherical harmonic Y_l^m(theta, phi).
    get(l, theta, phi)
        Compute all spherical harmonics of degree l for given angles.
    """

    def __init__(self):
        self.leg = {}

    def clear(self) -> None:
        """Clear cached Legendre polynomial values to save RAM memory."""
        self.leg = {}

    def get_element(self, d: int, m: int, theta: torch.Tensor,
                    phi: torch.Tensor) -> torch.Tensor:
        """
        Compute a single tesseral spherical harmonic Y_l^m(theta, phi).
        Parameters
        ----------
        l : int
            Degree of the spherical harmonic.
        m : int
            Order of the spherical harmonic.
        theta : torch.Tensor
            Tensor of polar angles (collatitude) in radians.
        phi : torch.Tensor
            Tensor of azimuthal angles (longitude) in radians.
        Returns
        -------
        torch.Tensor
            Tensor of the spherical harmonic values, same shape as `theta`.
        Examples
        --------
        >>> theta = torch.tensor([0.0, math.pi / 2])
        >>> phi = torch.tensor([0.0, math.pi])
        >>> SphericalHarmonics().get_element(1, 0, theta, phi)
        tensor([ 4.8860e-01, -2.1357e-08])
        """

        N = math.sqrt((2 * d + 1) / (4 * math.pi))
        leg = lpmv(d, abs(m), torch.cos(theta))
        if m == 0:
            return N * leg
        elif m > 0:
            Y = torch.cos(m * phi) * leg
        else:
            Y = torch.sin(abs(m) * phi) * leg
        N *= math.sqrt(2.0 / pochhammer(d - abs(m) + 1, 2 * abs(m)))
        Y *= N
        return Y

    def get(self,
            d: int,
            theta: torch.Tensor,
            phi: torch.Tensor,
            refresh=True) -> torch.Tensor:
        """
        Compute all spherical harmonics of degree l.
        Parameters
        ----------
        l : int
            Degree of the spherical harmonics.
        theta : torch.Tensor
            Tensor of polar angles (collatitude) in radians.
        phi : torch.Tensor
            Tensor of azimuthal angles (longitude) in radians.
        Returns
        -------
        torch.Tensor
            A tensor of shape [*theta.shape, 2 * l + 1].
        Examples
        --------
        >>> theta = torch.tensor([0.0, math.pi / 2])
        >>> phi = torch.tensor([0.0, math.pi])
        >>> SphericalHarmonics().get(1, theta, phi)
        tensor([[-0.0000e+00,  4.8860e-01, -0.0000e+00],
                [ 4.2715e-08, -2.1357e-08,  4.8860e-01]])
        """
        if refresh:
            self.clear()
        results = []
        for m in range(-d, d + 1):
            results.append(self.get_element(d, m, theta, phi))
        return torch.stack(results, dim=-1)


def irr_repr(order: int,
             alpha: float,
             beta: float,
             gamma: float,
             dtype: Optional[torch.dtype] = None) -> torch.Tensor:
    """
    Compute the irreducible representation of the special orthogonal group SO(3).
    This function computes the Wigner D-matrix for the given order and angles (alpha, beta, gamma).
    It is compatible with composition and spherical harmonics computations.
    The output is a 2D tensor of shape `(2*order + 1, 2*order + 1)` where each element represents
    the transformation coefficients of the Wigner D-matrix. The values correspond to the rotational
    components for the given order and Euler angles.
    Parameters
    ----------
    order : int
        The order of the representation.
    alpha : float
        The first Euler angle (rotation about the Y axis).
    beta : float
        The second Euler angle (rotation about the X axis).
    gamma : float
        The third Euler angle (rotation about the Y axis).
    dtype : torch.dtype, optional
        The desired data type of the resulting tensor. If None, the default dtype is used.
    Returns
    -------
    torch.Tensor
        The irreducible representation matrix of SO(3) for the specified order and angles.
    Examples
    --------
    >>> irr_repr(1, 0.1, 0.2, 0.3)
    tensor([[ 0.9216,  0.0587,  0.3836],
            [ 0.0198,  0.9801, -0.1977],
            [-0.3875,  0.1898,  0.9021]])
    """
    result = wigner_D(order, torch.tensor(alpha), torch.tensor(beta),
                      torch.tensor(gamma))[0]
    return result.clone().detach() if dtype is None else result.clone().detach(
    ).to(dtype)


def su2_generators(k: int) -> torch.Tensor:
    """Generate the generators of the special unitary group SU(2) in a given representation.

    The function computes the generators of the SU(2) group for a specific representation
    determined by the value of 'k'. These generators are commonly used in the study of
    quantum mechanics, angular momentum, and related areas of physics and mathematics.
    The generators are represented as matrices.

    The SU(2) group is a fundamental concept in quantum mechanics and symmetry theory.
    The generators of the group, denoted as J_x, J_y, and J_z, represent the three
    components of angular momentum operators. These generators play a key role in
    describing the transformation properties of physical systems under rotations.

    The returned tensor contains three matrices corresponding to the x, y, and z generators,
    usually denoted as J_x, J_y, and J_z. These matrices form a basis for the Lie algebra
    of the SU(2) group.

    In linear algebra, specifically within the context of quantum mechanics, lowering and
    raising operators are fundamental concepts that play a crucial role in altering the
    eigenvalues of certain operators while acting on quantum states. These operators are
    often referred to collectively as "ladder operators."

    A lowering operator is an operator that, when applied to a quantum state, reduces the
    eigenvalue associated with a particular observable. In the context of SU(2), the lowering
    operator corresponds to J_-.

    Conversely, a raising operator is an operator that increases the eigenvalue of an
    observable when applied to a quantum state. In the context of SU(2), the raising operator
    corresponds to J_+.

    The z-generator matrix represents the component of angular momentum along the z-axis,
    often denoted as J_z. It commutes with both J_x and J_y and is responsible for quantizing
    the angular momentum.

    Note that the dimensions of the returned tensor will be (3, 2j+1, 2j+1), where each matrix
    has a size of (2j+1) x (2j+1).
    Parameters
    ----------
    k : int
        The representation index, which determines the order of the representation.

    Returns
    -------
    torch.Tensor
        A stack of three SU(2) generators, corresponding to J_x, J_z, and J_y.

    Notes
    -----
    A generating set of a group is a subset $S$ of the group $G$ such that every element
    of $G$ can be expressed as a combination (under the group operation) of finitely many
    elements of the subset $S$ and their inverses.

    The special unitary group $SU_n(q)$ is the set of $n*n$ unitary matrices with determinant
    +1. $SU(2)$ is homeomorphic with the orthogonal group $O_3^+(2)$. It is also called the
    unitary unimodular group and is a Lie group.

    References
    ----------
    .. [1] https://en.wikipedia.org/wiki/Ladder_operator
    .. [2] https://en.wikipedia.org/wiki/Special_unitary_group#The_group_SU(2)
    .. [3] https://en.wikipedia.org/wiki/Generating_set_of_a_group
    .. [4] https://mathworld.wolfram.com/SpecialUnitaryGroup

    Examples
    --------
    >>> su2_generators(1)
    tensor([[[ 0.0000+0.0000j,  0.7071+0.0000j,  0.0000+0.0000j],
             [-0.7071+0.0000j,  0.0000+0.0000j,  0.7071+0.0000j],
             [ 0.0000+0.0000j, -0.7071+0.0000j,  0.0000+0.0000j]],
    <BLANKLINE>
            [[-0.0000-1.0000j,  0.0000+0.0000j,  0.0000+0.0000j],
             [ 0.0000+0.0000j,  0.0000+0.0000j,  0.0000+0.0000j],
             [ 0.0000+0.0000j,  0.0000+0.0000j,  0.0000+1.0000j]],
    <BLANKLINE>
            [[ 0.0000-0.0000j,  0.0000+0.7071j,  0.0000-0.0000j],
             [ 0.0000+0.7071j,  0.0000-0.0000j,  0.0000+0.7071j],
             [ 0.0000-0.0000j,  0.0000+0.7071j,  0.0000-0.0000j]]])
    """
    # Generate the raising operator matrix
    m = torch.arange(-k, k)
    raising = torch.diag(-torch.sqrt(k * (k + 1) - m * (m + 1)), diagonal=-1)

    # Generate the lowering operator matrix
    m = torch.arange(-k + 1, k + 1)
    lowering = torch.diag(torch.sqrt(k * (k + 1) - m * (m - 1)), diagonal=1)

    # Generate the z-generator matrix
    m = torch.arange(-k, k + 1)
    z_generator = torch.diag(1j * m)

    # Combine the matrices to form the x, z, and y generators
    x_generator = 0.5 * (raising + lowering)  # x (usually)
    y_generator = -0.5j * (raising - lowering)  # -y (usually)

    # Stack the generators along the first dimension to create a tensor
    generators = torch.stack([x_generator, z_generator, y_generator], dim=0)

    return generators


def get_matrix_kernel(A: torch.Tensor, eps: float = 1e-10) -> torch.Tensor:
    """
    Compute an orthonormal basis of the kernel (x_1, x_2, ...).

    This function calculates the null space (kernel) of the given matrix A, such that:
        A x_i = 0
        scalar_product(x_i, x_j) = delta_ij

    Parameters
    ----------
    A : torch.Tensor
        The input matrix.
    eps : float, optional
        Tolerance for singular values considered as zero (default is 1e-10).

    Returns
    -------
    torch.Tensor
        A matrix where each row is a basis vector of the kernel of A.

    Examples
    --------
    >>> from deepchem.utils.equivariance_utils import get_matrix_kernel
    >>> A = torch.tensor([[1.0, 2.0, 3.0],
    ...                   [2.0, 4.0, 6.0],
    ...                   [3.0, 6.0, 9.0]])
    >>> get_matrix_kernel(A)
    tensor([[ 0.0000, -0.8321,  0.5547],
            [ 0.9636, -0.1482, -0.2224]])
    """
    _, s, v = torch.svd(A)

    kernel = v.t()[s < eps]
    return kernel


def get_matrices_kernel(As: List[torch.Tensor],
                        eps: float = 1e-10) -> torch.Tensor:
    """
    Compute the common kernel of all the input matrices.

    This function computes the shared null space of a collection of matrices.

    Parameters
    ----------
    As : List[torch.Tensor]
        List of input matrices.
    eps : float, optional
        Tolerance for singular values considered as zero (default is 1e-10).

    Returns
    -------
    torch.Tensor
        A matrix where each row is a basis vector of the common kernel.
    """
    return get_matrix_kernel(torch.cat(As, dim=0), eps)


@cached_dirpklgz("cache/trans_Q")
def basis_transformation_Q_J(
    J: int,
    order_in: int,
    order_out: int,
) -> torch.Tensor:
    """
    Compute one part of the Q^-1 matrix for the article.

    This function computes the spherical harmonics projection matrix for the
    Sylvester equation in the subspace J needed in for the weight basis in SE(3)-Transformer model.

    References:
    -----------
    - SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
    Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, Max Welling
    NeurIPS 2020, https://arxiv.org/abs/2006.10503

    Parameters
    ----------
    J : int
        Order of the spherical harmonics.
    order_in : int
        Order of the input representation.
    order_out : int
        Order of the output representation.
    version : int, optional
        Version of the computation (default is 3).

    Returns
    -------
    torch.Tensor
        A tensor of shape [(m_out * m_in), m], where m = 2 * J + 1.

    Examples
    --------
    >>> from deepchem.utils.equivariance_utils import basis_transformation_Q_J
    >>> basis_transformation_Q_J(1, 1, 1).shape
    torch.Size([9, 3])
    """
    original_dtype = torch.get_default_dtype()
    torch.set_default_dtype(torch.float64)

    def _R_tensor(a: float, b: float, c: float) -> torch.Tensor:
        """
        Compute the Kronecker product of irreducible representations.

        This function calculates the Kronecker product of two irreducible
        representations (input and output orders) for a given set of rotation angles.

        Parameters
        ----------
        a : float
            Rotation angle around the x-axis.
        b : float
            Rotation angle around the y-axis.
        c : float
            Rotation angle around the z-axis.

        Returns
        -------
        torch.Tensor
            The Kronecker product of the irreducible representations of the
            given rotation angles.

        Examples
        --------
        >>> _R_tensor(1.0, 2.0, 3.0).shape
        torch.Size([...])
        return kron(irr_repr(order_out, a, b, c), irr_repr(order_in, a, b, c))
        """
        return kron(irr_repr(order_out, a, b, c), irr_repr(order_in, a, b, c))

    def _sylvester_submatrix(J: int, a: float, b: float,
                             c: float) -> torch.Tensor:
        """
        Generate the Kronecker product matrix for solving the Sylvester equation in subspace J.

        This function constructs a Kronecker product matrix that helps solve
        the Sylvester equation for the given subspace J. The equation ensures
        the transformation is valid within the subspace.

        Parameters
        ----------
        J : int
            Order of the spherical harmonics.
        a : float
            Rotation angle around the x-axis.
        b : float
            Rotation angle around the y-axis.
        c : float
            Rotation angle around the z-axis.

        Returns
        -------
        torch.Tensor
            A rank-deficient matrix for use in solving the Sylvester equation.

        Examples
        --------
        >>> _sylvester_submatrix(1, 1.0, 2.0, 3.0).shape
        torch.Size([...])
        """
        R_tensor = _R_tensor(a, b, c)
        R_irrep_J = irr_repr(J, a, b, c)
        return kron(R_tensor, torch.eye(R_irrep_J.size(0))) - \
                   kron(torch.eye(R_tensor.size(0)), R_irrep_J.t())

    random_angles = [[4.41301023, 5.56684102, 4.59384642],
                     [4.93325116, 6.12697327, 4.14574096],
                     [0.53878964, 4.09050444, 5.36539036],
                     [2.16017393, 3.48835314, 5.55174441],
                     [2.52385107, 0.2908958, 3.90040975]]

    null_space = get_matrices_kernel(
        [_sylvester_submatrix(J, a, b, c) for a, b, c in random_angles])
    Q_J = null_space[0]
    Q_J = Q_J.view((2 * order_out + 1) * (2 * order_in + 1), 2 * J + 1)
    torch.set_default_dtype(original_dtype)

    return Q_J


def get_spherical_from_cartesian(cartesian: torch.Tensor,
                                 divide_radius_by: float = 1.0) -> torch.Tensor:
    """
    Convert Cartesian coordinates to spherical coordinates.

    Parameters
    ----------
    cartesian : torch.Tensor
        Cartesian coordinates tensor of shape [..., 3].
    divide_radius_by : float, optional
        Factor by which to divide the radius (default is 1.0).

    Returns
    -------
    torch.Tensor
        Spherical coordinates tensor of shape [..., 3] with [radius, azimuth (phi), elevation (theta)].

    Examples
    --------
    >>> from deepchem.utils.equivariance_utils import get_spherical_from_cartesian
    >>> cartesian = torch.tensor([[1.0, 1.0, 1.0]])
    >>> get_spherical_from_cartesian(cartesian)
    tensor([[1.7321, 0.7854, 0.9553]])

    >>> cartesian = torch.tensor([[0.0, 0.0, 1.0]])  # Point on Z-axis
    >>> get_spherical_from_cartesian(cartesian)
    tensor([[1.0000, 0.0000, 1.5708]])

    >>> cartesian = torch.tensor([[0.0, 0.0, -1.0]])  # Point on negative Z-axis
    >>> get_spherical_from_cartesian(cartesian)
    tensor([[1.0000, 3.1416, 1.5708]])
    """
    spherical = torch.zeros_like(cartesian)

    ind_radius = 0
    ind_alpha = 1
    ind_beta = 2

    cartesian_x = 2
    cartesian_y = 0
    cartesian_z = 1

    r_xy = cartesian[..., cartesian_x]**2 + cartesian[..., cartesian_y]**2

    spherical[..., ind_beta] = torch.atan2(torch.sqrt(r_xy),
                                           cartesian[..., cartesian_z])
    spherical[..., ind_alpha] = torch.atan2(cartesian[..., cartesian_y],
                                            cartesian[..., cartesian_x])

    if divide_radius_by == 1.0:
        spherical[..., ind_radius] = torch.sqrt(r_xy +
                                                cartesian[..., cartesian_z]**2)
    else:
        spherical[..., ind_radius] = torch.sqrt(
            r_xy + cartesian[..., cartesian_z]**2) / divide_radius_by

    return spherical


def kron(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Computes the Kronecker product of two tensors needed to comput Q_J matrix.

    References:
    -----------
    - SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
    Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, Max Welling
    NeurIPS 2020, https://arxiv.org/abs/2006.10503

    Parameters
    ----------
    a : torch.Tensor
        First input tensor of shape [*, m, n].
    b : torch.Tensor
        Second input tensor of shape [*, p, q].

    Returns
    -------
    torch.Tensor
        The Kronecker product of `a` and `b`.

    Examples
    --------
    >>> from deepchem.utils.equivariance_utils import kron
    >>> A = torch.tensor([[1, 2], [3, 4]])
    >>> B = torch.tensor([[0, 5], [6, 7]])
    >>> kron(A, B)
    tensor([[ 0,  5,  0, 10],
            [ 6,  7, 12, 14],
            [ 0, 15,  0, 20],
            [18, 21, 24, 28]])
    """
    if a.ndimension() == 2 and b.ndimension() == 2:
        return torch.einsum("ij,kl->ikjl", (a, b)).reshape(
            a.size(0) * b.size(0),
            a.size(1) * b.size(1))

    siz1 = torch.Size(torch.tensor(a.shape[-2:]) * torch.tensor(b.shape[-2:]))
    res = a.unsqueeze(-1).unsqueeze(-3) * b.unsqueeze(-2).unsqueeze(-4)
    siz0 = res.shape[:-4]
    return res.reshape(siz0 + siz1)


def precompute_sh(r_ij: torch.Tensor, max_J: int) -> dict:
    """
    Precompute spherical harmonics up to a given order used in the forward pass of
    SE(3)-Transformer model.

    References:
    -----------
    - SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
    Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, Max Welling
    NeurIPS 2020, https://arxiv.org/abs/2006.10503

    Parameters
    ----------
    r_ij : torch.Tensor
        Relative positions tensor.
    max_J : int
        Maximum order of the spherical harmonics.

    Returns
    -------
    dict
        A dictionary where each key corresponds to an order J and the value is a tensor
        of shape [B, N, K, 2J+1].

    Examples
    --------
    >>> from deepchem.utils.equivariance_utils import precompute_sh
    >>> r_ij = torch.tensor([[1.0, 0.5, 1.0]])  # Example spherical coordinates (radius, phi, theta)
    >>> precompute_sh(r_ij, max_J=2)
    {0: tensor([[0.2821]]), 1: tensor([[-0.1971, -0.2640, -0.3608]]), 2: tensor([[ 0.3255,  0.2381, -0.0392,  0.4359,  0.2090]])}
    """
    i_alpha = 1
    i_beta = 2

    Y_Js = {}
    sh = SphericalHarmonics()

    for J in range(max_J + 1):
        Y_Js[J] = sh.get(J,
                         theta=math.pi - r_ij[..., i_beta],
                         phi=r_ij[..., i_alpha],
                         refresh=False)

    sh.clear()
    return Y_Js


def change_basis_real_to_complex(
        k: int,
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None) -> torch.Tensor:
    """Construct a transformation matrix to change the basis from real to complex spherical harmonics.

    This function constructs a transformation matrix Q that converts real spherical
    harmonics into complex spherical harmonics.
    It operates on the basis functions $Y_{\ell m}$ and $Y_{\ell}^{m}$, and accounts
    for the relationship between the real and complex forms of these harmonics
    as defined in the provided mathematical expressions.

    The resulting transformation matrix Q is used to change the basis of vectors or tensors of real spherical harmonics to
    their complex counterparts.

    Parameters
    ----------
    k : int
        The representation index, which determines the order of the representation.
    dtype : torch.dtype, optional
        The data type for the output tensor. If not provided, the
        function will infer it. Default is None.
    device : torch.device, optional
        The device where the output tensor will be placed. If not provided,
        the function will use the default device. Default is None.

    Returns
    -------
    torch.Tensor
        A transformation matrix Q that changes the basis from real to complex spherical harmonics.

    Notes
    -----
    Spherical harmonics Y_l^m are a family of functions that are defined on the surface of a
    unit sphere. They are used to represent various physical and mathematical phenomena that
    exhibit spherical symmetry. The indices l and m represent the degree and order of the
    spherical harmonics, respectively.

    The conversion from real to complex spherical harmonics is achieved by applying specific
    transformation coefficients to the real-valued harmonics. These coefficients are derived
    from the properties of spherical harmonics.

    References
    ----------
    .. [1] https://en.wikipedia.org/wiki/Spherical_harmonics#Real_form

    Examples
    --------
    # The transformation matrix generated is used to change the basis of a vector of
    # real spherical harmonics with representation index 1 to complex spherical harmonics.
    >>> change_basis_real_to_complex(1)
    tensor([[-0.7071+0.0000j,  0.0000+0.0000j,  0.0000-0.7071j],
            [ 0.0000+0.0000j,  0.0000-1.0000j,  0.0000+0.0000j],
            [-0.7071+0.0000j,  0.0000+0.0000j,  0.0000+0.7071j]])
    """
    q = torch.zeros((2 * k + 1, 2 * k + 1), dtype=torch.complex128)

    # Construct the transformation matrix Q for m in range(-k, 0)
    for m in range(-k, 0):
        q[k + m, k + abs(m)] = 1 / 2**0.5
        q[k + m, k - abs(m)] = complex(-1j / 2**0.5)  # type: ignore

    # Set the diagonal elements for m = 0
    q[k, k] = 1

    # Construct the transformation matrix Q for m in range(1, k + 1)
    for m in range(1, k + 1):
        q[k + m, k + abs(m)] = (-1)**m / 2**0.5
        q[k + m, k - abs(m)] = complex(1j * (-1)**m / 2**0.5)  # type: ignore

    # Apply the factor of (-1j)**k to make the Clebsch-Gordan coefficients real
    q = (-1j)**k * q

    # Handle dtype and device options
    if dtype is None:
        default_type = torch.empty(0).dtype
        if default_type == torch.float32:
            dtype = torch.complex64
        elif default_type == torch.float64:
            dtype = torch.complex128
    if device is None:
        device = torch.empty(0).device

    # Ensure the tensor is contiguous and on the specified device
    return q.to(
        dtype=dtype,
        device=device,
        copy=True,
        memory_format=torch.contiguous_format)  # type: ignore[call-overload]


def so3_generators(k: int) -> torch.Tensor:
    """Construct the generators of the SO(3) Lie algebra for a given quantum angular momentum.

    The function generates the generators of the special orthogonal group SO(3), which represents the group
    of rotations in three-dimensional space. Its Lie algebra, which consists of the generators of
    infinitesimal rotations, is often used in physics to describe angular momentum operators.
    The generators of the Lie algebra can be related to the SU(2) group, and this function uses
    a transformation to convert the SU(2) generators to the SO(3) basis.

    The primary significance of the SO(3) group lies in its representation of three-dimensional
    rotations. Each matrix in SO(3) corresponds to a unique rotation, capturing the intricate
    ways in which objects can be oriented in 3D space. This concept finds application in
    numerous fields, ranging from physics to engineering.

    Parameters
    ----------
     k : int
        The representation index, which determines the order of the representation.

    Returns
    -------
    torch.Tensor
        A stack of three SO(3) generators, corresponding to J_x, J_z, and J_y.

    Notes
    -----
    The special orthogonal group $SO_n(q)$ is the subgroup of the elements of general orthogonal
    group $GO_n(q)$ with determinant 1. $SO_3$ (often written $SO(3)$) is the rotation group
    for three-dimensional space.

    These matrices are orthogonal, which means their rows and columns form mutually perpendicular
    unit vectors. This preservation of angles and lengths makes orthogonal matrices fundamental
    in various mathematical and practical applications.

    The "special" part of $SO(3)$ refers to the determinant of these matrices being $+1$. The
    determinant is a scalar value that indicates how much a matrix scales volumes.
    A determinant of $+1$ ensures that the matrix represents a rotation in three-dimensional
    space without involving any reflection or scaling operations that would reverse the orientation of space.

    References
    ----------
    .. [1] https://en.wikipedia.org/wiki/Special_orthogonal_group
    .. [2] https://en.wikipedia.org/wiki/3D_rotation_group#Connection_between_SO(3)_and_SU(2)
    .. [3] https://www.pas.rochester.edu/assets/pdf/undergraduate/su-2s_double_covering_of_so-3.pdf

    Examples
    --------
    >>> so3_generators(1)
    tensor([[[ 0.0000,  0.0000,  0.0000],
             [ 0.0000,  0.0000, -1.0000],
             [ 0.0000,  1.0000,  0.0000]],
    <BLANKLINE>
            [[ 0.0000,  0.0000,  1.0000],
             [ 0.0000,  0.0000,  0.0000],
             [-1.0000,  0.0000,  0.0000]],
    <BLANKLINE>
            [[ 0.0000, -1.0000,  0.0000],
             [ 1.0000,  0.0000,  0.0000],
             [ 0.0000,  0.0000,  0.0000]]])
    """
    # Get the SU(2) generators for the given quantum angular momentum (spin) value.
    X = su2_generators(k)

    # Get the transformation matrix to change the basis from real to complex spherical harmonics.
    Q = change_basis_real_to_complex(k)

    # Convert the SU(2) generators to the SO(3) basis using the transformation matrix Q.
    # X represents the SU(2) generators, and Q is the transformation matrix from real to complex spherical harmonics.
    # The resulting X matrix will be the SO(3) generators in the complex basis.
    X = torch.conj(Q.T) @ X @ Q

    # Return the real part of the SO(3) generators to ensure they are purely real.
    return torch.real(X)


def wigner_D(k: int, alpha: torch.Tensor, beta: torch.Tensor,
             gamma: torch.Tensor) -> torch.Tensor:
    """Wigner D matrix representation of the SO(3) rotation group.

    The function computes the Wigner D matrix representation of the SO(3) rotation group
    for a given representation index 'k' and rotation angles 'alpha', 'beta', and 'gamma'.
    The resulting matrix satisfies properties of the SO(3) group representation.

    Parameters
    ----------
    k : int
        The representation index, which determines the order of the representation.
    alpha : torch.Tensor
        Rotation angles (in radians) around the Y axis, applied third.
    beta : torch.Tensor
        Rotation angles (in radians) around the X axis, applied second.
    gamma : torch.Tensor)
        Rotation angles (in radians) around the Y axis, applied first.

    Returns
    -------
    torch.Tensor
        The Wigner D matrix of shape (#angles, 2k+1, 2k+1).

    Notes
    -----
    The Wigner D-matrix is a unitary matrix in an irreducible representation
    of the groups SU(2) and SO(3).

    The Wigner D-matrix is used in quantum mechanics to describe the action
    of rotations on states of particles with angular momentum. It is a key
    concept in the representation theory of the rotation group SO(3), and
    it plays a crucial role in various physical contexts.

    Examples
    --------
    >>> k = 1
    >>> alpha = torch.tensor([0.1, 0.2])
    >>> beta = torch.tensor([0.3, 0.4])
    >>> gamma = torch.tensor([0.5, 0.6])
    >>> wigner_D_matrix = wigner_D(k, alpha, beta, gamma)
    >>> wigner_D_matrix
    tensor([[[ 0.8275,  0.1417,  0.5433],
             [ 0.0295,  0.9553, -0.2940],
             [-0.5607,  0.2593,  0.7863]],
    <BLANKLINE>
            [[ 0.7056,  0.2199,  0.6737],
             [ 0.0774,  0.9211, -0.3817],
             [-0.7044,  0.3214,  0.6329]]])
    """
    # Ensure that alpha, beta, and gamma have the same shape for broadcasting.
    alpha, beta, gamma = torch.broadcast_tensors(alpha, beta, gamma)

    # Ensure the angles are within the range [0, 2*pi) using modulo.
    alpha = alpha[..., None, None] % (2 * math.pi)
    beta = beta[..., None, None] % (2 * math.pi)
    gamma = gamma[..., None, None] % (2 * math.pi)

    # Get the SO(3) generators for the given quantum angular momentum (spin) value 'k'.
    X = so3_generators(k)

    # Calculate the Wigner D matrix using the matrix exponential of the generators
    # and the rotation angles alpha, beta, and gamma in the appropriate order.
    D_matrix = torch.matrix_exp(gamma * (X[1].unsqueeze(0))) @ torch.matrix_exp(
        beta * (X[0].unsqueeze(0))) @ torch.matrix_exp(alpha *
                                                       (X[1].unsqueeze(0)))
    return D_matrix


def commutator(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    """Compute the commutator of two matrices.

    Parameters
    ----------
    A : torch.Tensor
        The first matrix.
    B : torch.Tensor
        The second matrix.

    Returns
    -------
    torch.Tensor
        The commutator of the two matrices.

    Examples
    --------
    >>> A = torch.tensor([[1, 2], [3, 4]])
    >>> B = torch.tensor([[5, 6], [7, 8]])
    >>> commutator(A, B)
    tensor([[ -4, -12],
            [ 12,   4]])
    """
    return torch.matmul(A, B) - torch.matmul(B, A)


class LieGroup:
    """Abstract base class for Lie groups used in LieConv [1] (https://github.com/mfinzi/LieConv).

    This class represents continuous Lie groups acting on
    Euclidean data, as required by LieConv-style equivariant neural networks [1].
    Specific groups, like translations (T) or 3d roto-translations (SE(3)), or subclasses
    must implement the group exponential and logarithm maps, as well
    as a lifting procedure from Euclidean coordinates to group-aligned
    representations.

    Each group is characterized by:
    - A representation space on which the group acts.
    - A Lie algebra parameterization.
    - An optional quotient (orbit) embedding.

    Subclasses must implement at least:
    - :meth:`exp` - the group exponential map
    - :meth:`log` - the group logarithm map
    - :meth:`lifted_elems` - lifting from Euclidean coordinates to Lie algebra elements

    and define the attributes `rep_dim`, `lie_dim`, and `q_dim`.

    Attributes
    ----------
    rep_dim : int
        Dimension of the representation space on which the group acts
        (e.g., 2 for SO(2) acting on R^2).
    lie_dim : int
        Dimension of the Lie algebra parameterization
        (e.g., 1 for SO(2)).
    q_dim : int
        Dimension of the quotient space embedding X/G, if applicable
        (e.g., 1 for SO(2) acting on R^2).

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """

    rep_dim: int
    lie_dim: int
    q_dim: int

    def __init__(self, alpha: float = 0.2):
        """Initialize the Lie group.

        Parameters
        ----------
        alpha : float, optional (default 0.2)
            Parameter controlling the contribution of Lie group displacement
            versus orbit (quotient-space) displacement in the distance metric.
            Larger values emphasize group motion (orientation/rotation),
            while smaller values emphasize orbit or spatial displacement.
        """
        super().__init__()
        self.alpha = alpha

    def exp(self, a: torch.Tensor) -> torch.Tensor:
        """Exponential map from the Lie algebra to the Lie group.

        Computes the group element
        $\\exp\\left(\\sum_i a_i A_i\\right)$,
        where $\\{A_i\\}$ are the generators of the Lie algebra and
        $a \\in \\mathbb{R}^{\\text{lie\\_dim}}$ are the corresponding coefficients.

        Parameters
        ----------
        a : torch.Tensor
            Lie algebra coefficients taking values in
            $\mathbb{R}^{\\text{lie\\_dim}}$.

        Returns
        -------
        torch.Tensor
            Group elements represented as matrices in
            $\mathbb{R}^{\\text{rep\\_dim} \\times \\text{rep\\_dim}}$.
        """
        raise NotImplementedError

    def log(self, u: torch.Tensor) -> torch.Tensor:
        """Logarithm map from the Lie group to the Lie algebra.

        Computes the Lie algebra element corresponding to a group element,
        expressed in the chosen Lie algebra basis.

        Parameters
        ----------
        u : torch.Tensor
            Group elements represented as matrices in
            $\mathbb{R}^{\\text{rep\\_dim} \\times \\text{rep\\_dim}}$.

        Returns
        -------
        torch.Tensor
            Lie algebra coefficients taking values in
            $\mathbb{R}^{\\text{lie\\_dim}}$.
        """
        raise NotImplementedError

    def lifted_elems(
        self,
        points: torch.Tensor,
        n_samples: int,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Lift Euclidean coordinates into Lie algebra and quotient representations.

        This method maps input points in Euclidean space to elements of the Lie
        algebra, together with optional quotient (orbit) embeddings. For groups
        with multivalued lifts, multiple samples may be generated per input point.

        Parameters
        ----------
        points : torch.Tensor
            Input points taking values in $\mathbb{R}^{\\text{rep\\_dim}}$.
        n_samples : int
            Number of lifted samples per point.

        Returns
        -------
        a : torch.Tensor
            Lie algebra elements taking values in $\mathbb{R}^{\\text{lie\\_dim}}$.
        q : torch.Tensor or None
            Quotient (orbit) embeddings taking values in
            $\mathbb{R}^{\\text{q\\_dim}}$, or ``None`` if the quotient is trivial.
        """
        raise NotImplementedError

    def inv(self, g: torch.Tensor) -> torch.Tensor:
        """Compute the inverse of a group element.

        Uses the identity $g^{-1} = \\exp(-\\log(g))$.

        Parameters
        ----------
        g : torch.Tensor
            Group elements represented as matrices in
            $\mathbb{R}^{\\text{rep\\_dim} \\times \\text{rep\\_dim}}$.

        Returns
        -------
        torch.Tensor
            Inverse group elements in
            $\mathbb{R}^{\\text{rep\\_dim} \\times \\text{rep\\_dim}}$.
        """
        return self.exp(-self.log(g))

    def distance(self, abq_pairs: torch.Tensor) -> torch.Tensor:
        """Compute the combined group and orbit distance.

        The distance between two lifted points is defined as
        $$
        \\alpha \\, \\lVert \\Delta a \\rVert
        + (1 - \\alpha) \\, \\lVert q_1 - q_2 \\rVert,
        $$
        where $\\Delta a = \\log(v^{-1} u)$ is the relative Lie algebra element and
        $q_1, q_2$ are the corresponding quotient embeddings.

        Parameters
        ----------
        abq_pairs : torch.Tensor
            Concatenated relative Lie algebra and quotient components taking values in
            $\mathbb{R}^{\\text{lie\\_dim} + 2\\,\\text{q\\_dim}}$.

        Returns
        -------
        ArrayLike
            Scalar distance values.
        """
        ab_dist = torch.norm(abq_pairs[..., :self.lie_dim], dim=-1)
        if self.q_dim is None:
            raise ValueError("q_dim must be defined to compute distance.")
        if self.q_dim > 0:
            qa = abq_pairs[..., self.lie_dim:self.lie_dim + self.q_dim]
            qb = abq_pairs[..., self.lie_dim + self.q_dim:self.lie_dim +
                           2 * self.q_dim]
            qa_qb_dist = torch.norm(qa - qb, dim=-1)
        else:
            qa_qb_dist = torch.zeros_like(ab_dist)

        return self.alpha * ab_dist + (1.0 - self.alpha) * qa_qb_dist

    def lift(
        self,
        x: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
        nsamples: int,
        **kwargs,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Lift inputs to pairwise Lie group representations.

        This method lifts Euclidean points into the Lie algebra (and optional
        quotient space), expands associated features and masks, and constructs
        all pairwise relative group elements via the group logarithm.

        Parameters
        ----------
        x : Tuple
            Tuple ``(p, v, m)`` consisting of:
            - ``p``: input points taking values in $\mathbb{R}^{\\text{rep\\_dim}}$
            - ``v``: feature values associated with each point
            - ``m``: binary mask indicating valid points
        nsamples : int
            Number of lifted samples per point.

        Returns
        -------
        embedded_locations : torch.Tensor
            Pairwise relative group elements, optionally augmented with quotient
            embeddings, taking values in
            $\mathbb{R}^{\\text{lie\\_dim} + 2\\,\\text{q\\_dim}}$.
        expanded_v : torch.Tensor
            Feature values expanded to align with pairwise group elements.
        expanded_mask : torch.Tensor
            Mask expanded to align with pairwise group elements.
        """

        p, v, m = x

        expanded_a, expanded_q = self.lifted_elems(p, nsamples, **kwargs)

        nsamples = expanded_a.shape[-2] // m.shape[-1]

        # Expand features
        expanded_v = v[...,
                       None, :].repeat((1,) * len(v.shape[:-1]) + (nsamples, 1))
        expanded_v = expanded_v.reshape(*expanded_a.shape[:-1], v.shape[-1])

        # Expand mask
        expanded_mask = m[...,
                          None].repeat((1,) * len(v.shape[:-1]) + (nsamples,))
        expanded_mask = expanded_mask.reshape(*expanded_a.shape[:-1])

        paired_a = self.elems2pairs(expanded_a)

        if expanded_q is not None:
            q_in = expanded_q.unsqueeze(-2).expand(*paired_a.shape[:-1], 1)
            q_out = expanded_q.unsqueeze(-3).expand(*paired_a.shape[:-1], 1)
            embedded_locations = torch.cat([paired_a, q_in, q_out], dim=-1)
        else:
            embedded_locations = paired_a

        return embedded_locations, expanded_v, expanded_mask

    def expand_like(
        self,
        v: torch.Tensor,
        m: torch.Tensor,
        a: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Expand features and masks to align with lifted Lie algebra elements.

        This method replicates per-point feature values and masks so that they
        correspond to the lifted Lie algebra elements produced by the group
        lifting procedure.

        Parameters
        ----------
        v : torch.Tensor
            Feature values associated with input points.
        m : torch.Tensor
            Mask indicating valid input points.
        a : torch.Tensor
            Lifted Lie algebra elements used to determine the expansion.

        Returns
        -------
        expanded_v : torch.Tensor
            Feature values expanded to align with lifted Lie algebra elements.
        expanded_mask : torch.Tensor
            Mask expanded to align with lifted Lie algebra elements.
        """

        nsamples = a.shape[-2] // m.shape[-1]

        expanded_v = v[...,
                       None, :].repeat((1,) * len(v.shape[:-1]) + (nsamples, 1))
        expanded_v = expanded_v.reshape(*a.shape[:2], v.shape[-1])

        expanded_mask = m[...,
                          None].repeat((1,) * len(v.shape[:-1]) + (nsamples,))
        expanded_mask = expanded_mask.reshape(*a.shape[:2])

        return expanded_v, expanded_mask

    def elems2pairs(self, a: torch.Tensor) -> torch.Tensor:
        """Compute pairwise relative Lie algebra elements.

        For lifted elements $a_i, a_j \in \mathbb{R}^{\text{lie\_dim}}$, this method
        computes the relative displacement
        $\log\!\left(\exp(-a_j)\exp(a_i)\right)$
        for all ordered pairs $(i, j)$.

        Parameters
        ----------
        a : torch.Tensor
            Lie algebra elements taking values in $\mathbb{R}^{\text{lie\_dim}}$.

        Returns
        -------
        pairwise_lie_algebra : torch.Tensor
            Pairwise relative Lie algebra elements.
        """

        vinv = self.exp(-a.unsqueeze(-3))
        u = self.exp(a.unsqueeze(-2))
        pairwise_lie_algebra = self.log(vinv @ u)
        return pairwise_lie_algebra


class T(LieGroup):
    """k-dimensional translation group R^k adapted from [1] (https://github.com/mfinzi/LieConv).

    This Lie group (T) represents pure translations acting on R^k.

    Note that for the translation group the exponential and logarithm maps are
    identities. Consequently, relative group elements reduce to simple subtraction.
    In particular, in two dimensions this corresponds to the trivial difference
    x_i − x_j for $x \in \mathbb{R}^2$.


    The quotient space is trivial (q_dim = 0).

    Examples
    --------
    >>> from deepchem.utils.equivariance_utils import T
    >>> G = T(k=2)
    >>> x = torch.tensor([[0.0, 0.0],
    ...                   [1.0, 2.0]])
    >>> a, q = G.lifted_elems(x, n_samples=1)
    >>> a
    tensor([[0., 0.],
            [1., 2.]])
    >>> pairs = G.elems2pairs(a)
    >>> pairs
    tensor([[[ 0.,  0.],
             [-1., -2.]],
    <BLANKLINE>
            [[ 1.,  2.],
             [ 0.,  0.]]])
    >>> ab = pairs.reshape(-1, 2)
    >>> G.distance(ab)
    tensor([0.0000, 0.4472, 0.4472, 0.0000])

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """

    def __init__(self, k: int, alpha: float = 0.2) -> None:
        """Initialize the k-dimensional translation group.

        Parameters
        ----------
        k : int
            Dimension of the translation space R^k.
        alpha : float, optional
            Weighting factor for group displacement in the distance metric.
        """
        super().__init__(alpha=alpha)
        self.rep_dim: int = k
        self.lie_dim: int = k
        self.q_dim: int = 0

    def exp(self, a: torch.Tensor) -> torch.Tensor:
        """Exponential map (identity for translations).

        Parameters
        ----------
        a : torch.Tensor
            Lie algebra coefficients taking values in $\mathbb{R}^{\text{lie\_dim}}$.

        Returns
        -------
        torch.Tensor
            Group elements of shape $\mathbb{R}^{\text{k}}$.
        """
        return a

    def log(self, g: torch.Tensor) -> torch.Tensor:
        """Logarithm map from the Lie group to the Lie algebra.

        For the translation group, this map is the identity.

        Parameters
        ----------
        g : torch.Tensor
            Group elements taking values in $\mathbb{R}^k$.

        Returns
        -------
        torch.Tensor
            Lie algebra elements taking values in $\mathbb{R}^k$.
        """
        return g

    def lifted_elems(
        self,
        points: torch.Tensor,
        n_samples: int,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Lift Euclidean points into the Lie algebra.

        For the translation group, the lift is unique since the group is Abelian
        and simply connected; therefore `nsamples` must be 1.

        Parameters
        ----------
        points : torch.Tensor
            Input points taking values in $\mathbb{R}^k$.
        n_samples : int
            Number of lifted samples per point (must be 1).

        Returns
        -------
        a : torch.Tensor
            Lie algebra elements taking values in $\mathbb{R}^k$.
        q : None
            No quotient embedding for the translation group.
        """
        if n_samples != 1:
            raise ValueError("Translation group is Abelian; nsamples must be 1")
        return points, None

    def elems2pairs(self, a: torch.Tensor) -> torch.Tensor:
        """Compute pairwise Lie algebra differences.

        For translations:
            log(exp(-b) exp(a)) = a - b

        Parameters
        ----------
        a : torch.Tensor
            Lie algebra elements of shape (*, n, k).

        Returns
        -------
        torch.Tensor
            Pairwise differences of shape (*, n, n, k).
        """
        return a.unsqueeze(-2) - a.unsqueeze(-3)


def sinc(x: torch.Tensor, tresh: float = 7e-02) -> torch.Tensor:
    r"""Compute the normalized sinc function :math:`\mathrm{sinc}(x) = \sin(x) / x`.
    This code was adapted from [1] (https://github.com/mfinzi/LieConv).

    This implementation is numerically stable near :math:`x = 0` by using a
    Taylor series expansion for small values of :math:`|x|`.

    The Taylor approximation used is:

    :math:`\mathrm{sinc}(x) \approx
    1 - \frac{x^2}{6}
    \left(1 - \frac{x^2}{20}
    \left(1 - \frac{x^2}{42}\right)\right)`

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    tresh : float, optional (default 7e-02)
        Threshold below which the Taylor expansion is used.

    Returns
    -------
    torch.Tensor
        Elementwise values of :math:`\sin(x) / x`.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.equivariance_utils import sinc
    >>> x = torch.tensor([0.0, 1e-3, 1.0])
    >>> sinc(x)
    tensor([1.0000, 1.0000, 0.8415])

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """
    use_taylor = (x.abs() < tresh)
    return torch.where(
        use_taylor,
        1 - (x**2) / 6 * (1 - (x**2) / 20 * (1 - (x**2) / 42)),
        x.sin() / x,
    )


def sincc(x: torch.Tensor, tresh: float = 7e-02) -> torch.Tensor:
    r"""Compute :math:`(1 - \mathrm{sinc}(x)) / x^2`.

    This function arises frequently in Lie group exponential map expansions,
    particularly for rotation groups such as :math:`\mathrm{SO}(3)`.
    This code was adapted from [1] (https://github.com/mfinzi/LieConv).

    Near :math:`x = 0`, the Taylor approximation is:

    :math:`\frac{1 - \mathrm{sinc}(x)}{x^2} \approx
    \frac{1}{6}
    \left(1 - \frac{x^2}{20}
    \left(1 - \frac{x^2}{42}
    \left(1 - \frac{x^2}{72}\right)\right)\right)`

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    tresh : float, optional (default 7e-02)
        Threshold below which the Taylor expansion is used.

    Returns
    -------
    torch.Tensor
        Elementwise values of :math:`(1 - \mathrm{sinc}(x)) / x^2`.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.equivariance_utils import sincc
    >>> x = torch.tensor([0.0, 1e-3, 1.0])
    >>> sincc(x)
    tensor([0.1667, 0.1667, 0.1585])

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """
    use_taylor = (x.abs() < tresh)
    return torch.where(
        use_taylor,
        1 / 6 * (1 - (x**2) / 20 * (1 - (x**2) / 42 * (1 - (x**2) / 72))),
        (x - x.sin()) / x**3,
    )


def cosc(x: torch.Tensor, tresh: float = 7e-02) -> torch.Tensor:
    r"""Compute :math:`(1 - \cos(x)) / x^2`.
    This code was adapted from [1] (https://github.com/mfinzi/LieConv).

    This function appears in second-order expansions of rotation matrices
    and Lie group Jacobians.

    Near :math:`x = 0`, the Taylor approximation is:

    :math:`\frac{1 - \cos(x)}{x^2} \approx
    \frac{1}{2}
    \left(1 - \frac{x^2}{12}
    \left(1 - \frac{x^2}{30}
    \left(1 - \frac{x^2}{56}\right)\right)\right)`

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    tresh : float, optional (default 7e-02)
        Threshold below which the Taylor expansion is used.

    Returns
    -------
    torch.Tensor
        Elementwise values of :math:`(1 - \cos(x)) / x^2`.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.equivariance_utils import cosc
    >>> x = torch.tensor([0.0, 1e-3, 1.0])
    >>> cosc(x)
    tensor([0.5000, 0.5000, 0.4597])

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """
    use_taylor = (x.abs() < tresh)
    return torch.where(
        use_taylor,
        1 / 2 * (1 - (x**2) / 12 * (1 - (x**2) / 30 * (1 - (x**2) / 56))),
        (1 - x.cos()) / x**2,
    )


def coscc(x: torch.Tensor, tresh: float = 7e-02) -> torch.Tensor:
    r"""Compute a higher-order cosine correction term used in
    :math:`\mathrm{SO}(3)` Jacobians.
    This code was adapted from [1] (https://github.com/mfinzi/LieConv).

    This function corresponds to stabilized expressions involving
    :math:`(1 - \cos(x))` and :math:`\sin(x)` that appear in inverse Jacobians
    of the exponential map.

    Near :math:`x = 0`, the Taylor approximation is:

    :math:`\mathrm{coscc}(x) \approx
    \frac{1}{12}
    \left(1 + \frac{x^2}{60}
    \left(1 + \frac{x^2}{42}
    \left(1 + \frac{x^2}{40}\right)\right)\right)`

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    tresh : float, optional (default 7e-02)
        Threshold below which the Taylor expansion is used.

    Returns
    -------
    torch.Tensor
        Elementwise correction values.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.equivariance_utils import coscc
    >>> x = torch.tensor([0.0, 1e-3, 1.0])
    >>> coscc(x)
    tensor([0.0833, 0.0833, 0.0848])

    Notes
    -----
    The full expression includes safeguards against numerical instability
    when :math:`\cos(x) \approx 1`.

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """
    use_taylor = (x.abs() < tresh)

    t_expand = 1 / 12 * (1 + (x**2) / 60 * (1 + (x**2) / 42 * (1 +
                                                               (x**2) / 40)))
    cos_term = (2 * (1 - x.cos())).clamp(min=1e-6)
    full = (1 - x * x.sin() / cos_term) / x**2

    return torch.where(use_taylor, t_expand, full)


def sinc_inv(x: torch.Tensor, tresh: float = 7e-02) -> torch.Tensor:
    r"""Compute the inverse sinc function :math:`x / \sin(x)`.
    This code was adapted from [1] (https://github.com/mfinzi/LieConv).

    Near :math:`x = 0`, the Taylor approximation is:

    :math:`\frac{x}{\sin(x)} \approx
    1 + \frac{x^2}{6} + \frac{7 x^4}{360}`

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    tresh : float, optional (default 7e-02)
        Threshold below which the Taylor expansion is used.

    Returns
    -------
    torch.Tensor
        Elementwise values of :math:`x / \sin(x)`.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.equivariance_utils import sinc_inv
    >>> x = torch.tensor([0.0, 1e-3, 1.0])
    >>> sinc_inv(x)
    tensor([1.0000, 1.0000, 1.1884])

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """
    use_taylor = (x.abs() < tresh)
    t_expand = 1 + (1 / 6) * x**2 + (7 / 360) * x**4
    return torch.where(use_taylor, t_expand, x / x.sin())


def cross_matrix(k: torch.Tensor) -> torch.Tensor:
    r"""Construct the skew-symmetric cross-product matrix for vectors in
    :math:`\mathbb{R}^3`.
    This code was adapted from [1] (https://github.com/mfinzi/LieConv).

    This corresponds to the Hodge star operator mapping
    :math:`\Lambda^1 \mathbb{R}^3 \rightarrow \Lambda^2 \mathbb{R}^3`,
    such that for any :math:`v \in \mathbb{R}^3`,
    :math:`k \times v = [k]_\times v`.

    Parameters
    ----------
    k : torch.Tensor
        Input vectors of shape :math:`(n, 3)`, where n is the dimension of k.

    Returns
    -------
    torch.Tensor
        Skew-symmetric matrices of shape :math:`(n, 3, 3)`.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.equivariance_utils import cross_matrix
    >>> k = torch.tensor([1.0, 2.0, 3.0])
    >>> cross_matrix(k)
    tensor([[ 0., -3.,  2.],
            [ 3.,  0., -1.],
            [-2.,  1.,  0.]])

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """
    K = torch.zeros(*k.shape[:-1], 3, 3, device=k.device, dtype=k.dtype)
    # Construct the atrix of the form:
    # [k]_x =
    # [[  0,   -k_z,   k_y],
    #  [ k_z,    0,   -k_x],
    #  [-k_y,  k_x,    0 ]]
    K[..., 0, 1] = -k[..., 2]
    K[..., 0, 2] = k[..., 1]
    K[..., 1, 0] = k[..., 2]
    K[..., 1, 2] = -k[..., 0]
    K[..., 2, 0] = -k[..., 1]
    K[..., 2, 1] = k[..., 0]
    return K


def uncross_matrix(K: torch.Tensor) -> torch.Tensor:
    r"""Recover a vector in :math:`\mathbb{R}^3` from a skew-symmetric matrix.

    This code was adapted from [1] (https://github.com/mfinzi/LieConv).
    This is the inverse operation of :func:`cross_matrix`, corresponding to
    the Hodge star mapping
    :math:`\Lambda^2 \mathbb{R}^3 \rightarrow \Lambda^1 \mathbb{R}^3`.

    Parameters
    ----------
    K : torch.Tensor
        Skew-symmetric matrices of shape :math:`(n, 3, 3)`, where n is the dimension of K.

    Returns
    -------
    torch.Tensor
        Vectors of shape :math:`(n, 3)`.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.equivariance_utils import uncross_matrix
    >>> k = torch.tensor([1.0, 2.0, 3.0])
    >>> uncross_matrix(cross_matrix(k))
    tensor([1., 2., 3.])

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """
    k = torch.zeros(*K.shape[:-2], 3, device=K.device, dtype=K.dtype)
    k[..., 0] = (K[..., 2, 1] - K[..., 1, 2]) / 2  # k_x
    k[..., 1] = (K[..., 0, 2] - K[..., 2, 0]) / 2  # k_y
    k[..., 2] = (K[..., 1, 0] - K[..., 0, 1]) / 2  # k_z
    return k


class SO3(LieGroup):
    r"""Special Orthogonal Group in 3D :math:`\mathrm{SO}(3)`.
    This code was adapted from [1] (https://github.com/mfinzi/LieConv).

    This class implements the Lie group of 3D rotations with exponential
    and logarithmic maps between the Lie algebra
    :math:`\mathfrak{so}(3) \cong \mathbb{R}^3`
    and rotation matrices in :math:`\mathbb{R}^{3 \times 3}`.

    Attributes
    ----------
    lie_dim : int
        Dimension of the Lie algebra, equal to 3.
    rep_dim : int
        Representation dimension, equal to 3.
    q_dim : int
        Orbit identifier dimension, equal to 1.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.equivariance_utils import SO3
    >>> G = SO3()
    >>> w = torch.tensor([0.0, 0.0, 1.0])
    >>> R = G.exp(w)
    >>> R.shape
    torch.Size([3, 3])
    >>> # Recover the Lie algebra element using the logarithmic map:
    >>> w_rec = G.log(R)
    >>> torch.allclose(w, w_rec, atol=1e-5)
    True
    >>> # Sample a batch of random rotations:
    >>> R = G.sample(8)
    >>> R.shape
    torch.Size([8, 3, 3])
    >>> # Lift 3D points into so3 using stabilizer sampling:
    >>> points = torch.randn(4, 3)
    >>> a, q = G.lifted_elems(points, n_samples=6)
    >>> a.shape
    torch.Size([24, 3])
    >>> q.shape
    torch.Size([24, 1])

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """

    lie_dim: int = 3
    rep_dim: int = 3
    q_dim: int = 1

    def __init__(self, alpha: float = 0.2):
        """Initialize the :math:`\mathrm{SO}(3)` group.

        Parameters
        ----------
        alpha : float, optional (default 0.2)
            Weighting factor for the group displacement component
            in the distance metric.
        """
        super().__init__()
        self.alpha = alpha

    def exp(self, w: torch.Tensor) -> torch.Tensor:
        r"""Exponential map :math:`\exp : \mathfrak{so}(3) \rightarrow \mathrm{SO}(3)`.

        Implements Rodrigues' rotation formula. The input vectors encode
        axis–angle rotations, where the direction specifies the rotation
        axis and the magnitude specifies the rotation angle.

        Parameters
        ----------
        w : torch.Tensor
            Lie algebra elements of shape :math:`(n, 3)`, where n is the dimension of w.

        Returns
        -------
        torch.Tensor
            Rotation matrices of shape :math:`(n, 3, 3)`.

        Examples
        --------
        >>> from deepchem.utils.equivariance_utils import SO3
        >>> import torch
        >>> G = SO3()
        >>> w = torch.tensor([0.0, 0.0, 1.0])
        >>> R = G.exp(w)
        >>> R.shape
        torch.Size([3, 3])
        """
        theta = torch.norm(w, dim=-1)[..., None, None]
        K = cross_matrix(w)
        I_m = torch.eye(3, device=K.device, dtype=K.dtype)
        return I_m + K * sinc(theta) + (K @ K) * cosc(theta)

    def log(self, R: torch.Tensor) -> torch.Tensor:
        r"""Logarithmic map :math:`\log : \mathrm{SO}(3) \rightarrow \mathfrak{so}(3)`.

        Converts rotation matrices into axis–angle coordinates.

        Parameters
        ----------
        R : torch.Tensor
            Rotation matrices of shape :math:`(n, 3, 3)`, where n is the dimension of R.

        Returns
        -------
        torch.Tensor
            Lie algebra elements of shape :math:`(n, 3)`.

        Examples
        --------
        >>> from deepchem.utils.equivariance_utils import SO3
        >>> import torch
        >>> G = SO3()
        >>> w = torch.tensor([0.3, -0.2, 0.1])
        >>> R = G.exp(w)
        >>> w_rec = G.log(R)
        >>> torch.allclose(w, w_rec, atol=1e-5)
        True
        """
        trR = R[..., 0, 0] + R[..., 1, 1] + R[
            ..., 2, 2]  # to compute the trace: tr(R)= R11​+R22​+R33
        costheta = ((trR - 1) / 2).clamp(min=-1, max=1).unsqueeze(-1)
        theta = torch.acos(costheta)
        return uncross_matrix(R) * sinc_inv(theta)

    def sample(
        self,
        *shape: int,
        device: torch.device = torch.device("cpu"),
        dtype: torch.dtype = torch.float32,
    ) -> torch.Tensor:
        """Sample random rotations from :math:`\mathrm{SO}(3)`.

        Rotations are sampled by drawing random unit quaternions and
        converting them to axis–angle form.

        Parameters
        ----------
        *shape : int
            Batch shape.
        device : torch.device, optional
            Target device.
        dtype : torch.dtype, optional
            Tensor data type.

        Returns
        -------
        torch.Tensor
            Rotation matrices of shape :math:`(*shape, 3, 3)`.

        Examples
        --------
        >>> from deepchem.utils.equivariance_utils import SO3
        >>> G = SO3()
        >>> R = G.sample(8)
        >>> R.shape
        torch.Size([8, 3, 3])
        """
        q = torch.randn(*shape, 4, device=device, dtype=dtype)
        q /= torch.norm(q, dim=-1, keepdim=True)

        theta_2 = torch.atan2(
            torch.norm(q[..., 1:], dim=-1),
            q[..., 0],
        ).unsqueeze(-1)

        so3_elem = 2 * sinc_inv(theta_2) * q[..., 1:]
        return self.exp(so3_elem)

    def lifted_elems(
        self,
        points: torch.Tensor,
        n_samples: int,
        **kwargs,
    ):
        """Lift points in :math:`\mathbb{R}^3` to :math:`\mathrm{SO}(3)` elements.

        Each point is lifted by aligning a reference axis with the direction
        of the point and composing with samples from the stabilizer subgroup.

        Parameters
        ----------
        points : torch.Tensor
            Input points of shape :math:`(*, 3)`.
        n_samples : int
            Number of lifted group elements per point.

        Returns
        -------
        a : torch.Tensor
            Lie algebra elements of shape, where n is the dimension of `points`
            :math:`(n, n_samples, 3)`.
        q : torch.Tensor
            Orbit identifiers of shape
            :math:`(n, n_samples, 1)`.

        Examples
        --------
        >>> from deepchem.utils.equivariance_utils import SO3
        >>> import torch
        >>> G = SO3()
        >>> points = torch.randn(4, 3)
        >>> a, q = G.lifted_elems(points, n_samples=6)
        >>> a.shape
        torch.Size([24, 3])
        >>> q.shape
        torch.Size([24, 1])
        """
        d = self.rep_dim
        device, dtype = points.device, points.dtype

        q = torch.randn(*points.shape[:-1],
                        n_samples,
                        4,
                        device=device,
                        dtype=dtype)
        q /= torch.norm(q, dim=-1, keepdim=True)

        theta = 2 * torch.atan2(
            torch.norm(q[..., 1:], dim=-1),
            q[..., 0],
        ).unsqueeze(-1)

        zhat = torch.zeros(*points.shape[:-1],
                           n_samples,
                           3,
                           device=device,
                           dtype=dtype)
        zhat[..., 0] = 1.0
        Rz = self.exp(zhat * theta)

        r = torch.norm(points, dim=-1, keepdim=True)
        p_hat = points / r.clamp(min=1e-5)

        w = torch.cross(
            zhat,
            p_hat[..., None, :].expand_as(zhat),
        )

        sin = torch.norm(w, dim=-1)
        cos = p_hat[..., None, 0]

        angle = torch.atan2(sin, cos).unsqueeze(-1)
        Rp = self.exp(w * sinc_inv(angle))

        A = self.log(Rp @ Rz)

        q_id = r[..., None, :].expand(*r.shape[:-1], n_samples, 1)
        flat_q = q_id.reshape(*r.shape[:-2], r.shape[-2] * n_samples, 1)
        flat_a = A.reshape(*points.shape[:-2], points.shape[-2] * n_samples, d)

        return flat_a, flat_q


class SE3(SO3):
    r"""Special Euclidean Group in 3D, :math:`\mathrm{SE}(3)`.
    This code was adapted from [1] (https://github.com/mfinzi/LieConv).

    This class represents rigid body transformations in 3D, combining
    rotations in :math:`\mathrm{SO}(3)` and translations in
    :math:`\mathbb{R}^3`.

    Elements of the Lie algebra
    :math:`\mathfrak{se}(3)` are represented as 6D vectors
    :math:`(\omega, v)`, where :math:`\omega \in \mathbb{R}^3` is the rotation vector
    and :math:`v \in \mathbb{R}^3` is the translation vector.

    The group elements are represented as homogeneous transformation
    matrices in :math:`\mathbb{R}^{4 \times 4}`.

    Attributes
    ----------
    lie_dim : int
        Dimension of the Lie algebra, equal to 6.
    rep_dim : int
        Representation dimension, equal to 4.
    q_dim : int
        Quotient dimension, equal to 0 (trivial quotient).

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.equivariance_utils import SE3
    >>> G = SE3()
    >>> w = torch.tensor([0.0, 0.0, 1.0, 1.0, 0.0, 0.0])
    >>> T = G.exp(w)
    >>> T.shape
    torch.Size([4, 4])
    >>> # Recover the Lie algebra element using the logarithmic map:
    >>> w_rec = G.log(T)
    >>> torch.allclose(w, w_rec, atol=1e-5)
    True
    >>> # Lift a batch of 3D points into se3 elements:
    >>> points = torch.randn(2, 5, 3)
    >>> a, q = G.lifted_elems(points, n_samples=4)
    >>> a.shape
    torch.Size([2, 20, 6])
    >>> q is None
    True
    >>> # Compute a weighted distance between relative transformations:
    >>> G = SE3(alpha=0.5)
    >>> rel = torch.randn(10, 6)
    >>> d = G.distance(rel)
    >>> d.shape
    torch.Size([10])

    References
    ----------
    .. [1] Generalizing Convolutional Neural Networks for Equivariance
           to Lie Groups on Arbitrary Continuous Data
           Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
           NeurIPS 2020, https://arxiv.org/abs/2002.12880
    """

    lie_dim: int = 6
    rep_dim: int = 4
    q_dim: int = 0

    def __init__(self, alpha: float = 0.2, per_point: bool = True):
        """Initialize the :math:`\mathrm{SE}(3)` group.

        Parameters
        ----------
        alpha : float, optional (default 0.2)
            Weighting factor between rotational and translational
            components in the distance metric.
        per_point : bool, optional (default True)
            If ``True``, use a distinct random group element per point
            when lifting inputs.
        """
        super().__init__()
        self.alpha = alpha
        self.per_point = per_point

    def exp(self, w: torch.Tensor) -> torch.Tensor:
        r"""Exponential map :math:`\exp : \mathfrak{se}(3) \rightarrow \mathrm{SE}(3)`.

        Converts Lie algebra elements into homogeneous transformation matrices
        using the closed-form Rodrigues expansion and the associated
        left Jacobian.

        Parameters
        ----------
        w : torch.Tensor
            Lie algebra elements of shape :math:`(n, 6)`, where the first
            three components correspond to rotation and the last three
            to translation and n is the dimension of w.

        Returns
        -------
        torch.Tensor
            Homogeneous transformation matrices of shape
            :math:`(n, 4, 4)`.

        Examples
        --------
        >>> from deepchem.utils.equivariance_utils import SE3
        >>> import torch
        >>> G = SE3()
        >>> w = torch.tensor([0.0, 0.0, 1.0, 1.0, 0.0, 0.0])
        >>> T = G.exp(w)
        >>> T.shape
        torch.Size([4, 4])
        """
        theta = torch.norm(w[..., :3], dim=-1)[..., None, None]
        K = cross_matrix(w[..., :3])
        R = super().exp(w[..., :3])

        I_m = torch.eye(3, device=w.device, dtype=w.dtype)
        V = I_m + cosc(theta) * K + sincc(theta) * (K @ K)

        U = torch.zeros(*w.shape[:-1], 4, 4, device=w.device, dtype=w.dtype)
        U[..., :3, :3] = R
        U[..., :3, 3] = (V @ w[..., 3:].unsqueeze(-1)).squeeze(-1)
        U[..., 3, 3] = 1
        return U

    def log(self, U: torch.Tensor) -> torch.Tensor:
        r"""Logarithmic map :math:`\log : \mathrm{SE}(3) \rightarrow \mathfrak{se}(3)`.

        Converts homogeneous transformation matrices into Lie algebra
        elements using the inverse left Jacobian.

        Parameters
        ----------
        U : torch.Tensor
            Homogeneous transformation matrices of shape
            :math:`(n, 4, 4)`, where n is the dimension of U.

        Returns
        -------
        torch.Tensor
            Lie algebra elements of shape :math:`(n, 6)`.

        Examples
        --------
        >>> from deepchem.utils.equivariance_utils import SE3
        >>> import torch
        >>> G = SE3()
        >>> w = torch.randn(6)
        >>> T = G.exp(w)
        >>> w_rec = G.log(T)
        >>> torch.allclose(w, w_rec, atol=1e-5)
        True
        """
        w = super().log(U[..., :3, :3])

        I_m = torch.eye(3, device=w.device, dtype=w.dtype)
        K = cross_matrix(w[..., :3])
        theta = torch.norm(w, dim=-1)[..., None, None]

        cosccc = coscc(theta)
        V_inv = I_m - K / 2 + cosccc * (K @ K)

        u = (V_inv @ U[..., :3, 3].unsqueeze(-1)).squeeze(-1)
        return torch.cat([w, u], dim=-1)

    def lifted_elems(
        self,
        points: torch.Tensor,
        n_samples: int,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        r"""Lift Euclidean points into :math:`\mathfrak{se}(3)` elements.

        This method generates random rotations and combines them with
        translations defined by the input points.

        Parameters
        ----------
        points : torch.Tensor
            Input points of shape :math:`(B, N, 3)`.
        n_samples : int
            Number of lifted samples per point.

        Returns
        -------
        a : torch.Tensor
            Lie algebra elements of shape
            :math:`(B, N \cdot \text{n_samples}, 6)`.
        q : None
            No quotient embedding for :math:`\mathrm{SE}(3)`.

        Examples
        --------
        >>> from deepchem.utils.equivariance_utils import SE3
        >>> import torch
        >>> G = SE3()
        >>> points = torch.randn(2, 5, 3)
        >>> a, q = G.lifted_elems(points, n_samples=4)
        >>> a.shape
        torch.Size([2, 20, 6])
        >>> q is None
        True
        """
        bs, n = points.shape[:2]
        device = points.device
        dtype = points.dtype

        # Sample random unit quaternions to represent rotations in SO(3).
        # When lifting a point x ∈ R^3 to SE(3), we associate it with a
        # transformation g = (R, t), where R ∈ SO(3) is a rotation and t ∈ R^3 is a translation (point location).

        # The orientation R is not determined by the point x alone, so there are
        # infinitely many valid SE(3) elements corresponding to the same point.
        # To account for all possible orientations, LieConv approximates an
        # integral over SO(3) using Monte Carlo sampling.

        # If per_point=True, each point (molecule in a batch) receives its own independent set of
        # sampled rotations (local frames). Otherwise, all points in a batch
        # share the same sampled rotations (global frame).
        if self.per_point:
            q = torch.randn(bs, n, n_samples, 4, device=device, dtype=dtype)
        else:
            q = torch.randn(bs, 1, n_samples, 4, device=device, dtype=dtype)

        # Normalize 4D Gaussian vectors to obtain random unit quaternions on S^3,
        # which parameterize rotations in SO(3).
        # This corresponds to sampling rotations according to the Haar measure to not
        # depend on any coordinate system.

        # These samples are used to approximate integrals of the form:
        #   ∫_{SO(3)} f(R) dμ_Haar(R)
        # via Monte Carlo:
        #   ≈ (1 / n_samples) ∑_i f(R(q_i))
        q /= torch.norm(q, dim=-1, keepdim=True)

        # Convert unit quaternions to axis–angle (Lie algebra) representation.
        # Each unit quaternion q = (w, v) corresponds to a rotation with
        # angle θ and axis v / ||v||. This step maps quaternions into elements
        # of so(3).
        theta_2 = torch.atan2(
            torch.norm(q[..., 1:],
                       dim=-1),  # slicing vectorial component -> û*sin(theta/2)
            q[..., 0],  # slicing scalar component -> cos(theta/2)
        ).unsqueeze(-1)

        # Form se(3) elements with rotation part from so(3) and zero translation.
        # At this stage, we have pure rotational Lie algebra elements:
        #   (ω, 0) ∈ se(3)
        so3_elem = 2 * sinc_inv(theta_2) * q[
            ..., 1:]  # slicing vectorial component -> û*sin(theta/2)
        se3_elem = torch.cat(
            [so3_elem, torch.zeros_like(so3_elem)],
            dim=-1,
        )

        # Map the sampled Lie algebra elements to SE(3) via the exponential map.
        # This produces rigid transformations with rotation R and zero translation.
        R = self.exp(se3_elem)

        # Construct pure translation transforms:
        # Each point x ∈ R^3 is represented as an SE(3) element with:
        #   - identity rotation
        #   - translation equal to x
        T = torch.zeros(bs,
                        n,
                        n_samples,
                        4,
                        4,
                        device=points.device,
                        dtype=points.dtype)
        T[..., :, :] = torch.eye(
            4, device=points.device, dtype=points.dtype
        )  # Identity rotation and zero translation in remaining matrix components.
        T[..., :3,
          3] = points[:, :,
                      None, :]  # Fill the translation column of each SE(3) matrix with the point coordinates.

        # Compose translations with sampled rotations and map back to se(3).
        a = self.log(T @ R)

        # Flatten the samples into shape (B, N * n_samples, 6).
        return a.reshape(bs, n * n_samples, 6), None

    def distance(self, abq_pairs: torch.Tensor) -> torch.Tensor:
        r"""Compute a weighted distance in :math:`\mathfrak{se}(3)`.

        The distance is defined as:
        :math:`\alpha |\omega| + (1 - \alpha) |v|`,

        where :math:`\omega` is the rotational component and
        :math:`v` is the translational component.

        Parameters
        ----------
        abq_pairs : torch.Tensor
            Relative Lie algebra elements of shape :math:`(n, 6)`, where n is the dimension of `abq_pairs`.

        Returns
        -------
        torch.Tensor
            Scalar distance values of shape :math:`(n)`.

        Examples
        --------
        >>> from deepchem.utils.equivariance_utils import SE3
        >>> G = SE3(alpha=0.5)
        >>> d = G.distance(torch.randn(10, 6))
        >>> d.shape
        torch.Size([10])
        """
        dist_rot = torch.norm(abq_pairs[..., :3], dim=-1)
        dist_trans = torch.norm(abq_pairs[..., 3:], dim=-1)
        return self.alpha * dist_rot + (1.0 - self.alpha) * dist_trans
