import torch


def dihedral(
    x0: torch.Tensor, x1: torch.Tensor, x2: torch.Tensor, x3: torch.Tensor
) -> torch.Tensor:
    """Dihedral between four points.

    Reference
    ---------
    Closely follows implementation in Yutong Zhao's timemachine:
        https://github.com/proteneer/timemachine/blob/1a0ab45e605dc1e28c44ea90f38cb0dedce5c4db/timemachine/potentials/bonded.py#L152-L199
    """
    # check input shapes

    assert x0.shape == x1.shape == x2.shape == x3.shape

    # compute displacements 0->1, 2->1, 2->3
    r01 = x1 - x0 + torch.randn_like(x0) * 1e-5
    r21 = x1 - x2 + torch.randn_like(x0) * 1e-5
    r23 = x3 - x2 + torch.randn_like(x0) * 1e-5

    # compute normal planes
    n1 = torch.cross(r01, r21)
    n2 = torch.cross(r21, r23)

    rkj_normed = r21 / torch.norm(r21, dim=-1, keepdim=True)

    y = torch.sum(torch.mul(torch.cross(n1, n2), rkj_normed), dim=-1)
    x = torch.sum(torch.mul(n1, n2), dim=-1)

    # choose quadrant correctly
    theta = torch.atan2(y, x)

    return theta