import numpy as np

import autoarray as aa


def test__set_w_tilde():

    # fit inversion is None, so no need to bother with w_tilde.

    fit_0 = aa.m.MockFitImaging(inversion=None)
    fit_1 = aa.m.MockFitImaging(inversion=None)

    preloads = aa.Preloads(w_tilde=1, use_w_tilde=1)
    preloads.set_w_tilde_imaging(fit_0=fit_0, fit_1=fit_1)

    assert preloads.w_tilde is None
    assert preloads.use_w_tilde is False

    # Noise maps of fit are different but there is an inversion, so we should not preload w_tilde and use w_tilde.

    inversion = aa.m.MockInversion(
        leq=aa.m.MockLEq(linear_obj_list=[aa.m.MockMapper()])
    )

    fit_0 = aa.m.MockFitImaging(
        inversion=inversion,
        noise_map=aa.Array2D.zeros(shape_native=(3, 1), pixel_scales=0.1),
    )
    fit_1 = aa.m.MockFitImaging(
        inversion=inversion,
        noise_map=aa.Array2D.ones(shape_native=(3, 1), pixel_scales=0.1),
    )

    preloads = aa.Preloads(w_tilde=1, use_w_tilde=1)
    preloads.set_w_tilde_imaging(fit_0=fit_0, fit_1=fit_1)

    assert preloads.w_tilde is None
    assert preloads.use_w_tilde is False

    # Noise maps of fits are the same so preload w_tilde and use it.

    noise_map = aa.Array2D.ones(shape_native=(5, 5), pixel_scales=0.1, sub_size=1)

    mask = aa.m.MockMask(
        native_index_for_slim_index=noise_map.mask.native_index_for_slim_index
    )

    dataset = aa.m.MockDataset(psf=aa.Kernel2D.no_blur(pixel_scales=1.0), mask=mask)

    fit_0 = aa.m.MockFitImaging(
        inversion=inversion, dataset=dataset, noise_map=noise_map
    )
    fit_1 = aa.m.MockFitImaging(
        inversion=inversion, dataset=dataset, noise_map=noise_map
    )

    preloads = aa.Preloads(w_tilde=1, use_w_tilde=1)
    preloads.set_w_tilde_imaging(fit_0=fit_0, fit_1=fit_1)

    curvature_preload, indexes, lengths = aa.util.leq.w_tilde_curvature_preload_imaging_from(
        noise_map_native=fit_0.noise_map.native,
        kernel_native=fit_0.dataset.psf.native,
        native_index_for_slim_index=fit_0.dataset.mask.native_index_for_slim_index,
    )

    assert preloads.w_tilde.curvature_preload[0] == curvature_preload[0]
    assert preloads.w_tilde.indexes[0] == indexes[0]
    assert preloads.w_tilde.lengths[0] == lengths[0]
    assert preloads.w_tilde.noise_map_value == 1.0
    assert preloads.use_w_tilde == True


def test__set_relocated_grid():

    # LEq is None so there is no mapper, thus preload mapper to None.

    fit_0 = aa.m.MockFitImaging(inversion=None)
    fit_1 = aa.m.MockFitImaging(inversion=None)

    preloads = aa.Preloads(relocated_grid=1)
    preloads.set_relocated_grid(fit_0=fit_0, fit_1=fit_1)

    assert preloads.relocated_grid is None

    # Mapper's mapping matrices are different, thus preload mapper to None.

    leq_0 = aa.m.MockLEq(
        linear_obj_list=[aa.m.MockMapper(source_grid_slim=np.ones((3, 2)))]
    )
    leq_1 = aa.m.MockLEq(
        linear_obj_list=[aa.m.MockMapper(source_grid_slim=2.0 * np.ones((3, 2)))]
    )

    fit_0 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_0))
    fit_1 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_1))

    preloads = aa.Preloads(relocated_grid=1)
    preloads.set_relocated_grid(fit_0=fit_0, fit_1=fit_1)

    assert preloads.relocated_grid is None

    # Mapper's mapping matrices are the same, thus preload mapper.

    leq_0 = aa.m.MockLEq(
        linear_obj_list=[aa.m.MockMapper(source_grid_slim=np.ones((3, 2)))]
    )
    leq_1 = aa.m.MockLEq(
        linear_obj_list=[aa.m.MockMapper(source_grid_slim=np.ones((3, 2)))]
    )

    fit_0 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_0))
    fit_1 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_1))

    preloads = aa.Preloads(relocated_grid=1)
    preloads.set_relocated_grid(fit_0=fit_0, fit_1=fit_1)

    assert (preloads.relocated_grid == np.ones((3, 2))).all()


