import torch


def catmull_rom_spline(params: torch.Tensor, t: torch.Tensor, alpha: float = 0.5) -> torch.Tensor:
    """Evaluate a batch of Catmull-Rom Splines at different points.

    Args:
      params: torch.Tensor of size [batch, 4, D]. Each batch item is interpreted as four points corresponding to
              the points along the spline. The curve itself is bounded by the second and third points; the first and
              last only define the directions. D is the number of dimensions.
      t: torch.Tensor of size [batch, nvalues, 1] or [nvalues, 1] that defines the values of t to evaluate the curve at.
      alpha: number between 0 and 1 defining the type of spline knot parameterisation. A value of 0 is the standard
             uniform Catmull-Rom spline; 0.5 is centripetal; and, 1.0 is the chordal spline. Defaults to the centripetal
             spline.

    Returns:
      a torch.Tensor of size [batch, nvalues, 2] representing the points along each curve in the batch at the
      respective values of t.

    """

    if t.ndim == 2:
        t = t.expand(params.shape[0], *t.shape)

    p0 = params[:, 0].unsqueeze(1)
    p1 = params[:, 1].unsqueeze(1)
    p2 = params[:, 2].unsqueeze(1)
    p3 = params[:, 3].unsqueeze(1)

    # Premultiplied power constant for the following tj() function.
    alpha = alpha / 2

    def tj(ti, pi, pj):
        return ((pi - pj) ** 2).sum(dim=-1) ** alpha + ti

    t0 = torch.zeros(params.shape[0], 1, device=params.device)
    t1 = tj(t0, p0, p1)
    t2 = tj(t1, p1, p2)
    t3 = tj(t2, p2, p3)

    t0 = t0.unsqueeze(1)
    t1 = t1.unsqueeze(1)
    t2 = t2.unsqueeze(1)
    t3 = t3.unsqueeze(1)

    t = t1 + t * (t2 - t1)  # t is mapped to be between t1 and t2

    a1 = (t1 - t) / (t1 - t0) * p0 + (t - t0) / (t1 - t0) * p1
    a2 = (t2 - t) / (t2 - t1) * p1 + (t - t1) / (t2 - t1) * p2
    a3 = (t3 - t) / (t3 - t2) * p2 + (t - t2) / (t3 - t2) * p3

    b1 = (t2 - t) / (t2 - t0) * a1 + (t - t0) / (t2 - t0) * a2
    b2 = (t3 - t) / (t3 - t1) * a2 + (t - t1) / (t3 - t1) * a3

    return (t2 - t) / (t2 - t1) * b1 + (t - t1) / (t2 - t1) * b2