import numpy as np
import numpy.testing as npt
from lenstronomy.Data.imaging_data import ImageData


def test_interferometry_likelihood():
    """Test interferometry natural weighting likelihood function, test likelihood_method
    function output."""

    test_data = np.zeros((5, 5))
    test_data[0, :] = 1
    test_data[:, 4] = 2

    mask = np.ones((5, 5))
    model_unconvolved = np.zeros((5, 5))
    model_convolved = np.zeros((5, 5))
    model_unconvolved[0, 4] = 1.2
    model_convolved[0, :] = 1
    model_convolved[:, 4] = 1

    data_class = ImageData(
        image_data=test_data,
        background_rms=2.5,
        log_likelihood_constant=-1.0,
        likelihood_method="interferometry_natwt",
    )

    assert data_class.likelihood_method() == "interferometry_natwt"
    npt.assert_almost_equal(
        data_class.log_likelihood([model_unconvolved, model_convolved], mask),
        -0.712,
        decimal=8,
    )