def test__set_mapper_list():

    # LEq is None so there is no mapper, thus preload mapper to None.

    fit_0 = aa.m.MockFitImaging(inversion=None)
    fit_1 = aa.m.MockFitImaging(inversion=None)

    preloads = aa.Preloads(mapper_list=1)
    preloads.set_mapper_list(fit_0=fit_0, fit_1=fit_1)

    assert preloads.mapper_list is None

    # Mapper's mapping matrices are different, thus preload mapper to None.

    leq_0 = aa.m.MockLEq(
        linear_obj_list=[aa.m.MockMapper(mapping_matrix=np.ones((3, 2)))]
    )
    leq_1 = aa.m.MockLEq(
        linear_obj_list=[aa.m.MockMapper(mapping_matrix=2.0 * np.ones((3, 2)))]
    )

    fit_0 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_0))
    fit_1 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_1))

    preloads = aa.Preloads(mapper_list=1)
    preloads.set_mapper_list(fit_0=fit_0, fit_1=fit_1)

    assert preloads.mapper_list is None

    # Mapper's mapping matrices are the same, thus preload mapper.

    leq_0 = aa.m.MockLEq(
        linear_obj_list=[aa.m.MockMapper(mapping_matrix=np.ones((3, 2)))]
    )
    leq_1 = aa.m.MockLEq(
        linear_obj_list=[aa.m.MockMapper(mapping_matrix=np.ones((3, 2)))]
    )

    fit_0 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_0))
    fit_1 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_1))

    preloads = aa.Preloads(mapper_list=1)
    preloads.set_mapper_list(fit_0=fit_0, fit_1=fit_1)

    assert (preloads.mapper_list[0].mapping_matrix == np.ones((3, 2))).all()

    # Multiple mappers pre inversion still preloads full mapper list.

    leq_0 = aa.m.MockLEq(
        linear_obj_list=[
            aa.m.MockMapper(mapping_matrix=np.ones((3, 2))),
            aa.m.MockMapper(mapping_matrix=np.ones((3, 2))),
        ]
    )
    leq_1 = aa.m.MockLEq(
        linear_obj_list=[
            aa.m.MockMapper(mapping_matrix=np.ones((3, 2))),
            aa.m.MockMapper(mapping_matrix=np.ones((3, 2))),
        ]
    )

    fit_0 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_0))
    fit_1 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_1))

    preloads = aa.Preloads(mapper_list=1)
    preloads.set_mapper_list(fit_0=fit_0, fit_1=fit_1)

    assert (preloads.mapper_list[0].mapping_matrix == np.ones((3, 2))).all()
    assert (preloads.mapper_list[1].mapping_matrix == np.ones((3, 2))).all()


def test__set_operated_mapping_matrix_with_preloads():

    curvature_matrix_preload = np.array([[1.0]])
    curvature_matrix_counts = np.array([1.0])

    # LEq is None thus preload it to None.

    fit_0 = aa.m.MockFitImaging(inversion=None)
    fit_1 = aa.m.MockFitImaging(inversion=None)

    preloads = aa.Preloads(
        operated_mapping_matrix=1,
        curvature_matrix_preload=np.array([[1.0]]),
        curvature_matrix_counts=np.array([1.0]),
    )
    preloads.set_operated_mapping_matrix_with_preloads(fit_0=fit_0, fit_1=fit_1)

    assert preloads.operated_mapping_matrix is None
    assert preloads.curvature_matrix_preload is None
    assert preloads.curvature_matrix_counts is None

    # LEq's blurred mapping matrices are different thus no preloading.

    blurred_mapping_matrix_0 = np.array(
        [[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]
    )

    blurred_mapping_matrix_1 = np.array(
        [[0.0, 0.0, 1.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]
    )

    leq_0 = aa.m.MockLEqImaging(blurred_mapping_matrix=blurred_mapping_matrix_0)
    leq_1 = aa.m.MockLEqImaging(blurred_mapping_matrix=blurred_mapping_matrix_1)

    fit_0 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_0))
    fit_1 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_1))

    preloads = aa.Preloads(
        operated_mapping_matrix=1,
        curvature_matrix_preload=curvature_matrix_preload,
        curvature_matrix_counts=curvature_matrix_counts,
    )
    preloads.set_operated_mapping_matrix_with_preloads(fit_0=fit_0, fit_1=fit_1)

    assert preloads.operated_mapping_matrix is None
    assert preloads.curvature_matrix_preload is None
    assert preloads.curvature_matrix_counts is None

    # LEq's blurred mapping matrices are the same therefore preload it and the curvature sparse terms.

    leq_0 = aa.m.MockLEqImaging(blurred_mapping_matrix=blurred_mapping_matrix_0)
    leq_1 = aa.m.MockLEqImaging(blurred_mapping_matrix=blurred_mapping_matrix_0)

    inversion_0 = aa.m.MockInversion(
        leq=leq_0,
        curvature_matrix_preload=curvature_matrix_preload,
        curvature_matrix_counts=curvature_matrix_counts,
    )

    fit_0 = aa.m.MockFitImaging(inversion=inversion_0)
    fit_1 = aa.m.MockFitImaging(inversion=aa.m.MockInversion(leq=leq_1))

    preloads = aa.Preloads(
        operated_mapping_matrix=1,
        curvature_matrix_preload=curvature_matrix_preload,
        curvature_matrix_counts=curvature_matrix_counts,
    )
    preloads.set_operated_mapping_matrix_with_preloads(fit_0=fit_0, fit_1=fit_1)

    assert (preloads.operated_mapping_matrix == blurred_mapping_matrix_0).all()
    assert (
        preloads.curvature_matrix_preload == curvature_matrix_preload.astype("int")
    ).all()
    assert (
        preloads.curvature_matrix_counts == curvature_matrix_counts.astype("int")
    ).all()


