import torch


def run_mat_interp_back(kdat, coef_mat_real, coef_mat_imag):
    """Interpolates kdat to on-grid coordinates with input sparse matrices.

    Args:
        kdat (tensor): The off-grid frequency data.
        coef_mat_real (sparse tensor): The real interpolation coefficients
            stored as a sparse tensor.
        coef_mat_imag (sparse tensor): The imaginary interpolation coefficients
            stored as a sparse tensor.
        griddat (tensor): A tensor to store the outputs in.

    Returns:
        tensor: kdat interpolated to on-grid locations.
    """
    # we have to do these transposes because torch.mm requires first to be spmatrix
    real_kdat = kdat.select(1, 0).t().view(-1, kdat.shape[0])
    imag_kdat = kdat.select(1, 1).t().view(-1, kdat.shape[0])
    coef_mat_real = coef_mat_real.t_()
    coef_mat_imag = coef_mat_imag.t_()

    # apply multiplies with complex conjugate
    griddat = []

    griddat.append(
        (torch.mm(coef_mat_real, real_kdat) + torch.mm(coef_mat_imag, imag_kdat)).t()
    )

    griddat.append(
        (torch.mm(coef_mat_real, imag_kdat) - torch.mm(coef_mat_imag, real_kdat)).t()
    )

    griddat = torch.stack(griddat, dim=1)

    # put the matrices back in the order we were given
    coef_mat_real = coef_mat_real.t_()
    coef_mat_imag = coef_mat_imag.t_()

    return griddat