import torch


def curve_energy(c, model, eval_pts):
    """Computes curve energy (in ambient/embedding space) with
    Riemann sums.
    
    params:
        c:              geoml.curve.CubicSpline object - the curve in latent space
        model:          nn.Module object - the VAE containing the decoder mu/sigma
                        functions
        eval_pts:       int - the number of (ordered) discrete points representing 
                        the curve
    """
    c = c.view(-1, model.latent_dim)
    mu = model.dummy_pmu(c, False)
    mu = mu.view(-1, eval_pts, model.in_dim)
    delta_mu = (mu[:, 1:, :] - mu[:, :-1, :])
    sigma = model.p_sigma(c, False)
    sigma = sigma.view(-1, eval_pts, model.in_dim)
    delta_sig = (sigma[:, :-1, :] - sigma[:, 1:, :])

    d_mu = delta_mu.pow(2).sum(1)
    d_sig = delta_sig.pow(2).sum(1)
    
    return 0.5 * torch.sum(d_mu + d_sig, dim=-1)