def test__set_curvature_matrix():

    # Inversion is None thus preload curvature_matrix to None.

    fit_0 = aa.m.MockFitImaging(inversion=None)
    fit_1 = aa.m.MockFitImaging(inversion=None)

    preloads = aa.Preloads(curvature_matrix=1)
    preloads.set_curvature_matrix(fit_0=fit_0, fit_1=fit_1)

    assert preloads.curvature_matrix is None

    # Inversion's curvature matrices are different thus no preloading.

    curvature_matrix_0 = np.array([[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])

    curvature_matrix_1 = np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])

    fit_0 = aa.m.MockFitImaging(
        inversion=aa.m.MockInversion(curvature_matrix=curvature_matrix_0)
    )
    fit_1 = aa.m.MockFitImaging(
        inversion=aa.m.MockInversion(curvature_matrix=curvature_matrix_1)
    )

    preloads = aa.Preloads(curvature_matrix=1)
    preloads.set_curvature_matrix(fit_0=fit_0, fit_1=fit_1)

    assert preloads.curvature_matrix is None

    # LEq's curvature matrices are the same therefore preload it and the curvature sparse terms.

    preloads = aa.Preloads(curvature_matrix=2)

    fit_0 = aa.m.MockFitImaging(
        inversion=aa.m.MockInversion(curvature_matrix=curvature_matrix_0)
    )
    fit_1 = aa.m.MockFitImaging(
        inversion=aa.m.MockInversion(curvature_matrix=curvature_matrix_0)
    )

    preloads.set_curvature_matrix(fit_0=fit_0, fit_1=fit_1)

    assert (preloads.curvature_matrix == curvature_matrix_0).all()


def test__set_regularization_matrix_and_term():

    regularization = aa.m.MockRegularization(regularization_matrix=np.eye(2))

    # Inversion is None thus preload log_det_regularization_matrix_term to None.

    fit_0 = aa.m.MockFitImaging(inversion=None)
    fit_1 = aa.m.MockFitImaging(inversion=None)

    preloads = aa.Preloads(log_det_regularization_matrix_term=1)
    preloads.set_regularization_matrix_and_term(fit_0=fit_0, fit_1=fit_1)

    assert preloads.regularization_matrix is None
    assert preloads.log_det_regularization_matrix_term is None

    # Inversion's log_det_regularization_matrix_term are different thus no preloading.

    leq = aa.m.MockLEq(linear_obj_list=[None])

    fit_0 = aa.m.MockFitImaging(
        inversion=aa.m.MockInversion(
            log_det_regularization_matrix_term=0,
            regularization_list=[regularization],
            leq=leq,
        )
    )
    fit_1 = aa.m.MockFitImaging(
        inversion=aa.m.MockInversion(
            log_det_regularization_matrix_term=1,
            regularization_list=[regularization],
            leq=leq,
        )
    )

    preloads = aa.Preloads(log_det_regularization_matrix_term=1)
    preloads.set_regularization_matrix_and_term(fit_0=fit_0, fit_1=fit_1)

    assert preloads.regularization_matrix is None
    assert preloads.log_det_regularization_matrix_term is None

    # LEq's blurred mapping matrices are the same therefore preload it and the curvature sparse terms.

    preloads = aa.Preloads(log_det_regularization_matrix_term=2)

    fit_0 = aa.m.MockFitImaging(
        inversion=aa.m.MockInversion(
            leq=aa.m.MockLEq(linear_obj_list=[aa.m.MockMapper()]),
            log_det_regularization_matrix_term=1,
            regularization_list=[regularization],
            preloads=preloads,
        )
    )
    fit_1 = aa.m.MockFitImaging(
        inversion=aa.m.MockInversion(
            leq=aa.m.MockLEq(linear_obj_list=[aa.m.MockMapper()]),
            log_det_regularization_matrix_term=1,
            regularization_list=[regularization],
            preloads=preloads,
        )
    )

    preloads.set_regularization_matrix_and_term(fit_0=fit_0, fit_1=fit_1)

    assert (preloads.regularization_matrix == np.eye(2)).all()
    assert preloads.log_det_regularization_matrix_term == 1
