import torch


def directional_derivatives(lin_op, directions, device):
    """``lin_op`` represents a curvature matrix (either ``GGNLinearOperator`` or
    ``HessianLinearOperator`` from ``vivit.hessianfree``). ``directions`` is a
    ``D x nof_directions`` matrix, where ``nof_directions`` directions are stored
    column-wise. For every direction ``d``, we compute ``d.T @ lin_op_matrix @ d``.
    """

    nof_directions = directions.shape[1]

    derivs = torch.zeros(nof_directions)
    for d_idx in range(nof_directions):

        # Pick direction and convert to numpy
        d = directions[:, d_idx]
        d_np = torch.clone(d).cpu().numpy()

        # Compute directional derivative
        mat_vec_product = torch.from_numpy(lin_op.matvec(d_np)).to(device)
        derivs[d_idx] = torch.inner(mat_vec_product, d).item()

    return derivs