import numpy as np
import numpy.testing as npt
import scipy.signal
from lenstronomy.ImSim.image_model import ImageModel
from lenstronomy.Data.imaging_data import ImageData
from lenstronomy.Data.psf import PSF
from lenstronomy.LightModel.light_model import LightModel
from lenstronomy.LensModel.lens_model import LensModel
import lenstronomy.Util.simulation_util as sim_util
from lenstronomy.Util import kernel_util
import lenstronomy.Util.util as util

from lenstronomy.ImSim.image_linear_solve import ImageLinearFit

"""
Test the linear solver for natwt (natural weighting) interferometric data.
Test the _image_linear_solve function of ImageLinearFit class.
The idea is to define data, psf, source, lens, lens light classes respectively, and run the linear solving
inside and outside of the _image_linear_solve function. Verify the 1st and 4th output of _image_linear_solve.
The test should be independent of the specific definitions of the light and lens profiles.
"""


def test_image_linear_solve_with_primary_beam_and_interferometry_psf():
    background_rms = 0.05
    exp_time = np.inf
    numPix = 80
    deltaPix = 0.05
    psf_type = "PIXEL"
    kernel_size = 161

    # simulate a primary beam (pb)
    primary_beam = np.zeros((numPix, numPix))
    for i in range(numPix):
        for j in range(numPix):
            primary_beam[i, j] = np.exp(-1e-4 * ((i - 78) ** 2 + (j - 56) ** 2))
    primary_beam /= np.max(primary_beam)

    # simulate a spherical sinc function as psf, which contains negative pixels
    psf_test = np.zeros((221, 221))
    for i in range(221):
        for j in range(221):
            if i > j:
                psf_test[i, j] = psf_test[j, i]
            r = np.sqrt((i - 110) ** 2 + (j - 110) ** 2)
            if r == 0:
                psf_test[i, j] = 1
            else:
                psf_test[i, j] = np.sin(r * 0.5) / (r * 0.5)

    # note that the simulated noise here is not the interferometric noise. we just use it to test the numerics
    test_noise = scipy.signal.fftconvolve(
        np.random.normal(0, 1, (numPix, numPix)), psf_test, mode="same"
    )

    kwargs_data = sim_util.data_configure_simple(
        numPix, deltaPix, exp_time, background_rms
    )
    kwargs_data["ra_at_xy_0"] = -(40) * deltaPix
    kwargs_data["dec_at_xy_0"] = -(40) * deltaPix
    kwargs_data["antenna_primary_beam"] = primary_beam
    kwargs_data[
        "likelihood_method"
    ] = "interferometry_natwt"  # testing just for interferometry natwt method
    data_class = ImageData(**kwargs_data)

    kernel_cut = kernel_util.cut_psf(psf_test, kernel_size, normalisation=False)
    kwargs_psf = {
        "psf_type": psf_type,
        "pixel_size": deltaPix,
        "kernel_point_source": kernel_cut,
        "kernel_point_source_normalisation": False,
    }
    psf_class = PSF(**kwargs_psf)

    # define lens model and source model
    kwargs_shear = {"gamma1": 0.01, "gamma2": 0.01}
    kwargs_spemd = {
        "theta_E": 1.0,
        "gamma": 1.8,
        "center_x": 0,
        "center_y": 0,
        "e1": 0.1,
        "e2": 0.04,
    }
    lens_model_list = ["SPEP", "SHEAR"]
    kwargs_lens = [kwargs_spemd, kwargs_shear]
    lens_model_class = LensModel(lens_model_list=lens_model_list)

    kwargs_sersic = {
        "amp": 25.0,
        "R_sersic": 0.3,
        "n_sersic": 2,
        "center_x": 0,
        "center_y": 0,
    }
    lens_light_model_list = ["SERSIC"]
    kwargs_lens_light = [kwargs_sersic]
    lens_light_model_class = LightModel(light_model_list=lens_light_model_list)

    kwargs_sersic_ellipse = {
        "amp": 10.0,
        "R_sersic": 0.6,
        "n_sersic": 7,
        "center_x": 0,
        "center_y": 0,
        "e1": 0.05,
        "e2": 0.02,
    }
    source_model_list = ["SERSIC_ELLIPSE"]
    kwargs_source = [kwargs_sersic_ellipse]
    source_model_class = LightModel(light_model_list=source_model_list)

    kwargs_numerics = {"supersampling_factor": 1, "supersampling_convolution": False}

    imageModel = ImageModel(
        data_class,
        psf_class,
        lens_model_class,
        source_model_class,
        lens_light_model_class,
        kwargs_numerics=kwargs_numerics,
    )
    image_sim = imageModel.image(kwargs_lens, kwargs_source, kwargs_lens_light)

    # normalize the noise to make it small compared to the model image
    test_noise *= 1e-2 * (np.max(image_sim) / np.std(test_noise))
    sim_data = image_sim + test_noise
    data_class.update_data(sim_data)

    # define the ImageLinearFit class using the materials defined above, run the _image_linear_solve function
    imageLinearFit = ImageLinearFit(
        data_class,
        psf_class,
        lens_model_class,
        source_model_class,
        lens_light_model_class,
        kwargs_numerics=kwargs_numerics,
    )
    model, _, _, amps = imageLinearFit._image_linear_solve(
        kwargs_lens, kwargs_source, kwargs_lens_light
    )

    # execute the same linear solving outside of the _image_linear_solve function
    A = imageLinearFit._linear_response_matrix(
        kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps=None, unconvolved=True
    )
    A0 = util.array2image(A[0])
    A1 = util.array2image(A[1])
    A0c = scipy.signal.fftconvolve(A0, psf_test, mode="same")
    A1c = scipy.signal.fftconvolve(A1, psf_test, mode="same")
    M = np.zeros((2, 2))
    b = np.zeros((2))
    M[0, 0] = np.sum(A0c * A0)
    M[0, 1] = np.sum(A0c * A1)
    M[1, 0] = np.sum(A1c * A0)
    M[1, 1] = np.sum(A1c * A1)
    b[0] = np.sum(A0 * sim_data)
    b[1] = np.sum(A1 * sim_data)

    amps0 = np.linalg.lstsq(M, b)[0]
    clean_model = amps0[0] * A0 + amps0[1] * A1
    dirty_model = amps0[0] * A0c + amps0[1] * A1c

    npt.assert_almost_equal([clean_model, dirty_model], model, decimal=8)
    npt.assert_almost_equal(amps0, amps, decimal=8)
