import torch


def compute_jacobian_on_map(x, y, forward_transform, eps=0.01):
    """
    Computes the differentials dx/du, dy/du, dx/dv, dy/dv for the given projection function) using central differences. For spherical projections, (lon, lat = u, v)

    Returns (x.shape[0], x.shape[1], 2, 2)
    """

    # Compute dx/dv, dy/dv
    x0, y0 = forward_transform(x - eps, y)
    x1, y1 = forward_transform(x + eps, y)
    dx_du = (x1 - x0) / (2 * eps)
    dy_du = (y1 - y0) / (2 * eps)

    # Compute dx/du, dy/du
    x2, y2 = forward_transform(x, y - eps)
    x3, y3 = forward_transform(x, y + eps)
    dx_dv = (x3 - x2) / (2 * eps)
    dy_dv = (y3 - y2) / (2 * eps)

    return torch.stack((torch.stack(
        (dx_du, dy_du), -1), torch.stack((dx_dv, dy_dv), -1)), -1)