import os
import sys
import time
import copy

import numpy as np

import matplotlib
import matplotlib.pyplot as plt

from scipy import (
    sparse, interpolate, optimize
)

from itertools import compress

import pylops

import autolens as al


# NOTE:
sys.path.append(
    "{}/utils".format(os.environ["GitHub"])
)
import pylops_utils as pylops_utils

import autolens_utils.autolens_plot_utils as autolens_plot_utils


def visibilities_from_cube(cube, transformers, shape):

    # NOTE: Check that the shape of the uv_wavelengths in the transformers is the
    # same as the input shape.

    class Image:
        def __init__(self, array_2d):
            self.array_2d = array_2d

        @property
        def in_1d_binned(self):
            return np.ndarray.flatten(self.array_2d)

        @property
        def in_2d_binned(self):
            return self.array_2d

    if len(transformers) == shape[0]:

        visibilities = np.zeros(shape=shape)
        for i in range(shape[0]):
            visibilities[i] = transformers[i].visibilities_from_image(
                image=Image(array_2d=cube[i])
            )

    else:
        raise ValueError("shape mismatch")

    return visibilities


def visibilities_from_image_and_uv_wavelengths(image, uv_wavelengths, grid, transformer_class):

    # NOTE: Image must be a class

    if len(uv_wavelengths.shape) == 2:

        transformer = transformer_class(
            uv_wavelengths=uv_wavelengths,
            grid=grid.in_radians,
        )

        visibilities = transformer.visibilities_from_image(
            image=image
        )

    if len(uv_wavelengths.shape) == 3:

        transformers = [
            transformer_class(
                uv_wavelengths=uv_wavelengths[i],
                grid=grid.in_radians,
            )
            for i in range(uv_wavelengths.shape[0])
        ]

        visibilities = np.zeros(
            shape=uv_wavelengths.shape
        )
        for i in range(uv_wavelengths.shape[0]):
            visibilities[i, ...] = transformers[i].visibilities_from_image(
                image=image
            )

    if len(uv_wavelengths.shape) == 4:

        raise ValueError("This has not been implemented yet")

    return visibilities















# NOTE: This works regardless of how many galaxies are used in the tracer
# NOTE: The spped of this function can not improve by means of saving the weights of the interpolation
def ray_trace(traced_grids_of_planes, image, grid, interpolator=None):

    traced_grid_i = traced_grids_of_planes[0] # This is the regular grid.
    traced_grid_j = traced_grids_of_planes[1]

    # plt.figure()
    # plt.plot(traced_grid_i[:, 0], traced_grid_i[:, 1], linestyle="None", marker="o", color="b")
    # plt.plot(traced_grid_j[:, 0], traced_grid_j[:, 1], linestyle="None", marker="o", color="r")
    # plt.show()
    # exit()

    x_interp = np.unique(traced_grid_i[:, 0])
    y_interp = np.unique(traced_grid_i[:, 1])
    # # NOTE: MADE A CHANGE HERE!
    # x_interp = np.unique(grid[:, 0])
    # y_interp = np.unique(grid[:, 1])

    # plt.figure()
    # plt.plot(x_interp, y_interp, linestyle="None", marker="o")
    # plt.show()
    # exit()

    #print(image.shape);exit()

    #start = time.time()
    image_interp = interpolate.RegularGridInterpolator(
        points=(y_interp, x_interp),
        values=image[::-1, :],
        method="linear",
        bounds_error=False,
        fill_value=0.0
    )

    lensed_image = image_interp(traced_grid_j).reshape(image.shape)
    # # NOTE: MADE A CHANGE HERE!
    # lensed_image = image_interp(traced_grid_j).reshape(grid.shape_2d)
    #end = time.time()
    #print(end - start)

    # NOTE: check autoarray/grids.py -> interpolator (I am not sure if it's going to be faster)
    # lensed_image = interpolate.griddata(
    #     points=traced_grid_i,
    #     values=np.ndarray.flatten(image[:, :]),
    #     xi=traced_grid_j,
    #     method="linear",
    # )
    # lensed_image = lensed_image.reshape(image.shape)

    return lensed_image





# def tracer(lens_galaxy, source_redshift):
#
#     # NOTE: Do I need to check anything about the lens_galaxy?
#
#     tracer = al.Tracer.from_galaxies(
#         galaxies=[
#             lens_galaxy,
#             al.Galaxy(
#                 redshift=source_redshift,
#                 light=al.lp.LightProfile()
#             )
#         ]
#     )
#
#     return tracer

"""
# NOTE: Rename to get_lensed_image_from_image_and_tracer
def func_image(lens_galaxy, source_redshift, grid, image):

    tracer = al.Tracer.from_galaxies(
        galaxies=[
            lens_galaxy,
            al.Galaxy(
                redshift=source_redshift,
                light=al.lp.LightProfile()
            )
        ]
    )

    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
        grid=grid,
        plane_index_limit=tracer.upper_plane_index_with_light_profile
    )

    lensed_image = ray_trace(
        traced_grids_of_planes=traced_grids_of_planes,
        image=image
    )

    return lensed_image


def func_cube(lens_galaxy, source_redshift, grid, cube, interpolator=None):

    tracer = al.Tracer.from_galaxies(
        galaxies=[
            lens_galaxy,
            al.Galaxy(
                redshift=source_redshift,
                light=al.lp.LightProfile()
            )
        ]
    )

    return lensed_cube_from_tracer(
        tracer=tracer,
        grid=grid,
        cube=cube,
        interpolator=interpolator
    )
"""


def lensed_image_from_tracer(tracer, grid, image, interpolator=None):

    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
        grid=grid,
        plane_index_limit=tracer.upper_plane_index_with_light_profile
    )
    # # NOTE: MADE A CHANGE HERE!
    # traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
    #     grid=grid.in_1d_binned,
    #     plane_index_limit=tracer.upper_plane_index_with_light_profile
    # )

    # plt.figure()
    # plt.plot(
    #     traced_grids_of_planes[0][:, 0],
    #     traced_grids_of_planes[0][:, 1],
    #     linestyle="None",
    #     marker="o",
    #     markersize=2,
    #     color="b"
    # )
    # plt.plot(
    #     traced_grids_of_planes[1][:, 0],
    #     traced_grids_of_planes[1][:, 1],
    #     linestyle="None",
    #     marker="o",
    #     markersize=2,
    #     color="r"
    # )
    # # plt.plot(
    # #     grid[:, 0],
    # #     grid[:, 1],
    # #     linestyle="None",
    # #     marker="o",
    # #     markersize=1,
    # #     color="black"
    # # )
    # plt.show()
    # exit()

    # print(
    #     traced_grids_of_planes[0].shape,
    #     traced_grids_of_planes[1].shape,
    #     grid.shape,
    #     grid.shape_2d
    # )
    # exit()

    lensed_image = ray_trace(
        traced_grids_of_planes=traced_grids_of_planes,
        image=image,
        grid=grid
    )

    return lensed_image


def lensed_cube_from_tracer(tracer, grid, cube, z_mask=None, interpolator=None):

    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
        grid=grid,
        plane_index_limit=tracer.upper_plane_index_with_light_profile
    )

    # start = time.time()
    lensed_cube = np.zeros(
        shape=cube.shape
    )
    for i in range(lensed_cube.shape[0]):

        lensed_cube[i] = ray_trace(
            traced_grids_of_planes=traced_grids_of_planes,
            image=cube[i],
            grid=grid
        )
    # end = time.time()
    # print("Total time:", end - start)
    # exit()

    return lensed_cube


# def lensed_cube_from_galaxies(galaxies, grid, cube, interpolator=None):
#     pass


def mapper_from_tracer_and_grid(tracer, grid):

    # if not isinstance(grid, al.MaskedGrid):
    #     raise ValueError("This has not been implemented yet.")

    mappers_of_planes = tracer.mappers_of_planes_from_grid(
        grid=grid,
    )

    return mappers_of_planes[-1]

def test__regularization_matrix(tracer, grid, pixelization, ):
    pass


def transformed_mapping_matrices_from_transformer_grid_and_tracer(transformer, grid, tracer):

    mapper = mapper_from_tracer_and_grid(
        tracer=tracer,
        grid=grid
    )

    return transformer.transformed_mapping_matrices_from_mapping_matrix(
        mapping_matrix=mapper.mapping_matrix
    )


# TODO: A LOT TO-DO
def get_parameter_from_tracer(tracer, parameter_name="centre"):

    for mass_profile in tracer.mass_profiles:

        if mass_profile.__class__.__name__ in [
            "EllipticalPowerLaw",
            "EllipticalIsothermal"
        ]:

            parameter = getattr(
                mass_profile,
                parameter_name
            )

    return parameter


def galaxies_with_mass_profiles_from_tracer(tracer):

    return list(
        compress(
            tracer.galaxies,
            list(map(lambda galaxy: galaxy.has_mass_profile, tracer.galaxies)))
    )

def galaxies_with_light_profiles_from_tracer(tracer):

    return list(
        compress(
            tracer.galaxies,
            list(map(lambda galaxy: galaxy.has_light_profile, tracer.galaxies)))
    )



def galaxies_have_mass_profiles(galaxies):

    return all(list(map(lambda galaxy: galaxy.has_mass_profile, galaxies)))


def tracer_from_len_galaxies_and_src_galaxy_with_pixelization(len_galaxies, src_galaxy):

    if not isinstance(src_galaxy, al.Galaxy):
        raise ValueError(
            "The input must be an al.Galaxy class"
        )

    if src_galaxy.has_pixelization:
        if galaxies_have_mass_profiles(
            galaxies=len_galaxies
        ):
            galaxies = len_galaxies + [src_galaxy, ]
        else:
            raise ValueError(
                "Not all galaxies have mass profiles"
            )

            # NOTE: Diagnostics
    else:
        raise ValueError(
            "The \'src_galaxy\' does not have a pixelization"
        )

    return al.Tracer.from_galaxies(
        galaxies=galaxies
    )


def tracer_from_instance(instance):

    if hasattr(instance, "galaxies"):
        tracer = al.Tracer.from_galaxies(
            galaxies=instance.galaxies
        )
    else:
        raise AttributeError("...")

    return tracer


def tracer_derivative(tracer, transformer, grid, constants=None, epsilon=1e-8, reshaped=True):

    # if tracer.has_pixelization:
    #     raise ValueError(
    #         "This has not been implemented yet"
    #     )
    # else:
    #     raise ValueError(
    #         "In development ..."
    #     )

    # TODO:
    #
    # 1) Return a list of all the model parameters for each galaxy
    #
    # 2) Iterate through this list and adjust each parameter by 1e-8
    #
    # 3) For each adjustment, define a new tracer and produce visibilities. This
    # will be one column of the derivative matrix

    # derivative = np.zeros(
    #     shape=(np.prod(transformer.uv_wavelengths.shape), )
    # )


    # NOTE: For analytic profiles ...
    def galaxy_profiles(galaxy):

        def is_analytic_profile(obj):
            return isinstance(
                obj, (al.mp.MassProfile, al.lp.LightProfile)
            )

        return [key for key, value in galaxy.__dict__.items()
            if is_analytic_profile(value)
        ]



    galaxies = copy.deepcopy(tracer.galaxies)

    if constants is None:
        #raise ValueError("...")
        constants = {
            "mass":[
                "core_radius"
            ],
            "bulge":[],
            "light":[],
        }

        # constants = {
        #     "mass":[
        #         "core_radius"
        #     ],
        #     "bulge":[],
        #     "light":["intensity", "radius_break", "alpha", "gamma"],
        # }

    # def perturb(galaxies, index, ):
    #     pass
    #
    # def perturb_value(value):
    #
    #     pass

    n = 0

    perturbed_tracers = []
    for i, galaxy in enumerate(tracer.galaxies):

        for profile in galaxy_profiles(galaxy=galaxy):
            #print(profile)

            # if isinstance(profile, al.mp.SphericalNFWMCRLudlow):
            #     print("OK")

            for j, (name, value) in enumerate(
                getattr(galaxy, profile).__dict__.items()
            ):
                print(name, value)

                # NOTE: Check
                if name in constants[profile]:
                    pass
                    # print(
                    #     "The parameter \'{}\' is kept constant".format(name)
                    # )

                elif name in ["centre", "elliptical_comps"]:

                    for k, value_k in enumerate(value):

                        perturbed_galaxies = copy.deepcopy(galaxies)

                        if isinstance(value, tuple):
                            value_as_list = list(value)
                        else:
                            raise ValueError("This has not been implemented")

                        value_as_list[k] += epsilon

                        setattr(
                            getattr(perturbed_galaxies[i], profile), name, tuple(value_as_list)
                        )

                        perturbed_tracer = al.Tracer.from_galaxies(
                            galaxies=perturbed_galaxies
                        )
                        perturbed_tracers.append(perturbed_tracer)

                        n += 1
                else:

                    perturbed_galaxies = copy.deepcopy(galaxies)

                    setattr(
                        getattr(perturbed_galaxies[i], profile), name, value + epsilon
                    )

                    perturbed_tracer = al.Tracer.from_galaxies(
                        galaxies=perturbed_galaxies
                    )
                    perturbed_tracers.append(perturbed_tracer)

                    n += 1

    #exit()

    # 0 centre
    # 2 axis_ratio
    # 3 phi
    # 4 einstein_radius
    # 5 slope
    # 6 core_radius
    # The parameter 'core_radius' is kept constant
    # 6 centre
    # 8 axis_ratio
    # 9 phi
    # 10 intensity
    # The parameter 'intensity' is kept constant
    # 10 effective_radius
    # 11 sersic_index

    visibilities = tracer.profile_visibilities_from_grid_and_transformer(
        grid=grid, transformer=transformer
    )

    derivative = np.zeros(
        shape=((len(perturbed_tracers), ) + transformer.uv_wavelengths.shape)
    )
    for i in range(derivative.shape[0]):
        visibilities_i = perturbed_tracers[i].profile_visibilities_from_grid_and_transformer(
            grid=grid, transformer=transformer
        )

        derivative[i] = np.divide(
            np.subtract(visibilities_i, visibilities),
            epsilon
        )

    # NOTE: Plot
    # plt.figure()
    # plt.imshow(derivative[:, :, 0], aspect="auto", norm=matplotlib.colors.SymLogNorm(10**-5.0))
    # plt.show()
    # exit()


    if reshaped:
        derivative = np.concatenate(
            (derivative[..., 0], derivative[..., 1]),
            axis=-1
        )

    return derivative


def gradients(a, grid):

    gradient_y = np.gradient(a.in_2d, grid.in_2d[:, 0, 0], axis=0)
    gradient_x = np.gradient(a.in_2d, grid.in_2d[0, :, 1], axis=1)

    return (gradient_y, gradient_x)


def deflections_via_potential_from_grid(potential, grid):

    return gradients(a=potential, grid=grid)


def gradients_from_galaxy(galaxy, grid):
    raise ValueError(
        "This has not been implemented yet."
    )

    # if galaxy.has_light_profile:
    #     pass
    # else:
    #     pass

def replace_light_profiles_in_tracer_with_pixelization(tracer, pixelization, regularization=None):

    galaxies_with_mass_profiles = galaxies_with_mass_profiles_from_tracer(
        tracer=tracer
    )

    planes_with_light_profile = tracer.planes_with_light_profile

    if not planes_with_light_profile:
        raise ValueError("...")
    elif len(planes_with_light_profile) == 1:
        redshift = planes_with_light_profile[0].redshift
    elif len(planes_with_light_profile) > 1:
        raise ValueError(
            "This has not been implemented yet"
        )
    else:
        raise ValueError("...")

    return al.Tracer.from_galaxies(
        galaxies=[
            *galaxies_with_mass_profiles_from_tracer(
                tracer=tracer
            ),
            al.Galaxy(
                redshift=redshift,
                pixelization=pixelization,
                regularization=regularization if regularization is not None else al.reg.Constant(
                    coefficient=1.0
                ),
            )
        ]
    )


def plot_mapping_matrix(mapping_matrix, grid, nrows, ncols, figsize=None):


    def array_wrapper(array, grid):

        # NOTE: This has changed from the 0.45.0 version.
        return grid.mask.mapping.array_stored_2d_from_array_1d(
            array_1d=array
        )

    n = mapping_matrix.shape[-1]

    if nrows * ncols < n:
        raise ValueError("The total number of panels {} is less than {}".format(
                nrows * ncols, n
            )
        )

    figure, axes = plt.subplots(
        nrows=nrows,
        ncols=ncols,
        figsize=figsize if figsize is not None else (16, 8)
    )

    k = 0
    for i in range(nrows):
        for j in range(ncols):

            if k < n:
                axes[i, j].imshow(
                    array_wrapper(array=mapping_matrix[:, k], grid=grid)
                )
            else:
                axes[i, j].axis("off")

            axes[i, j].set_xticks([])
            axes[i, j].set_yticks([])

            k += 1

    plt.subplots_adjust(wspace=0.0, hspace=0.0)
    plt.show()

def mappers_of_planes_from_grid(tracer, source_redshift, mask, pixelization):
    raise ValueError(
        "This has not been implemented yet."
    )
    # tracer = al.Tracer.from_galaxies(
    #     galaxies=[
    #         *galaxies_with_mass_profiles_from_tracer(
    #             tracer=tracer
    #         ),
    #         al.Galaxy(redshift=source_redshift)
    #     ]
    # )
    #
    # grid = al.MaskedGrid.from_mask(mask=mask)
    #
    # # plt.figure()
    # # plt.plot(grid[:, 0], grid[:, 1], linestyle="None", marker="o")
    # # plt.show()
    # # exit()
    #
    # mappers_of_planes = []
    #
    # traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
    #     grid=grid
    # )
    #
    # traced_sparse_grids_of_planes = tracer.traced_sparse_grids_of_planes_from_grid(
    #     grid=grid,
    #     preload_sparse_grids_of_planes=None
    # )
    #
    # for (plane_index, plane) in enumerate(tracer.planes):
    #     pass
    #     print(plane.galaxies)
    #     # print(plane_index)
    #     # if not plane.has_pixelization:
    #     #     mappers_of_planes.append(None)
    #     # # else:
    #     # #     mapper = plane.mapper_from_grid_and_sparse_grid(
    #     # #         grid=traced_grids_of_planes[plane_index],
    #     # #         sparse_grid=traced_sparse_grids_of_planes[plane_index],
    #     # #         inversion_uses_border=inversion_uses_border,
    #     # #     )
    #     # #     mappers_of_planes.append(mapper)
    #
    # #return mappers_of_planes


def planes_with_light_profiles(tracer):

    """

    1) Find all planes with light profiles.

    2) Get the galaxies of these planes.

    3) Get the light profiles of those galaxies.

    4) Make images of all those light profiles.

    """
    # NOTE: Find all planes with light profiles
    # NOTE:
    # # image = tracer.galaxies[1].profile_image_from_grid(grid=grid)
    # # plt.figure()
    # # plt.imshow(image.in_2d)
    # # plt.show()
    #
    # # asd = list(map(lambda plane: plane.has_light_profile, tracer.planes))
    # # print(asd)
    # #print(tracer.planes)
    #
    # #print(tracer.has_light_profile)
    # #print(list(map(lambda plane: plane.has_light_profile, tracer.planes)))
    # #print(tracer.galaxies_with_light_profile)
    #
    # #print(tracer.planes[1].galaxies_with_light_profile)
    #
    # galaxies = tracer.planes[1].galaxies_with_light_profile
    # galaxy = galaxies[0]
    #
    # galaxy_light_profiles = galaxy.light_profiles
    #
    # image_0 = galaxy_light_profiles[0].profile_image_from_grid(grid=grid)
    # image_0_in_2d = image_0.in_2d
    #
    # image_1 = galaxy_light_profiles[1].profile_image_from_grid(grid=grid)
    # image_1_in_2d = image_1.in_2d


# al.Galaxy(
#     redshift=src_redshift,
#     pixelization=pixelization,
#     regularization=al.reg.Constant(
#         coefficient=1000000.0
#     ),
# )


class inversion_settings:

    def __init__(self, method="operators"):

        self.method = method


class inversion_abstract:

    def __init__(self, Op, reconstruction, P, transformer, tracer, mapper, regularization_matrix, shape):

        self.Op = Op

        self._reconstruction = reconstruction

        self.P = P

        self.tracer = tracer
        self.mapper = mapper
        self.transformer = transformer

        self.regularization_matrix = regularization_matrix

        self.shape = shape


    @property
    def reconstruction(self):
        return self._reconstruction

    @reconstruction.setter
    def reconstruction(self, x):
         self._reconstruction = x


    def reshape(self, x):

        return x.reshape(self.shape)

    @property
    def mapped_reconstructed_image(self):

        array_1d = self.mapper.mapping_matrix @ self.reconstruction

        return self.mapper.grid.mapping.array_stored_1d_from_array_1d(
            array_1d=array_1d
        )

    @property
    def mapped_reconstructed_visibilities(self):

        return self.transformer.visibilities_from_image(
            image=self.mapped_reconstructed_image.in_2d
        )

        # return self.reshape(
        #     x=self.transformer._matvec(x=self.mapped_reconstructed_image)
        # )

    # @property
    # def model_data(self):
    #
    #     # return self.reshape(
    #     #     x=self.Op._matvec(x=self.reconstruction),
    #     # )
    #
    #     return self.reshape(
    #         x=self.Op * self.reconstruction,
    #     )

    @property
    def model_data(self):

        return self.mapped_reconstructed_visibilities


    def log_det(self, x):

        return np.linalg.slogdet(x)[1]


    @property
    def regularization_term(self):
        return np.matmul(
            self.reconstruction.T,
            np.matmul(self.regularization_matrix, self.reconstruction),
        )

    @property
    def log_det_curvature_reg_matrix_term(self):
        return self.log_det(x=self.P)

    @property
    def log_det_regularization_matrix_term(self):
        return self.log_det(x=self.regularization_matrix)



def wrap_mapping_matrix(mapping_matrix):

    return pylops.MatrixMult(
        sparse.bsr_matrix(mapper.mapping_matrix), dtype="float64"
    )

# def adsd(Op_1, Op_2):
#
#     return Op_1 * Op_2

def preconditioner_from(mapping_matrix, weights, regularization_matrix):
    pass


class AbstractMaskedDatasetSettings:
    def __init__(self):
        pass

class AbstractMaskedDataset:

    def __init__(self, data_1d, weights_1d):

        self.data_1d = data_1d
        self.weights_1d = weights_1d

def AbstractMaskedDatasetConverter(cls):

    # NOTE: Make it a dictionary wh
    attributes = ["data_1d", "weights_1d"]
    # attributes = {
    #     "data_1d":[], "weights_1d":[]
    # }
    for attr in attributes:
        if hasattr(cls, attr):
            pass
        else:
            if attr == "data_1":
                _attributes = []



class inversion_settings:

    def __init__(self, method="operators", use_preconditioner=True, preconditioner_method="powell"):

        self.method = method
        self.use_preconditioner = use_preconditioner
        self.preconditioner_method = preconditioner_method


def inversion_from_(
    masked_dataset,
    tracer,
    grid,
    transformer,
    transformer_from_averaged=None,
    inversion_uses_border=True,
    preload_sparse_grids_of_planes=None,
    settings=inversion_settings()
):


    def reconstruction_from_(masked_dataset, mapper, regularization_matrix, transformer, settings, transformer_from_averaged=None):

        mapping_matrix_as_operator = pylops.MatrixMult(
            sparse.bsr_matrix(mapper.mapping_matrix), dtype="float64"
        )

        Op = transformer * mapping_matrix_as_operator

        # Op = Op.toreal(
        #     forw=False, adj=True
        # )

        #print(settings.use_preconditioner, settings.preconditioner_method);exit()
        if settings.use_preconditioner:

            if settings.preconditioner_method == "approx":

                if transformer_from_averaged is None:
                    raise ValueError("This is a requirement for this method")

                start = time.time()
                transformed_mapping_matrices = transformer_from_averaged.transformed_mapping_matrices_from_mapping_matrix(
                    mapping_matrix=mapper.mapping_matrix
                )
                end = time.time()
                print(
                    "It took t={} to compute thetransformed mapping matrices".format(end - start)
                )

                F_real = np.dot(
                    transformed_mapping_matrices[0].T * masked_dataset.weights_averaged_from_reshaped[:, 0],
                    transformed_mapping_matrices[0]
                )
                F_imag = np.dot(
                    transformed_mapping_matrices[1].T * masked_dataset.weights_averaged_from_reshaped[:, 1],
                    transformed_mapping_matrices[1]
                )

                F_approx = F_real + F_imag

            else:

                F_approx = np.multiply(
                    np.sum(masked_dataset.weights_1d),
                    mapper.mapping_matrix.T @ mapper.mapping_matrix
                )


            P = np.add(F_approx, regularization_matrix)

            P_inv = np.linalg.inv(P)

            # # NOTE: converting the inverse of P to a diagonal matrix
            # P_inv = np.diag(np.diag(P_inv))

            # plt.figure()
            # plt.imshow(P_inv)
            # plt.show()
            # exit()

            # NOTE: The inverse of the P matrix is not sparse, but rather has very low value off-diagonal terms.
            MOp = pylops.MatrixMult(
                sparse.bsr_matrix(P_inv), dtype="float64"
            )

        else:

            MOp = None

        # values = pylops_utils.NormalEquationsInversion(
        #     Op=Op,
        #     data=masked_dataset.data_1d,
        #     WOp=pylops.Diagonal(
        #         diag=masked_dataset.weights_1d, dtype="float64"
        #     ),
        #     HOp=pylops.MatrixMult(
        #         sparse.bsr_matrix(regularization_matrix), dtype="float64"
        #     ),
        #     M=MOp,
        #     tol=1e-08,
        #     atol=1e-08,
        #     **dict(
        #         maxiter=mapper.mapping_matrix.shape[-1]
        #     )
        # )

        print(
            "Performing inversion ..."
        )

        start = time.time()
        values = pylops.NormalEquationsInversion(
            Op=Op,
            Regs=None,
            data=masked_dataset.data_1d,
            Weight=pylops.Diagonal(
                diag=masked_dataset.weights_1d, dtype="float64"
            ),
            epsNRs=[1.0],
            NRegs=[
                pylops.MatrixMult(
                    sparse.bsr_matrix(regularization_matrix), dtype="float64"
                )
            ],
            M=MOp,
            tol=1e-08
        )
        end = time.time()
        print(
            "It took t={} to perform the reconstruction".format(end - start)
        )

        # plt.figure()
        # autolens_plot_utils.draw_voronoi_pixels(mapper=mapper, values=values)
        # plt.show()
        # exit()

        return values, Op, P


    if tracer.has_pixelization:
        mappers_of_planes = tracer.mappers_of_planes_from_grid(
            grid=grid,
            inversion_uses_border=inversion_uses_border,
            preload_sparse_grids_of_planes=preload_sparse_grids_of_planes,
        )

        mapper = mappers_of_planes[-1]

        regularization = tracer.regularizations_of_planes[-1]
        # print(
        #     "coefficient = ", regularization.coefficient
        # )

        regularization_matrix = regularization.regularization_matrix_from_mapper(
            mapper=mapper
        )

        # plt.figure()
        # plt.imshow(regularization_matrix)
        # plt.show()
        # exit()

        start = time.time()
        reconstruction, Op, P = reconstruction_from_(
            masked_dataset=masked_dataset,
            mapper=mapper,
            regularization_matrix=regularization_matrix,
            transformer=transformer,
            settings=settings,
            transformer_from_averaged=transformer_from_averaged,
        )
        end = time.time()
        print(
            "It took t={} to perform the inversion".format(end - start)
        )

    else:
        raise ValueError(
            "This has not been implemented yet"
        )

    return inversion_abstract(
        Op=Op,
        reconstruction=reconstruction,
        P=P,
        mapper=mapper,
        tracer=tracer,
        transformer=transformer,
        regularization_matrix=regularization_matrix,
        shape=masked_dataset.data_shape,
    )


def redshifts_from_list_of_galaxies(galaxies):

    return [galaxy.redshift for galaxy in galaxies]


def optimize_(masked_dataset, transformer, lens_galaxies, source_redshift, pixelization, x0=0.0, regularization_class=al.reg.Constant):

    import autoarray_classes.autoarray_fit_classes as autoarray_fit_classes

    class optimization:

        def __init__(self, source_redshift, pixelization, lens_galaxies, regularization_class):

            self.source_redshift = source_redshift

            # NOTE:
            self.pixelization = pixelization

            if all(redshift < self.source_redshift
                for redshift in redshifts_from_list_of_galaxies(galaxies=lens_galaxies)
            ):
                self.lens_galaxies = lens_galaxies
            else:
                raise ValueError("...")

            self.regularization_class = regularization_class


        def source_galaxy(self, coefficient):

            return al.Galaxy(
                redshift=self.source_redshift,
                pixelization=self.pixelization,
                regularization=self.regularization_class(
                    coefficient=coefficient
                ),
            )

        def tracer_from(self, coefficient):

            return al.Tracer.from_galaxies(
                galaxies=[
                    *self.lens_galaxies, self.source_galaxy(coefficient=coefficient)
                ]
            )



    def func(log_coefficient, masked_dataset, grid, transformer, optimization_class):

        coefficient = 10.0**log_coefficient[0];print(log_coefficient)

        inversion = inversion_from_(
            masked_dataset=masked_dataset,
            grid=grid,
            transformer=transformer,
            tracer=optimization_class.tracer_from(
                coefficient=coefficient
            ),
        )

        fit = autoarray_fit_classes.FitInterferometer(
            masked_dataset=masked_dataset,
            model_data=inversion.model_data,
            inversion=inversion
        )

        return -fit.evidence

    # NOTE:
    # if isinstance(regularization_class, al.reg.Constant):
    #     print("OK")
    # else:
    #     raise TypeError(
    #         "This has not been implemented yet"
    #     )



    optimization_class = optimization(
        source_redshift=source_redshift,
        pixelization=pixelization,
        lens_galaxies=lens_galaxies,
        regularization_class=regularization_class
    )

    grid = al.Grid.from_mask(
        mask=transformer.real_space_mask
    )

    # log_coefficient_arr = np.linspace(-3.0, 8.0, 20)
    # evidences = np.zeros(shape=log_coefficient_arr.shape)
    # for i, log_coefficient in enumerate(log_coefficient_arr):
    #
    #     evidence = func(
    #         log_coefficient=log_coefficient,
    #         masked_dataset=masked_dataset,
    #         grid=grid,
    #         transformer=transformer,
    #         optimization_class=optimization_class
    #     )
    #
    #     evidences[i] = evidence
    #
    # plt.figure()
    # plt.plot(log_coefficient_arr, evidences)
    # plt.show()
    # exit()

    res = optimize.minimize(
        func,
        x0=x0,
        args=(masked_dataset, grid, transformer, optimization_class),
        method='COBYLA'
    )
    #print(res)

    return res.x

3a7f1f212c7fb8149af459386b68c5a8529b0e22
<mask token>


def lensed_image_from_tracer(tracer, grid, image, interpolator=None):
    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(grid=
        grid, plane_index_limit=tracer.upper_plane_index_with_light_profile)
    lensed_image = ray_trace(traced_grids_of_planes=traced_grids_of_planes,
        image=image, grid=grid)
    return lensed_image


<mask token>


def gradients_from_galaxy(galaxy, grid):
    raise ValueError('This has not been implemented yet.')


<mask token>


class inversion_settings:

    def __init__(self, method='operators'):
        self.method = method


class inversion_abstract:

    def __init__(self, Op, reconstruction, P, transformer, tracer, mapper,
        regularization_matrix, shape):
        self.Op = Op
        self._reconstruction = reconstruction
        self.P = P
        self.tracer = tracer
        self.mapper = mapper
        self.transformer = transformer
        self.regularization_matrix = regularization_matrix
        self.shape = shape

    @property
    def reconstruction(self):
        return self._reconstruction

    @reconstruction.setter
    def reconstruction(self, x):
        self._reconstruction = x

    def reshape(self, x):
        return x.reshape(self.shape)

    @property
    def mapped_reconstructed_image(self):
        array_1d = self.mapper.mapping_matrix @ self.reconstruction
        return self.mapper.grid.mapping.array_stored_1d_from_array_1d(array_1d
            =array_1d)

    @property
    def mapped_reconstructed_visibilities(self):
        return self.transformer.visibilities_from_image(image=self.
            mapped_reconstructed_image.in_2d)

    @property
    def model_data(self):
        return self.mapped_reconstructed_visibilities

    def log_det(self, x):
        return np.linalg.slogdet(x)[1]

    @property
    def regularization_term(self):
        return np.matmul(self.reconstruction.T, np.matmul(self.
            regularization_matrix, self.reconstruction))

    @property
    def log_det_curvature_reg_matrix_term(self):
        return self.log_det(x=self.P)

    @property
    def log_det_regularization_matrix_term(self):
        return self.log_det(x=self.regularization_matrix)


<mask token>


def preconditioner_from(mapping_matrix, weights, regularization_matrix):
    pass


class AbstractMaskedDatasetSettings:

    def __init__(self):
        pass


class AbstractMaskedDataset:

    def __init__(self, data_1d, weights_1d):
        self.data_1d = data_1d
        self.weights_1d = weights_1d


<mask token>


class inversion_settings:

    def __init__(self, method='operators', use_preconditioner=True,
        preconditioner_method='powell'):
        self.method = method
        self.use_preconditioner = use_preconditioner
        self.preconditioner_method = preconditioner_method


<mask token>

<mask token>


def visibilities_from_cube(cube, transformers, shape):


    class Image:

        def __init__(self, array_2d):
            self.array_2d = array_2d

        @property
        def in_1d_binned(self):
            return np.ndarray.flatten(self.array_2d)

        @property
        def in_2d_binned(self):
            return self.array_2d
    if len(transformers) == shape[0]:
        visibilities = np.zeros(shape=shape)
        for i in range(shape[0]):
            visibilities[i] = transformers[i].visibilities_from_image(image
                =Image(array_2d=cube[i]))
    else:
        raise ValueError('shape mismatch')
    return visibilities


<mask token>


def lensed_image_from_tracer(tracer, grid, image, interpolator=None):
    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(grid=
        grid, plane_index_limit=tracer.upper_plane_index_with_light_profile)
    lensed_image = ray_trace(traced_grids_of_planes=traced_grids_of_planes,
        image=image, grid=grid)
    return lensed_image


def lensed_cube_from_tracer(tracer, grid, cube, z_mask=None, interpolator=None
    ):
    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(grid=
        grid, plane_index_limit=tracer.upper_plane_index_with_light_profile)
    lensed_cube = np.zeros(shape=cube.shape)
    for i in range(lensed_cube.shape[0]):
        lensed_cube[i] = ray_trace(traced_grids_of_planes=
            traced_grids_of_planes, image=cube[i], grid=grid)
    return lensed_cube


<mask token>


def galaxies_with_light_profiles_from_tracer(tracer):
    return list(compress(tracer.galaxies, list(map(lambda galaxy: galaxy.
        has_light_profile, tracer.galaxies))))


<mask token>


def gradients(a, grid):
    gradient_y = np.gradient(a.in_2d, grid.in_2d[:, 0, 0], axis=0)
    gradient_x = np.gradient(a.in_2d, grid.in_2d[0, :, 1], axis=1)
    return gradient_y, gradient_x


<mask token>


def gradients_from_galaxy(galaxy, grid):
    raise ValueError('This has not been implemented yet.')


def replace_light_profiles_in_tracer_with_pixelization(tracer, pixelization,
    regularization=None):
    galaxies_with_mass_profiles = galaxies_with_mass_profiles_from_tracer(
        tracer=tracer)
    planes_with_light_profile = tracer.planes_with_light_profile
    if not planes_with_light_profile:
        raise ValueError('...')
    elif len(planes_with_light_profile) == 1:
        redshift = planes_with_light_profile[0].redshift
    elif len(planes_with_light_profile) > 1:
        raise ValueError('This has not been implemented yet')
    else:
        raise ValueError('...')
    return al.Tracer.from_galaxies(galaxies=[*
        galaxies_with_mass_profiles_from_tracer(tracer=tracer), al.Galaxy(
        redshift=redshift, pixelization=pixelization, regularization=
        regularization if regularization is not None else al.reg.Constant(
        coefficient=1.0))])


def plot_mapping_matrix(mapping_matrix, grid, nrows, ncols, figsize=None):

    def array_wrapper(array, grid):
        return grid.mask.mapping.array_stored_2d_from_array_1d(array_1d=array)
    n = mapping_matrix.shape[-1]
    if nrows * ncols < n:
        raise ValueError('The total number of panels {} is less than {}'.
            format(nrows * ncols, n))
    figure, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize if
        figsize is not None else (16, 8))
    k = 0
    for i in range(nrows):
        for j in range(ncols):
            if k < n:
                axes[i, j].imshow(array_wrapper(array=mapping_matrix[:, k],
                    grid=grid))
            else:
                axes[i, j].axis('off')
            axes[i, j].set_xticks([])
            axes[i, j].set_yticks([])
            k += 1
    plt.subplots_adjust(wspace=0.0, hspace=0.0)
    plt.show()


<mask token>


class inversion_settings:

    def __init__(self, method='operators'):
        self.method = method


class inversion_abstract:

    def __init__(self, Op, reconstruction, P, transformer, tracer, mapper,
        regularization_matrix, shape):
        self.Op = Op
        self._reconstruction = reconstruction
        self.P = P
        self.tracer = tracer
        self.mapper = mapper
        self.transformer = transformer
        self.regularization_matrix = regularization_matrix
        self.shape = shape

    @property
    def reconstruction(self):
        return self._reconstruction

    @reconstruction.setter
    def reconstruction(self, x):
        self._reconstruction = x

    def reshape(self, x):
        return x.reshape(self.shape)

    @property
    def mapped_reconstructed_image(self):
        array_1d = self.mapper.mapping_matrix @ self.reconstruction
        return self.mapper.grid.mapping.array_stored_1d_from_array_1d(array_1d
            =array_1d)

    @property
    def mapped_reconstructed_visibilities(self):
        return self.transformer.visibilities_from_image(image=self.
            mapped_reconstructed_image.in_2d)

    @property
    def model_data(self):
        return self.mapped_reconstructed_visibilities

    def log_det(self, x):
        return np.linalg.slogdet(x)[1]

    @property
    def regularization_term(self):
        return np.matmul(self.reconstruction.T, np.matmul(self.
            regularization_matrix, self.reconstruction))

    @property
    def log_det_curvature_reg_matrix_term(self):
        return self.log_det(x=self.P)

    @property
    def log_det_regularization_matrix_term(self):
        return self.log_det(x=self.regularization_matrix)


def wrap_mapping_matrix(mapping_matrix):
    return pylops.MatrixMult(sparse.bsr_matrix(mapper.mapping_matrix),
        dtype='float64')


def preconditioner_from(mapping_matrix, weights, regularization_matrix):
    pass


class AbstractMaskedDatasetSettings:

    def __init__(self):
        pass


class AbstractMaskedDataset:

    def __init__(self, data_1d, weights_1d):
        self.data_1d = data_1d
        self.weights_1d = weights_1d


def AbstractMaskedDatasetConverter(cls):
    attributes = ['data_1d', 'weights_1d']
    for attr in attributes:
        if hasattr(cls, attr):
            pass
        elif attr == 'data_1':
            _attributes = []


class inversion_settings:

    def __init__(self, method='operators', use_preconditioner=True,
        preconditioner_method='powell'):
        self.method = method
        self.use_preconditioner = use_preconditioner
        self.preconditioner_method = preconditioner_method


def inversion_from_(masked_dataset, tracer, grid, transformer,
    transformer_from_averaged=None, inversion_uses_border=True,
    preload_sparse_grids_of_planes=None, settings=inversion_settings()):

    def reconstruction_from_(masked_dataset, mapper, regularization_matrix,
        transformer, settings, transformer_from_averaged=None):
        mapping_matrix_as_operator = pylops.MatrixMult(sparse.bsr_matrix(
            mapper.mapping_matrix), dtype='float64')
        Op = transformer * mapping_matrix_as_operator
        if settings.use_preconditioner:
            if settings.preconditioner_method == 'approx':
                if transformer_from_averaged is None:
                    raise ValueError('This is a requirement for this method')
                start = time.time()
                transformed_mapping_matrices = (transformer_from_averaged.
                    transformed_mapping_matrices_from_mapping_matrix(
                    mapping_matrix=mapper.mapping_matrix))
                end = time.time()
                print('It took t={} to compute thetransformed mapping matrices'
                    .format(end - start))
                F_real = np.dot(transformed_mapping_matrices[0].T *
                    masked_dataset.weights_averaged_from_reshaped[:, 0],
                    transformed_mapping_matrices[0])
                F_imag = np.dot(transformed_mapping_matrices[1].T *
                    masked_dataset.weights_averaged_from_reshaped[:, 1],
                    transformed_mapping_matrices[1])
                F_approx = F_real + F_imag
            else:
                F_approx = np.multiply(np.sum(masked_dataset.weights_1d), 
                    mapper.mapping_matrix.T @ mapper.mapping_matrix)
            P = np.add(F_approx, regularization_matrix)
            P_inv = np.linalg.inv(P)
            MOp = pylops.MatrixMult(sparse.bsr_matrix(P_inv), dtype='float64')
        else:
            MOp = None
        print('Performing inversion ...')
        start = time.time()
        values = pylops.NormalEquationsInversion(Op=Op, Regs=None, data=
            masked_dataset.data_1d, Weight=pylops.Diagonal(diag=
            masked_dataset.weights_1d, dtype='float64'), epsNRs=[1.0],
            NRegs=[pylops.MatrixMult(sparse.bsr_matrix(
            regularization_matrix), dtype='float64')], M=MOp, tol=1e-08)
        end = time.time()
        print('It took t={} to perform the reconstruction'.format(end - start))
        return values, Op, P
    if tracer.has_pixelization:
        mappers_of_planes = tracer.mappers_of_planes_from_grid(grid=grid,
            inversion_uses_border=inversion_uses_border,
            preload_sparse_grids_of_planes=preload_sparse_grids_of_planes)
        mapper = mappers_of_planes[-1]
        regularization = tracer.regularizations_of_planes[-1]
        regularization_matrix = (regularization.
            regularization_matrix_from_mapper(mapper=mapper))
        start = time.time()
        reconstruction, Op, P = reconstruction_from_(masked_dataset=
            masked_dataset, mapper=mapper, regularization_matrix=
            regularization_matrix, transformer=transformer, settings=
            settings, transformer_from_averaged=transformer_from_averaged)
        end = time.time()
        print('It took t={} to perform the inversion'.format(end - start))
    else:
        raise ValueError('This has not been implemented yet')
    return inversion_abstract(Op=Op, reconstruction=reconstruction, P=P,
        mapper=mapper, tracer=tracer, transformer=transformer,
        regularization_matrix=regularization_matrix, shape=masked_dataset.
        data_shape)


def redshifts_from_list_of_galaxies(galaxies):
    return [galaxy.redshift for galaxy in galaxies]


def optimize_(masked_dataset, transformer, lens_galaxies, source_redshift,
    pixelization, x0=0.0, regularization_class=al.reg.Constant):
    import autoarray_classes.autoarray_fit_classes as autoarray_fit_classes


    class optimization:

        def __init__(self, source_redshift, pixelization, lens_galaxies,
            regularization_class):
            self.source_redshift = source_redshift
            self.pixelization = pixelization
            if all(redshift < self.source_redshift for redshift in
                redshifts_from_list_of_galaxies(galaxies=lens_galaxies)):
                self.lens_galaxies = lens_galaxies
            else:
                raise ValueError('...')
            self.regularization_class = regularization_class

        def source_galaxy(self, coefficient):
            return al.Galaxy(redshift=self.source_redshift, pixelization=
                self.pixelization, regularization=self.regularization_class
                (coefficient=coefficient))

        def tracer_from(self, coefficient):
            return al.Tracer.from_galaxies(galaxies=[*self.lens_galaxies,
                self.source_galaxy(coefficient=coefficient)])

    def func(log_coefficient, masked_dataset, grid, transformer,
        optimization_class):
        coefficient = 10.0 ** log_coefficient[0]
        print(log_coefficient)
        inversion = inversion_from_(masked_dataset=masked_dataset, grid=
            grid, transformer=transformer, tracer=optimization_class.
            tracer_from(coefficient=coefficient))
        fit = autoarray_fit_classes.FitInterferometer(masked_dataset=
            masked_dataset, model_data=inversion.model_data, inversion=
            inversion)
        return -fit.evidence
    optimization_class = optimization(source_redshift=source_redshift,
        pixelization=pixelization, lens_galaxies=lens_galaxies,
        regularization_class=regularization_class)
    grid = al.Grid.from_mask(mask=transformer.real_space_mask)
    res = optimize.minimize(func, x0=x0, args=(masked_dataset, grid,
        transformer, optimization_class), method='COBYLA')
    return res.x

<mask token>


def visibilities_from_cube(cube, transformers, shape):


    class Image:

        def __init__(self, array_2d):
            self.array_2d = array_2d

        @property
        def in_1d_binned(self):
            return np.ndarray.flatten(self.array_2d)

        @property
        def in_2d_binned(self):
            return self.array_2d
    if len(transformers) == shape[0]:
        visibilities = np.zeros(shape=shape)
        for i in range(shape[0]):
            visibilities[i] = transformers[i].visibilities_from_image(image
                =Image(array_2d=cube[i]))
    else:
        raise ValueError('shape mismatch')
    return visibilities


<mask token>


def ray_trace(traced_grids_of_planes, image, grid, interpolator=None):
    traced_grid_i = traced_grids_of_planes[0]
    traced_grid_j = traced_grids_of_planes[1]
    x_interp = np.unique(traced_grid_i[:, 0])
    y_interp = np.unique(traced_grid_i[:, 1])
    image_interp = interpolate.RegularGridInterpolator(points=(y_interp,
        x_interp), values=image[::-1, :], method='linear', bounds_error=
        False, fill_value=0.0)
    lensed_image = image_interp(traced_grid_j).reshape(image.shape)
    return lensed_image


<mask token>


def lensed_image_from_tracer(tracer, grid, image, interpolator=None):
    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(grid=
        grid, plane_index_limit=tracer.upper_plane_index_with_light_profile)
    lensed_image = ray_trace(traced_grids_of_planes=traced_grids_of_planes,
        image=image, grid=grid)
    return lensed_image


def lensed_cube_from_tracer(tracer, grid, cube, z_mask=None, interpolator=None
    ):
    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(grid=
        grid, plane_index_limit=tracer.upper_plane_index_with_light_profile)
    lensed_cube = np.zeros(shape=cube.shape)
    for i in range(lensed_cube.shape[0]):
        lensed_cube[i] = ray_trace(traced_grids_of_planes=
            traced_grids_of_planes, image=cube[i], grid=grid)
    return lensed_cube


def mapper_from_tracer_and_grid(tracer, grid):
    mappers_of_planes = tracer.mappers_of_planes_from_grid(grid=grid)
    return mappers_of_planes[-1]


def test__regularization_matrix(tracer, grid, pixelization):
    pass


<mask token>


def galaxies_with_mass_profiles_from_tracer(tracer):
    return list(compress(tracer.galaxies, list(map(lambda galaxy: galaxy.
        has_mass_profile, tracer.galaxies))))


def galaxies_with_light_profiles_from_tracer(tracer):
    return list(compress(tracer.galaxies, list(map(lambda galaxy: galaxy.
        has_light_profile, tracer.galaxies))))


<mask token>


def gradients(a, grid):
    gradient_y = np.gradient(a.in_2d, grid.in_2d[:, 0, 0], axis=0)
    gradient_x = np.gradient(a.in_2d, grid.in_2d[0, :, 1], axis=1)
    return gradient_y, gradient_x


def deflections_via_potential_from_grid(potential, grid):
    return gradients(a=potential, grid=grid)


def gradients_from_galaxy(galaxy, grid):
    raise ValueError('This has not been implemented yet.')


def replace_light_profiles_in_tracer_with_pixelization(tracer, pixelization,
    regularization=None):
    galaxies_with_mass_profiles = galaxies_with_mass_profiles_from_tracer(
        tracer=tracer)
    planes_with_light_profile = tracer.planes_with_light_profile
    if not planes_with_light_profile:
        raise ValueError('...')
    elif len(planes_with_light_profile) == 1:
        redshift = planes_with_light_profile[0].redshift
    elif len(planes_with_light_profile) > 1:
        raise ValueError('This has not been implemented yet')
    else:
        raise ValueError('...')
    return al.Tracer.from_galaxies(galaxies=[*
        galaxies_with_mass_profiles_from_tracer(tracer=tracer), al.Galaxy(
        redshift=redshift, pixelization=pixelization, regularization=
        regularization if regularization is not None else al.reg.Constant(
        coefficient=1.0))])


def plot_mapping_matrix(mapping_matrix, grid, nrows, ncols, figsize=None):

    def array_wrapper(array, grid):
        return grid.mask.mapping.array_stored_2d_from_array_1d(array_1d=array)
    n = mapping_matrix.shape[-1]
    if nrows * ncols < n:
        raise ValueError('The total number of panels {} is less than {}'.
            format(nrows * ncols, n))
    figure, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize if
        figsize is not None else (16, 8))
    k = 0
    for i in range(nrows):
        for j in range(ncols):
            if k < n:
                axes[i, j].imshow(array_wrapper(array=mapping_matrix[:, k],
                    grid=grid))
            else:
                axes[i, j].axis('off')
            axes[i, j].set_xticks([])
            axes[i, j].set_yticks([])
            k += 1
    plt.subplots_adjust(wspace=0.0, hspace=0.0)
    plt.show()


<mask token>


class inversion_settings:

    def __init__(self, method='operators'):
        self.method = method


class inversion_abstract:

    def __init__(self, Op, reconstruction, P, transformer, tracer, mapper,
        regularization_matrix, shape):
        self.Op = Op
        self._reconstruction = reconstruction
        self.P = P
        self.tracer = tracer
        self.mapper = mapper
        self.transformer = transformer
        self.regularization_matrix = regularization_matrix
        self.shape = shape

    @property
    def reconstruction(self):
        return self._reconstruction

    @reconstruction.setter
    def reconstruction(self, x):
        self._reconstruction = x

    def reshape(self, x):
        return x.reshape(self.shape)

    @property
    def mapped_reconstructed_image(self):
        array_1d = self.mapper.mapping_matrix @ self.reconstruction
        return self.mapper.grid.mapping.array_stored_1d_from_array_1d(array_1d
            =array_1d)

    @property
    def mapped_reconstructed_visibilities(self):
        return self.transformer.visibilities_from_image(image=self.
            mapped_reconstructed_image.in_2d)

    @property
    def model_data(self):
        return self.mapped_reconstructed_visibilities

    def log_det(self, x):
        return np.linalg.slogdet(x)[1]

    @property
    def regularization_term(self):
        return np.matmul(self.reconstruction.T, np.matmul(self.
            regularization_matrix, self.reconstruction))

    @property
    def log_det_curvature_reg_matrix_term(self):
        return self.log_det(x=self.P)

    @property
    def log_det_regularization_matrix_term(self):
        return self.log_det(x=self.regularization_matrix)


def wrap_mapping_matrix(mapping_matrix):
    return pylops.MatrixMult(sparse.bsr_matrix(mapper.mapping_matrix),
        dtype='float64')


def preconditioner_from(mapping_matrix, weights, regularization_matrix):
    pass


class AbstractMaskedDatasetSettings:

    def __init__(self):
        pass


class AbstractMaskedDataset:

    def __init__(self, data_1d, weights_1d):
        self.data_1d = data_1d
        self.weights_1d = weights_1d


def AbstractMaskedDatasetConverter(cls):
    attributes = ['data_1d', 'weights_1d']
    for attr in attributes:
        if hasattr(cls, attr):
            pass
        elif attr == 'data_1':
            _attributes = []


class inversion_settings:

    def __init__(self, method='operators', use_preconditioner=True,
        preconditioner_method='powell'):
        self.method = method
        self.use_preconditioner = use_preconditioner
        self.preconditioner_method = preconditioner_method


def inversion_from_(masked_dataset, tracer, grid, transformer,
    transformer_from_averaged=None, inversion_uses_border=True,
    preload_sparse_grids_of_planes=None, settings=inversion_settings()):

    def reconstruction_from_(masked_dataset, mapper, regularization_matrix,
        transformer, settings, transformer_from_averaged=None):
        mapping_matrix_as_operator = pylops.MatrixMult(sparse.bsr_matrix(
            mapper.mapping_matrix), dtype='float64')
        Op = transformer * mapping_matrix_as_operator
        if settings.use_preconditioner:
            if settings.preconditioner_method == 'approx':
                if transformer_from_averaged is None:
                    raise ValueError('This is a requirement for this method')
                start = time.time()
                transformed_mapping_matrices = (transformer_from_averaged.
                    transformed_mapping_matrices_from_mapping_matrix(
                    mapping_matrix=mapper.mapping_matrix))
                end = time.time()
                print('It took t={} to compute thetransformed mapping matrices'
                    .format(end - start))
                F_real = np.dot(transformed_mapping_matrices[0].T *
                    masked_dataset.weights_averaged_from_reshaped[:, 0],
                    transformed_mapping_matrices[0])
                F_imag = np.dot(transformed_mapping_matrices[1].T *
                    masked_dataset.weights_averaged_from_reshaped[:, 1],
                    transformed_mapping_matrices[1])
                F_approx = F_real + F_imag
            else:
                F_approx = np.multiply(np.sum(masked_dataset.weights_1d), 
                    mapper.mapping_matrix.T @ mapper.mapping_matrix)
            P = np.add(F_approx, regularization_matrix)
            P_inv = np.linalg.inv(P)
            MOp = pylops.MatrixMult(sparse.bsr_matrix(P_inv), dtype='float64')
        else:
            MOp = None
        print('Performing inversion ...')
        start = time.time()
        values = pylops.NormalEquationsInversion(Op=Op, Regs=None, data=
            masked_dataset.data_1d, Weight=pylops.Diagonal(diag=
            masked_dataset.weights_1d, dtype='float64'), epsNRs=[1.0],
            NRegs=[pylops.MatrixMult(sparse.bsr_matrix(
            regularization_matrix), dtype='float64')], M=MOp, tol=1e-08)
        end = time.time()
        print('It took t={} to perform the reconstruction'.format(end - start))
        return values, Op, P
    if tracer.has_pixelization:
        mappers_of_planes = tracer.mappers_of_planes_from_grid(grid=grid,
            inversion_uses_border=inversion_uses_border,
            preload_sparse_grids_of_planes=preload_sparse_grids_of_planes)
        mapper = mappers_of_planes[-1]
        regularization = tracer.regularizations_of_planes[-1]
        regularization_matrix = (regularization.
            regularization_matrix_from_mapper(mapper=mapper))
        start = time.time()
        reconstruction, Op, P = reconstruction_from_(masked_dataset=
            masked_dataset, mapper=mapper, regularization_matrix=
            regularization_matrix, transformer=transformer, settings=
            settings, transformer_from_averaged=transformer_from_averaged)
        end = time.time()
        print('It took t={} to perform the inversion'.format(end - start))
    else:
        raise ValueError('This has not been implemented yet')
    return inversion_abstract(Op=Op, reconstruction=reconstruction, P=P,
        mapper=mapper, tracer=tracer, transformer=transformer,
        regularization_matrix=regularization_matrix, shape=masked_dataset.
        data_shape)


def redshifts_from_list_of_galaxies(galaxies):
    return [galaxy.redshift for galaxy in galaxies]


def optimize_(masked_dataset, transformer, lens_galaxies, source_redshift,
    pixelization, x0=0.0, regularization_class=al.reg.Constant):
    import autoarray_classes.autoarray_fit_classes as autoarray_fit_classes


    class optimization:

        def __init__(self, source_redshift, pixelization, lens_galaxies,
            regularization_class):
            self.source_redshift = source_redshift
            self.pixelization = pixelization
            if all(redshift < self.source_redshift for redshift in
                redshifts_from_list_of_galaxies(galaxies=lens_galaxies)):
                self.lens_galaxies = lens_galaxies
            else:
                raise ValueError('...')
            self.regularization_class = regularization_class

        def source_galaxy(self, coefficient):
            return al.Galaxy(redshift=self.source_redshift, pixelization=
                self.pixelization, regularization=self.regularization_class
                (coefficient=coefficient))

        def tracer_from(self, coefficient):
            return al.Tracer.from_galaxies(galaxies=[*self.lens_galaxies,
                self.source_galaxy(coefficient=coefficient)])

    def func(log_coefficient, masked_dataset, grid, transformer,
        optimization_class):
        coefficient = 10.0 ** log_coefficient[0]
        print(log_coefficient)
        inversion = inversion_from_(masked_dataset=masked_dataset, grid=
            grid, transformer=transformer, tracer=optimization_class.
            tracer_from(coefficient=coefficient))
        fit = autoarray_fit_classes.FitInterferometer(masked_dataset=
            masked_dataset, model_data=inversion.model_data, inversion=
            inversion)
        return -fit.evidence
    optimization_class = optimization(source_redshift=source_redshift,
        pixelization=pixelization, lens_galaxies=lens_galaxies,
        regularization_class=regularization_class)
    grid = al.Grid.from_mask(mask=transformer.real_space_mask)
    res = optimize.minimize(func, x0=x0, args=(masked_dataset, grid,
        transformer, optimization_class), method='COBYLA')
    return res.x

<mask token>


def visibilities_from_cube(cube, transformers, shape):


    class Image:

        def __init__(self, array_2d):
            self.array_2d = array_2d

        @property
        def in_1d_binned(self):
            return np.ndarray.flatten(self.array_2d)

        @property
        def in_2d_binned(self):
            return self.array_2d
    if len(transformers) == shape[0]:
        visibilities = np.zeros(shape=shape)
        for i in range(shape[0]):
            visibilities[i] = transformers[i].visibilities_from_image(image
                =Image(array_2d=cube[i]))
    else:
        raise ValueError('shape mismatch')
    return visibilities


def visibilities_from_image_and_uv_wavelengths(image, uv_wavelengths, grid,
    transformer_class):
    if len(uv_wavelengths.shape) == 2:
        transformer = transformer_class(uv_wavelengths=uv_wavelengths, grid
            =grid.in_radians)
        visibilities = transformer.visibilities_from_image(image=image)
    if len(uv_wavelengths.shape) == 3:
        transformers = [transformer_class(uv_wavelengths=uv_wavelengths[i],
            grid=grid.in_radians) for i in range(uv_wavelengths.shape[0])]
        visibilities = np.zeros(shape=uv_wavelengths.shape)
        for i in range(uv_wavelengths.shape[0]):
            visibilities[i, ...] = transformers[i].visibilities_from_image(
                image=image)
    if len(uv_wavelengths.shape) == 4:
        raise ValueError('This has not been implemented yet')
    return visibilities


def ray_trace(traced_grids_of_planes, image, grid, interpolator=None):
    traced_grid_i = traced_grids_of_planes[0]
    traced_grid_j = traced_grids_of_planes[1]
    x_interp = np.unique(traced_grid_i[:, 0])
    y_interp = np.unique(traced_grid_i[:, 1])
    image_interp = interpolate.RegularGridInterpolator(points=(y_interp,
        x_interp), values=image[::-1, :], method='linear', bounds_error=
        False, fill_value=0.0)
    lensed_image = image_interp(traced_grid_j).reshape(image.shape)
    return lensed_image


<mask token>


def lensed_image_from_tracer(tracer, grid, image, interpolator=None):
    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(grid=
        grid, plane_index_limit=tracer.upper_plane_index_with_light_profile)
    lensed_image = ray_trace(traced_grids_of_planes=traced_grids_of_planes,
        image=image, grid=grid)
    return lensed_image


def lensed_cube_from_tracer(tracer, grid, cube, z_mask=None, interpolator=None
    ):
    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(grid=
        grid, plane_index_limit=tracer.upper_plane_index_with_light_profile)
    lensed_cube = np.zeros(shape=cube.shape)
    for i in range(lensed_cube.shape[0]):
        lensed_cube[i] = ray_trace(traced_grids_of_planes=
            traced_grids_of_planes, image=cube[i], grid=grid)
    return lensed_cube


def mapper_from_tracer_and_grid(tracer, grid):
    mappers_of_planes = tracer.mappers_of_planes_from_grid(grid=grid)
    return mappers_of_planes[-1]


def test__regularization_matrix(tracer, grid, pixelization):
    pass


def transformed_mapping_matrices_from_transformer_grid_and_tracer(transformer,
    grid, tracer):
    mapper = mapper_from_tracer_and_grid(tracer=tracer, grid=grid)
    return transformer.transformed_mapping_matrices_from_mapping_matrix(
        mapping_matrix=mapper.mapping_matrix)


def get_parameter_from_tracer(tracer, parameter_name='centre'):
    for mass_profile in tracer.mass_profiles:
        if mass_profile.__class__.__name__ in ['EllipticalPowerLaw',
            'EllipticalIsothermal']:
            parameter = getattr(mass_profile, parameter_name)
    return parameter


def galaxies_with_mass_profiles_from_tracer(tracer):
    return list(compress(tracer.galaxies, list(map(lambda galaxy: galaxy.
        has_mass_profile, tracer.galaxies))))


def galaxies_with_light_profiles_from_tracer(tracer):
    return list(compress(tracer.galaxies, list(map(lambda galaxy: galaxy.
        has_light_profile, tracer.galaxies))))


def galaxies_have_mass_profiles(galaxies):
    return all(list(map(lambda galaxy: galaxy.has_mass_profile, galaxies)))


def tracer_from_len_galaxies_and_src_galaxy_with_pixelization(len_galaxies,
    src_galaxy):
    if not isinstance(src_galaxy, al.Galaxy):
        raise ValueError('The input must be an al.Galaxy class')
    if src_galaxy.has_pixelization:
        if galaxies_have_mass_profiles(galaxies=len_galaxies):
            galaxies = len_galaxies + [src_galaxy]
        else:
            raise ValueError('Not all galaxies have mass profiles')
    else:
        raise ValueError("The 'src_galaxy' does not have a pixelization")
    return al.Tracer.from_galaxies(galaxies=galaxies)


def tracer_from_instance(instance):
    if hasattr(instance, 'galaxies'):
        tracer = al.Tracer.from_galaxies(galaxies=instance.galaxies)
    else:
        raise AttributeError('...')
    return tracer


def tracer_derivative(tracer, transformer, grid, constants=None, epsilon=
    1e-08, reshaped=True):

    def galaxy_profiles(galaxy):

        def is_analytic_profile(obj):
            return isinstance(obj, (al.mp.MassProfile, al.lp.LightProfile))
        return [key for key, value in galaxy.__dict__.items() if
            is_analytic_profile(value)]
    galaxies = copy.deepcopy(tracer.galaxies)
    if constants is None:
        constants = {'mass': ['core_radius'], 'bulge': [], 'light': []}
    n = 0
    perturbed_tracers = []
    for i, galaxy in enumerate(tracer.galaxies):
        for profile in galaxy_profiles(galaxy=galaxy):
            for j, (name, value) in enumerate(getattr(galaxy, profile).
                __dict__.items()):
                print(name, value)
                if name in constants[profile]:
                    pass
                elif name in ['centre', 'elliptical_comps']:
                    for k, value_k in enumerate(value):
                        perturbed_galaxies = copy.deepcopy(galaxies)
                        if isinstance(value, tuple):
                            value_as_list = list(value)
                        else:
                            raise ValueError('This has not been implemented')
                        value_as_list[k] += epsilon
                        setattr(getattr(perturbed_galaxies[i], profile),
                            name, tuple(value_as_list))
                        perturbed_tracer = al.Tracer.from_galaxies(galaxies
                            =perturbed_galaxies)
                        perturbed_tracers.append(perturbed_tracer)
                        n += 1
                else:
                    perturbed_galaxies = copy.deepcopy(galaxies)
                    setattr(getattr(perturbed_galaxies[i], profile), name, 
                        value + epsilon)
                    perturbed_tracer = al.Tracer.from_galaxies(galaxies=
                        perturbed_galaxies)
                    perturbed_tracers.append(perturbed_tracer)
                    n += 1
    visibilities = tracer.profile_visibilities_from_grid_and_transformer(grid
        =grid, transformer=transformer)
    derivative = np.zeros(shape=(len(perturbed_tracers),) + transformer.
        uv_wavelengths.shape)
    for i in range(derivative.shape[0]):
        visibilities_i = perturbed_tracers[i
            ].profile_visibilities_from_grid_and_transformer(grid=grid,
            transformer=transformer)
        derivative[i] = np.divide(np.subtract(visibilities_i, visibilities),
            epsilon)
    if reshaped:
        derivative = np.concatenate((derivative[..., 0], derivative[..., 1]
            ), axis=-1)
    return derivative


def gradients(a, grid):
    gradient_y = np.gradient(a.in_2d, grid.in_2d[:, 0, 0], axis=0)
    gradient_x = np.gradient(a.in_2d, grid.in_2d[0, :, 1], axis=1)
    return gradient_y, gradient_x


def deflections_via_potential_from_grid(potential, grid):
    return gradients(a=potential, grid=grid)


def gradients_from_galaxy(galaxy, grid):
    raise ValueError('This has not been implemented yet.')


def replace_light_profiles_in_tracer_with_pixelization(tracer, pixelization,
    regularization=None):
    galaxies_with_mass_profiles = galaxies_with_mass_profiles_from_tracer(
        tracer=tracer)
    planes_with_light_profile = tracer.planes_with_light_profile
    if not planes_with_light_profile:
        raise ValueError('...')
    elif len(planes_with_light_profile) == 1:
        redshift = planes_with_light_profile[0].redshift
    elif len(planes_with_light_profile) > 1:
        raise ValueError('This has not been implemented yet')
    else:
        raise ValueError('...')
    return al.Tracer.from_galaxies(galaxies=[*
        galaxies_with_mass_profiles_from_tracer(tracer=tracer), al.Galaxy(
        redshift=redshift, pixelization=pixelization, regularization=
        regularization if regularization is not None else al.reg.Constant(
        coefficient=1.0))])


def plot_mapping_matrix(mapping_matrix, grid, nrows, ncols, figsize=None):

    def array_wrapper(array, grid):
        return grid.mask.mapping.array_stored_2d_from_array_1d(array_1d=array)
    n = mapping_matrix.shape[-1]
    if nrows * ncols < n:
        raise ValueError('The total number of panels {} is less than {}'.
            format(nrows * ncols, n))
    figure, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize if
        figsize is not None else (16, 8))
    k = 0
    for i in range(nrows):
        for j in range(ncols):
            if k < n:
                axes[i, j].imshow(array_wrapper(array=mapping_matrix[:, k],
                    grid=grid))
            else:
                axes[i, j].axis('off')
            axes[i, j].set_xticks([])
            axes[i, j].set_yticks([])
            k += 1
    plt.subplots_adjust(wspace=0.0, hspace=0.0)
    plt.show()


def mappers_of_planes_from_grid(tracer, source_redshift, mask, pixelization):
    raise ValueError('This has not been implemented yet.')


def planes_with_light_profiles(tracer):
    """

    1) Find all planes with light profiles.

    2) Get the galaxies of these planes.

    3) Get the light profiles of those galaxies.

    4) Make images of all those light profiles.

    """


class inversion_settings:

    def __init__(self, method='operators'):
        self.method = method


class inversion_abstract:

    def __init__(self, Op, reconstruction, P, transformer, tracer, mapper,
        regularization_matrix, shape):
        self.Op = Op
        self._reconstruction = reconstruction
        self.P = P
        self.tracer = tracer
        self.mapper = mapper
        self.transformer = transformer
        self.regularization_matrix = regularization_matrix
        self.shape = shape

    @property
    def reconstruction(self):
        return self._reconstruction

    @reconstruction.setter
    def reconstruction(self, x):
        self._reconstruction = x

    def reshape(self, x):
        return x.reshape(self.shape)

    @property
    def mapped_reconstructed_image(self):
        array_1d = self.mapper.mapping_matrix @ self.reconstruction
        return self.mapper.grid.mapping.array_stored_1d_from_array_1d(array_1d
            =array_1d)

    @property
    def mapped_reconstructed_visibilities(self):
        return self.transformer.visibilities_from_image(image=self.
            mapped_reconstructed_image.in_2d)

    @property
    def model_data(self):
        return self.mapped_reconstructed_visibilities

    def log_det(self, x):
        return np.linalg.slogdet(x)[1]

    @property
    def regularization_term(self):
        return np.matmul(self.reconstruction.T, np.matmul(self.
            regularization_matrix, self.reconstruction))

    @property
    def log_det_curvature_reg_matrix_term(self):
        return self.log_det(x=self.P)

    @property
    def log_det_regularization_matrix_term(self):
        return self.log_det(x=self.regularization_matrix)


def wrap_mapping_matrix(mapping_matrix):
    return pylops.MatrixMult(sparse.bsr_matrix(mapper.mapping_matrix),
        dtype='float64')


def preconditioner_from(mapping_matrix, weights, regularization_matrix):
    pass


class AbstractMaskedDatasetSettings:

    def __init__(self):
        pass


class AbstractMaskedDataset:

    def __init__(self, data_1d, weights_1d):
        self.data_1d = data_1d
        self.weights_1d = weights_1d


def AbstractMaskedDatasetConverter(cls):
    attributes = ['data_1d', 'weights_1d']
    for attr in attributes:
        if hasattr(cls, attr):
            pass
        elif attr == 'data_1':
            _attributes = []


class inversion_settings:

    def __init__(self, method='operators', use_preconditioner=True,
        preconditioner_method='powell'):
        self.method = method
        self.use_preconditioner = use_preconditioner
        self.preconditioner_method = preconditioner_method


def inversion_from_(masked_dataset, tracer, grid, transformer,
    transformer_from_averaged=None, inversion_uses_border=True,
    preload_sparse_grids_of_planes=None, settings=inversion_settings()):

    def reconstruction_from_(masked_dataset, mapper, regularization_matrix,
        transformer, settings, transformer_from_averaged=None):
        mapping_matrix_as_operator = pylops.MatrixMult(sparse.bsr_matrix(
            mapper.mapping_matrix), dtype='float64')
        Op = transformer * mapping_matrix_as_operator
        if settings.use_preconditioner:
            if settings.preconditioner_method == 'approx':
                if transformer_from_averaged is None:
                    raise ValueError('This is a requirement for this method')
                start = time.time()
                transformed_mapping_matrices = (transformer_from_averaged.
                    transformed_mapping_matrices_from_mapping_matrix(
                    mapping_matrix=mapper.mapping_matrix))
                end = time.time()
                print('It took t={} to compute thetransformed mapping matrices'
                    .format(end - start))
                F_real = np.dot(transformed_mapping_matrices[0].T *
                    masked_dataset.weights_averaged_from_reshaped[:, 0],
                    transformed_mapping_matrices[0])
                F_imag = np.dot(transformed_mapping_matrices[1].T *
                    masked_dataset.weights_averaged_from_reshaped[:, 1],
                    transformed_mapping_matrices[1])
                F_approx = F_real + F_imag
            else:
                F_approx = np.multiply(np.sum(masked_dataset.weights_1d), 
                    mapper.mapping_matrix.T @ mapper.mapping_matrix)
            P = np.add(F_approx, regularization_matrix)
            P_inv = np.linalg.inv(P)
            MOp = pylops.MatrixMult(sparse.bsr_matrix(P_inv), dtype='float64')
        else:
            MOp = None
        print('Performing inversion ...')
        start = time.time()
        values = pylops.NormalEquationsInversion(Op=Op, Regs=None, data=
            masked_dataset.data_1d, Weight=pylops.Diagonal(diag=
            masked_dataset.weights_1d, dtype='float64'), epsNRs=[1.0],
            NRegs=[pylops.MatrixMult(sparse.bsr_matrix(
            regularization_matrix), dtype='float64')], M=MOp, tol=1e-08)
        end = time.time()
        print('It took t={} to perform the reconstruction'.format(end - start))
        return values, Op, P
    if tracer.has_pixelization:
        mappers_of_planes = tracer.mappers_of_planes_from_grid(grid=grid,
            inversion_uses_border=inversion_uses_border,
            preload_sparse_grids_of_planes=preload_sparse_grids_of_planes)
        mapper = mappers_of_planes[-1]
        regularization = tracer.regularizations_of_planes[-1]
        regularization_matrix = (regularization.
            regularization_matrix_from_mapper(mapper=mapper))
        start = time.time()
        reconstruction, Op, P = reconstruction_from_(masked_dataset=
            masked_dataset, mapper=mapper, regularization_matrix=
            regularization_matrix, transformer=transformer, settings=
            settings, transformer_from_averaged=transformer_from_averaged)
        end = time.time()
        print('It took t={} to perform the inversion'.format(end - start))
    else:
        raise ValueError('This has not been implemented yet')
    return inversion_abstract(Op=Op, reconstruction=reconstruction, P=P,
        mapper=mapper, tracer=tracer, transformer=transformer,
        regularization_matrix=regularization_matrix, shape=masked_dataset.
        data_shape)


def redshifts_from_list_of_galaxies(galaxies):
    return [galaxy.redshift for galaxy in galaxies]


def optimize_(masked_dataset, transformer, lens_galaxies, source_redshift,
    pixelization, x0=0.0, regularization_class=al.reg.Constant):
    import autoarray_classes.autoarray_fit_classes as autoarray_fit_classes


    class optimization:

        def __init__(self, source_redshift, pixelization, lens_galaxies,
            regularization_class):
            self.source_redshift = source_redshift
            self.pixelization = pixelization
            if all(redshift < self.source_redshift for redshift in
                redshifts_from_list_of_galaxies(galaxies=lens_galaxies)):
                self.lens_galaxies = lens_galaxies
            else:
                raise ValueError('...')
            self.regularization_class = regularization_class

        def source_galaxy(self, coefficient):
            return al.Galaxy(redshift=self.source_redshift, pixelization=
                self.pixelization, regularization=self.regularization_class
                (coefficient=coefficient))

        def tracer_from(self, coefficient):
            return al.Tracer.from_galaxies(galaxies=[*self.lens_galaxies,
                self.source_galaxy(coefficient=coefficient)])

    def func(log_coefficient, masked_dataset, grid, transformer,
        optimization_class):
        coefficient = 10.0 ** log_coefficient[0]
        print(log_coefficient)
        inversion = inversion_from_(masked_dataset=masked_dataset, grid=
            grid, transformer=transformer, tracer=optimization_class.
            tracer_from(coefficient=coefficient))
        fit = autoarray_fit_classes.FitInterferometer(masked_dataset=
            masked_dataset, model_data=inversion.model_data, inversion=
            inversion)
        return -fit.evidence
    optimization_class = optimization(source_redshift=source_redshift,
        pixelization=pixelization, lens_galaxies=lens_galaxies,
        regularization_class=regularization_class)
    grid = al.Grid.from_mask(mask=transformer.real_space_mask)
    res = optimize.minimize(func, x0=x0, args=(masked_dataset, grid,
        transformer, optimization_class), method='COBYLA')
    return res.x

import os
import sys
import time
import copy

import numpy as np

import matplotlib
import matplotlib.pyplot as plt

from scipy import (
    sparse, interpolate, optimize
)

from itertools import compress

import pylops

import autolens as al


# NOTE:
sys.path.append(
    "{}/utils".format(os.environ["GitHub"])
)
import pylops_utils as pylops_utils

import autolens_utils.autolens_plot_utils as autolens_plot_utils


def visibilities_from_cube(cube, transformers, shape):

    # NOTE: Check that the shape of the uv_wavelengths in the transformers is the
    # same as the input shape.

    class Image:
        def __init__(self, array_2d):
            self.array_2d = array_2d

        @property
        def in_1d_binned(self):
            return np.ndarray.flatten(self.array_2d)

        @property
        def in_2d_binned(self):
            return self.array_2d

    if len(transformers) == shape[0]:

        visibilities = np.zeros(shape=shape)
        for i in range(shape[0]):
            visibilities[i] = transformers[i].visibilities_from_image(
                image=Image(array_2d=cube[i])
            )

    else:
        raise ValueError("shape mismatch")

    return visibilities


def visibilities_from_image_and_uv_wavelengths(image, uv_wavelengths, grid, transformer_class):

    # NOTE: Image must be a class

    if len(uv_wavelengths.shape) == 2:

        transformer = transformer_class(
            uv_wavelengths=uv_wavelengths,
            grid=grid.in_radians,
        )

        visibilities = transformer.visibilities_from_image(
            image=image
        )

    if len(uv_wavelengths.shape) == 3:

        transformers = [
            transformer_class(
                uv_wavelengths=uv_wavelengths[i],
                grid=grid.in_radians,
            )
            for i in range(uv_wavelengths.shape[0])
        ]

        visibilities = np.zeros(
            shape=uv_wavelengths.shape
        )
        for i in range(uv_wavelengths.shape[0]):
            visibilities[i, ...] = transformers[i].visibilities_from_image(
                image=image
            )

    if len(uv_wavelengths.shape) == 4:

        raise ValueError("This has not been implemented yet")

    return visibilities















# NOTE: This works regardless of how many galaxies are used in the tracer
# NOTE: The spped of this function can not improve by means of saving the weights of the interpolation
def ray_trace(traced_grids_of_planes, image, grid, interpolator=None):

    traced_grid_i = traced_grids_of_planes[0] # This is the regular grid.
    traced_grid_j = traced_grids_of_planes[1]

    # plt.figure()
    # plt.plot(traced_grid_i[:, 0], traced_grid_i[:, 1], linestyle="None", marker="o", color="b")
    # plt.plot(traced_grid_j[:, 0], traced_grid_j[:, 1], linestyle="None", marker="o", color="r")
    # plt.show()
    # exit()

    x_interp = np.unique(traced_grid_i[:, 0])
    y_interp = np.unique(traced_grid_i[:, 1])
    # # NOTE: MADE A CHANGE HERE!
    # x_interp = np.unique(grid[:, 0])
    # y_interp = np.unique(grid[:, 1])

    # plt.figure()
    # plt.plot(x_interp, y_interp, linestyle="None", marker="o")
    # plt.show()
    # exit()

    #print(image.shape);exit()

    #start = time.time()
    image_interp = interpolate.RegularGridInterpolator(
        points=(y_interp, x_interp),
        values=image[::-1, :],
        method="linear",
        bounds_error=False,
        fill_value=0.0
    )

    lensed_image = image_interp(traced_grid_j).reshape(image.shape)
    # # NOTE: MADE A CHANGE HERE!
    # lensed_image = image_interp(traced_grid_j).reshape(grid.shape_2d)
    #end = time.time()
    #print(end - start)

    # NOTE: check autoarray/grids.py -> interpolator (I am not sure if it's going to be faster)
    # lensed_image = interpolate.griddata(
    #     points=traced_grid_i,
    #     values=np.ndarray.flatten(image[:, :]),
    #     xi=traced_grid_j,
    #     method="linear",
    # )
    # lensed_image = lensed_image.reshape(image.shape)

    return lensed_image





# def tracer(lens_galaxy, source_redshift):
#
#     # NOTE: Do I need to check anything about the lens_galaxy?
#
#     tracer = al.Tracer.from_galaxies(
#         galaxies=[
#             lens_galaxy,
#             al.Galaxy(
#                 redshift=source_redshift,
#                 light=al.lp.LightProfile()
#             )
#         ]
#     )
#
#     return tracer

"""
# NOTE: Rename to get_lensed_image_from_image_and_tracer
def func_image(lens_galaxy, source_redshift, grid, image):

    tracer = al.Tracer.from_galaxies(
        galaxies=[
            lens_galaxy,
            al.Galaxy(
                redshift=source_redshift,
                light=al.lp.LightProfile()
            )
        ]
    )

    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
        grid=grid,
        plane_index_limit=tracer.upper_plane_index_with_light_profile
    )

    lensed_image = ray_trace(
        traced_grids_of_planes=traced_grids_of_planes,
        image=image
    )

    return lensed_image


def func_cube(lens_galaxy, source_redshift, grid, cube, interpolator=None):

    tracer = al.Tracer.from_galaxies(
        galaxies=[
            lens_galaxy,
            al.Galaxy(
                redshift=source_redshift,
                light=al.lp.LightProfile()
            )
        ]
    )

    return lensed_cube_from_tracer(
        tracer=tracer,
        grid=grid,
        cube=cube,
        interpolator=interpolator
    )
"""


def lensed_image_from_tracer(tracer, grid, image, interpolator=None):

    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
        grid=grid,
        plane_index_limit=tracer.upper_plane_index_with_light_profile
    )
    # # NOTE: MADE A CHANGE HERE!
    # traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
    #     grid=grid.in_1d_binned,
    #     plane_index_limit=tracer.upper_plane_index_with_light_profile
    # )

    # plt.figure()
    # plt.plot(
    #     traced_grids_of_planes[0][:, 0],
    #     traced_grids_of_planes[0][:, 1],
    #     linestyle="None",
    #     marker="o",
    #     markersize=2,
    #     color="b"
    # )
    # plt.plot(
    #     traced_grids_of_planes[1][:, 0],
    #     traced_grids_of_planes[1][:, 1],
    #     linestyle="None",
    #     marker="o",
    #     markersize=2,
    #     color="r"
    # )
    # # plt.plot(
    # #     grid[:, 0],
    # #     grid[:, 1],
    # #     linestyle="None",
    # #     marker="o",
    # #     markersize=1,
    # #     color="black"
    # # )
    # plt.show()
    # exit()

    # print(
    #     traced_grids_of_planes[0].shape,
    #     traced_grids_of_planes[1].shape,
    #     grid.shape,
    #     grid.shape_2d
    # )
    # exit()

    lensed_image = ray_trace(
        traced_grids_of_planes=traced_grids_of_planes,
        image=image,
        grid=grid
    )

    return lensed_image


def lensed_cube_from_tracer(tracer, grid, cube, z_mask=None, interpolator=None):

    traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
        grid=grid,
        plane_index_limit=tracer.upper_plane_index_with_light_profile
    )

    # start = time.time()
    lensed_cube = np.zeros(
        shape=cube.shape
    )
    for i in range(lensed_cube.shape[0]):

        lensed_cube[i] = ray_trace(
            traced_grids_of_planes=traced_grids_of_planes,
            image=cube[i],
            grid=grid
        )
    # end = time.time()
    # print("Total time:", end - start)
    # exit()

    return lensed_cube


# def lensed_cube_from_galaxies(galaxies, grid, cube, interpolator=None):
#     pass


def mapper_from_tracer_and_grid(tracer, grid):

    # if not isinstance(grid, al.MaskedGrid):
    #     raise ValueError("This has not been implemented yet.")

    mappers_of_planes = tracer.mappers_of_planes_from_grid(
        grid=grid,
    )

    return mappers_of_planes[-1]

def test__regularization_matrix(tracer, grid, pixelization, ):
    pass


def transformed_mapping_matrices_from_transformer_grid_and_tracer(transformer, grid, tracer):

    mapper = mapper_from_tracer_and_grid(
        tracer=tracer,
        grid=grid
    )

    return transformer.transformed_mapping_matrices_from_mapping_matrix(
        mapping_matrix=mapper.mapping_matrix
    )


# TODO: A LOT TO-DO
def get_parameter_from_tracer(tracer, parameter_name="centre"):

    for mass_profile in tracer.mass_profiles:

        if mass_profile.__class__.__name__ in [
            "EllipticalPowerLaw",
            "EllipticalIsothermal"
        ]:

            parameter = getattr(
                mass_profile,
                parameter_name
            )

    return parameter


def galaxies_with_mass_profiles_from_tracer(tracer):

    return list(
        compress(
            tracer.galaxies,
            list(map(lambda galaxy: galaxy.has_mass_profile, tracer.galaxies)))
    )

def galaxies_with_light_profiles_from_tracer(tracer):

    return list(
        compress(
            tracer.galaxies,
            list(map(lambda galaxy: galaxy.has_light_profile, tracer.galaxies)))
    )



def galaxies_have_mass_profiles(galaxies):

    return all(list(map(lambda galaxy: galaxy.has_mass_profile, galaxies)))


def tracer_from_len_galaxies_and_src_galaxy_with_pixelization(len_galaxies, src_galaxy):

    if not isinstance(src_galaxy, al.Galaxy):
        raise ValueError(
            "The input must be an al.Galaxy class"
        )

    if src_galaxy.has_pixelization:
        if galaxies_have_mass_profiles(
            galaxies=len_galaxies
        ):
            galaxies = len_galaxies + [src_galaxy, ]
        else:
            raise ValueError(
                "Not all galaxies have mass profiles"
            )

            # NOTE: Diagnostics
    else:
        raise ValueError(
            "The \'src_galaxy\' does not have a pixelization"
        )

    return al.Tracer.from_galaxies(
        galaxies=galaxies
    )


def tracer_from_instance(instance):

    if hasattr(instance, "galaxies"):
        tracer = al.Tracer.from_galaxies(
            galaxies=instance.galaxies
        )
    else:
        raise AttributeError("...")

    return tracer


def tracer_derivative(tracer, transformer, grid, constants=None, epsilon=1e-8, reshaped=True):

    # if tracer.has_pixelization:
    #     raise ValueError(
    #         "This has not been implemented yet"
    #     )
    # else:
    #     raise ValueError(
    #         "In development ..."
    #     )

    # TODO:
    #
    # 1) Return a list of all the model parameters for each galaxy
    #
    # 2) Iterate through this list and adjust each parameter by 1e-8
    #
    # 3) For each adjustment, define a new tracer and produce visibilities. This
    # will be one column of the derivative matrix

    # derivative = np.zeros(
    #     shape=(np.prod(transformer.uv_wavelengths.shape), )
    # )


    # NOTE: For analytic profiles ...
    def galaxy_profiles(galaxy):

        def is_analytic_profile(obj):
            return isinstance(
                obj, (al.mp.MassProfile, al.lp.LightProfile)
            )

        return [key for key, value in galaxy.__dict__.items()
            if is_analytic_profile(value)
        ]



    galaxies = copy.deepcopy(tracer.galaxies)

    if constants is None:
        #raise ValueError("...")
        constants = {
            "mass":[
                "core_radius"
            ],
            "bulge":[],
            "light":[],
        }

        # constants = {
        #     "mass":[
        #         "core_radius"
        #     ],
        #     "bulge":[],
        #     "light":["intensity", "radius_break", "alpha", "gamma"],
        # }

    # def perturb(galaxies, index, ):
    #     pass
    #
    # def perturb_value(value):
    #
    #     pass

    n = 0

    perturbed_tracers = []
    for i, galaxy in enumerate(tracer.galaxies):

        for profile in galaxy_profiles(galaxy=galaxy):
            #print(profile)

            # if isinstance(profile, al.mp.SphericalNFWMCRLudlow):
            #     print("OK")

            for j, (name, value) in enumerate(
                getattr(galaxy, profile).__dict__.items()
            ):
                print(name, value)

                # NOTE: Check
                if name in constants[profile]:
                    pass
                    # print(
                    #     "The parameter \'{}\' is kept constant".format(name)
                    # )

                elif name in ["centre", "elliptical_comps"]:

                    for k, value_k in enumerate(value):

                        perturbed_galaxies = copy.deepcopy(galaxies)

                        if isinstance(value, tuple):
                            value_as_list = list(value)
                        else:
                            raise ValueError("This has not been implemented")

                        value_as_list[k] += epsilon

                        setattr(
                            getattr(perturbed_galaxies[i], profile), name, tuple(value_as_list)
                        )

                        perturbed_tracer = al.Tracer.from_galaxies(
                            galaxies=perturbed_galaxies
                        )
                        perturbed_tracers.append(perturbed_tracer)

                        n += 1
                else:

                    perturbed_galaxies = copy.deepcopy(galaxies)

                    setattr(
                        getattr(perturbed_galaxies[i], profile), name, value + epsilon
                    )

                    perturbed_tracer = al.Tracer.from_galaxies(
                        galaxies=perturbed_galaxies
                    )
                    perturbed_tracers.append(perturbed_tracer)

                    n += 1

    #exit()

    # 0 centre
    # 2 axis_ratio
    # 3 phi
    # 4 einstein_radius
    # 5 slope
    # 6 core_radius
    # The parameter 'core_radius' is kept constant
    # 6 centre
    # 8 axis_ratio
    # 9 phi
    # 10 intensity
    # The parameter 'intensity' is kept constant
    # 10 effective_radius
    # 11 sersic_index

    visibilities = tracer.profile_visibilities_from_grid_and_transformer(
        grid=grid, transformer=transformer
    )

    derivative = np.zeros(
        shape=((len(perturbed_tracers), ) + transformer.uv_wavelengths.shape)
    )
    for i in range(derivative.shape[0]):
        visibilities_i = perturbed_tracers[i].profile_visibilities_from_grid_and_transformer(
            grid=grid, transformer=transformer
        )

        derivative[i] = np.divide(
            np.subtract(visibilities_i, visibilities),
            epsilon
        )

    # NOTE: Plot
    # plt.figure()
    # plt.imshow(derivative[:, :, 0], aspect="auto", norm=matplotlib.colors.SymLogNorm(10**-5.0))
    # plt.show()
    # exit()


    if reshaped:
        derivative = np.concatenate(
            (derivative[..., 0], derivative[..., 1]),
            axis=-1
        )

    return derivative


def gradients(a, grid):

    gradient_y = np.gradient(a.in_2d, grid.in_2d[:, 0, 0], axis=0)
    gradient_x = np.gradient(a.in_2d, grid.in_2d[0, :, 1], axis=1)

    return (gradient_y, gradient_x)


def deflections_via_potential_from_grid(potential, grid):

    return gradients(a=potential, grid=grid)


def gradients_from_galaxy(galaxy, grid):
    raise ValueError(
        "This has not been implemented yet."
    )

    # if galaxy.has_light_profile:
    #     pass
    # else:
    #     pass

def replace_light_profiles_in_tracer_with_pixelization(tracer, pixelization, regularization=None):

    galaxies_with_mass_profiles = galaxies_with_mass_profiles_from_tracer(
        tracer=tracer
    )

    planes_with_light_profile = tracer.planes_with_light_profile

    if not planes_with_light_profile:
        raise ValueError("...")
    elif len(planes_with_light_profile) == 1:
        redshift = planes_with_light_profile[0].redshift
    elif len(planes_with_light_profile) > 1:
        raise ValueError(
            "This has not been implemented yet"
        )
    else:
        raise ValueError("...")

    return al.Tracer.from_galaxies(
        galaxies=[
            *galaxies_with_mass_profiles_from_tracer(
                tracer=tracer
            ),
            al.Galaxy(
                redshift=redshift,
                pixelization=pixelization,
                regularization=regularization if regularization is not None else al.reg.Constant(
                    coefficient=1.0
                ),
            )
        ]
    )


def plot_mapping_matrix(mapping_matrix, grid, nrows, ncols, figsize=None):


    def array_wrapper(array, grid):

        # NOTE: This has changed from the 0.45.0 version.
        return grid.mask.mapping.array_stored_2d_from_array_1d(
            array_1d=array
        )

    n = mapping_matrix.shape[-1]

    if nrows * ncols < n:
        raise ValueError("The total number of panels {} is less than {}".format(
                nrows * ncols, n
            )
        )

    figure, axes = plt.subplots(
        nrows=nrows,
        ncols=ncols,
        figsize=figsize if figsize is not None else (16, 8)
    )

    k = 0
    for i in range(nrows):
        for j in range(ncols):

            if k < n:
                axes[i, j].imshow(
                    array_wrapper(array=mapping_matrix[:, k], grid=grid)
                )
            else:
                axes[i, j].axis("off")

            axes[i, j].set_xticks([])
            axes[i, j].set_yticks([])

            k += 1

    plt.subplots_adjust(wspace=0.0, hspace=0.0)
    plt.show()

def mappers_of_planes_from_grid(tracer, source_redshift, mask, pixelization):
    raise ValueError(
        "This has not been implemented yet."
    )
    # tracer = al.Tracer.from_galaxies(
    #     galaxies=[
    #         *galaxies_with_mass_profiles_from_tracer(
    #             tracer=tracer
    #         ),
    #         al.Galaxy(redshift=source_redshift)
    #     ]
    # )
    #
    # grid = al.MaskedGrid.from_mask(mask=mask)
    #
    # # plt.figure()
    # # plt.plot(grid[:, 0], grid[:, 1], linestyle="None", marker="o")
    # # plt.show()
    # # exit()
    #
    # mappers_of_planes = []
    #
    # traced_grids_of_planes = tracer.traced_grids_of_planes_from_grid(
    #     grid=grid
    # )
    #
    # traced_sparse_grids_of_planes = tracer.traced_sparse_grids_of_planes_from_grid(
    #     grid=grid,
    #     preload_sparse_grids_of_planes=None
    # )
    #
    # for (plane_index, plane) in enumerate(tracer.planes):
    #     pass
    #     print(plane.galaxies)
    #     # print(plane_index)
    #     # if not plane.has_pixelization:
    #     #     mappers_of_planes.append(None)
    #     # # else:
    #     # #     mapper = plane.mapper_from_grid_and_sparse_grid(
    #     # #         grid=traced_grids_of_planes[plane_index],
    #     # #         sparse_grid=traced_sparse_grids_of_planes[plane_index],
    #     # #         inversion_uses_border=inversion_uses_border,
    #     # #     )
    #     # #     mappers_of_planes.append(mapper)
    #
    # #return mappers_of_planes


def planes_with_light_profiles(tracer):

    """

    1) Find all planes with light profiles.

    2) Get the galaxies of these planes.

    3) Get the light profiles of those galaxies.

    4) Make images of all those light profiles.

    """
    # NOTE: Find all planes with light profiles
    # NOTE:
    # # image = tracer.galaxies[1].profile_image_from_grid(grid=grid)
    # # plt.figure()
    # # plt.imshow(image.in_2d)
    # # plt.show()
    #
    # # asd = list(map(lambda plane: plane.has_light_profile, tracer.planes))
    # # print(asd)
    # #print(tracer.planes)
    #
    # #print(tracer.has_light_profile)
    # #print(list(map(lambda plane: plane.has_light_profile, tracer.planes)))
    # #print(tracer.galaxies_with_light_profile)
    #
    # #print(tracer.planes[1].galaxies_with_light_profile)
    #
    # galaxies = tracer.planes[1].galaxies_with_light_profile
    # galaxy = galaxies[0]
    #
    # galaxy_light_profiles = galaxy.light_profiles
    #
    # image_0 = galaxy_light_profiles[0].profile_image_from_grid(grid=grid)
    # image_0_in_2d = image_0.in_2d
    #
    # image_1 = galaxy_light_profiles[1].profile_image_from_grid(grid=grid)
    # image_1_in_2d = image_1.in_2d


# al.Galaxy(
#     redshift=src_redshift,
#     pixelization=pixelization,
#     regularization=al.reg.Constant(
#         coefficient=1000000.0
#     ),
# )


class inversion_settings:

    def __init__(self, method="operators"):

        self.method = method


class inversion_abstract:

    def __init__(self, Op, reconstruction, P, transformer, tracer, mapper, regularization_matrix, shape):

        self.Op = Op

        self._reconstruction = reconstruction

        self.P = P

        self.tracer = tracer
        self.mapper = mapper
        self.transformer = transformer

        self.regularization_matrix = regularization_matrix

        self.shape = shape


    @property
    def reconstruction(self):
        return self._reconstruction

    @reconstruction.setter
    def reconstruction(self, x):
         self._reconstruction = x


    def reshape(self, x):

        return x.reshape(self.shape)

    @property
    def mapped_reconstructed_image(self):

        array_1d = self.mapper.mapping_matrix @ self.reconstruction

        return self.mapper.grid.mapping.array_stored_1d_from_array_1d(
            array_1d=array_1d
        )

    @property
    def mapped_reconstructed_visibilities(self):

        return self.transformer.visibilities_from_image(
            image=self.mapped_reconstructed_image.in_2d
        )

        # return self.reshape(
        #     x=self.transformer._matvec(x=self.mapped_reconstructed_image)
        # )

    # @property
    # def model_data(self):
    #
    #     # return self.reshape(
    #     #     x=self.Op._matvec(x=self.reconstruction),
    #     # )
    #
    #     return self.reshape(
    #         x=self.Op * self.reconstruction,
    #     )

    @property
    def model_data(self):

        return self.mapped_reconstructed_visibilities


    def log_det(self, x):

        return np.linalg.slogdet(x)[1]


    @property
    def regularization_term(self):
        return np.matmul(
            self.reconstruction.T,
            np.matmul(self.regularization_matrix, self.reconstruction),
        )

    @property
    def log_det_curvature_reg_matrix_term(self):
        return self.log_det(x=self.P)

    @property
    def log_det_regularization_matrix_term(self):
        return self.log_det(x=self.regularization_matrix)



def wrap_mapping_matrix(mapping_matrix):

    return pylops.MatrixMult(
        sparse.bsr_matrix(mapper.mapping_matrix), dtype="float64"
    )

# def adsd(Op_1, Op_2):
#
#     return Op_1 * Op_2

def preconditioner_from(mapping_matrix, weights, regularization_matrix):
    pass


class AbstractMaskedDatasetSettings:
    def __init__(self):
        pass

class AbstractMaskedDataset:

    def __init__(self, data_1d, weights_1d):

        self.data_1d = data_1d
        self.weights_1d = weights_1d

def AbstractMaskedDatasetConverter(cls):

    # NOTE: Make it a dictionary wh
    attributes = ["data_1d", "weights_1d"]
    # attributes = {
    #     "data_1d":[], "weights_1d":[]
    # }
    for attr in attributes:
        if hasattr(cls, attr):
            pass
        else:
            if attr == "data_1":
                _attributes = []



class inversion_settings:

    def __init__(self, method="operators", use_preconditioner=True, preconditioner_method="powell"):

        self.method = method
        self.use_preconditioner = use_preconditioner
        self.preconditioner_method = preconditioner_method


def inversion_from_(
    masked_dataset,
    tracer,
    grid,
    transformer,
    transformer_from_averaged=None,
    inversion_uses_border=True,
    preload_sparse_grids_of_planes=None,
    settings=inversion_settings()
):


    def reconstruction_from_(masked_dataset, mapper, regularization_matrix, transformer, settings, transformer_from_averaged=None):

        mapping_matrix_as_operator = pylops.MatrixMult(
            sparse.bsr_matrix(mapper.mapping_matrix), dtype="float64"
        )

        Op = transformer * mapping_matrix_as_operator

        # Op = Op.toreal(
        #     forw=False, adj=True
        # )

        #print(settings.use_preconditioner, settings.preconditioner_method);exit()
        if settings.use_preconditioner:

            if settings.preconditioner_method == "approx":

                if transformer_from_averaged is None:
                    raise ValueError("This is a requirement for this method")

                start = time.time()
                transformed_mapping_matrices = transformer_from_averaged.transformed_mapping_matrices_from_mapping_matrix(
                    mapping_matrix=mapper.mapping_matrix
                )
                end = time.time()
                print(
                    "It took t={} to compute thetransformed mapping matrices".format(end - start)
                )

                F_real = np.dot(
                    transformed_mapping_matrices[0].T * masked_dataset.weights_averaged_from_reshaped[:, 0],
                    transformed_mapping_matrices[0]
                )
                F_imag = np.dot(
                    transformed_mapping_matrices[1].T * masked_dataset.weights_averaged_from_reshaped[:, 1],
                    transformed_mapping_matrices[1]
                )

                F_approx = F_real + F_imag

            else:

                F_approx = np.multiply(
                    np.sum(masked_dataset.weights_1d),
                    mapper.mapping_matrix.T @ mapper.mapping_matrix
                )


            P = np.add(F_approx, regularization_matrix)

            P_inv = np.linalg.inv(P)

            # # NOTE: converting the inverse of P to a diagonal matrix
            # P_inv = np.diag(np.diag(P_inv))

            # plt.figure()
            # plt.imshow(P_inv)
            # plt.show()
            # exit()

            # NOTE: The inverse of the P matrix is not sparse, but rather has very low value off-diagonal terms.
            MOp = pylops.MatrixMult(
                sparse.bsr_matrix(P_inv), dtype="float64"
            )

        else:

            MOp = None

        # values = pylops_utils.NormalEquationsInversion(
        #     Op=Op,
        #     data=masked_dataset.data_1d,
        #     WOp=pylops.Diagonal(
        #         diag=masked_dataset.weights_1d, dtype="float64"
        #     ),
        #     HOp=pylops.MatrixMult(
        #         sparse.bsr_matrix(regularization_matrix), dtype="float64"
        #     ),
        #     M=MOp,
        #     tol=1e-08,
        #     atol=1e-08,
        #     **dict(
        #         maxiter=mapper.mapping_matrix.shape[-1]
        #     )
        # )

        print(
            "Performing inversion ..."
        )

        start = time.time()
        values = pylops.NormalEquationsInversion(
            Op=Op,
            Regs=None,
            data=masked_dataset.data_1d,
            Weight=pylops.Diagonal(
                diag=masked_dataset.weights_1d, dtype="float64"
            ),
            epsNRs=[1.0],
            NRegs=[
                pylops.MatrixMult(
                    sparse.bsr_matrix(regularization_matrix), dtype="float64"
                )
            ],
            M=MOp,
            tol=1e-08
        )
        end = time.time()
        print(
            "It took t={} to perform the reconstruction".format(end - start)
        )

        # plt.figure()
        # autolens_plot_utils.draw_voronoi_pixels(mapper=mapper, values=values)
        # plt.show()
        # exit()

        return values, Op, P


    if tracer.has_pixelization:
        mappers_of_planes = tracer.mappers_of_planes_from_grid(
            grid=grid,
            inversion_uses_border=inversion_uses_border,
            preload_sparse_grids_of_planes=preload_sparse_grids_of_planes,
        )

        mapper = mappers_of_planes[-1]

        regularization = tracer.regularizations_of_planes[-1]
        # print(
        #     "coefficient = ", regularization.coefficient
        # )

        regularization_matrix = regularization.regularization_matrix_from_mapper(
            mapper=mapper
        )

        # plt.figure()
        # plt.imshow(regularization_matrix)
        # plt.show()
        # exit()

        start = time.time()
        reconstruction, Op, P = reconstruction_from_(
            masked_dataset=masked_dataset,
            mapper=mapper,
            regularization_matrix=regularization_matrix,
            transformer=transformer,
            settings=settings,
            transformer_from_averaged=transformer_from_averaged,
        )
        end = time.time()
        print(
            "It took t={} to perform the inversion".format(end - start)
        )

    else:
        raise ValueError(
            "This has not been implemented yet"
        )

    return inversion_abstract(
        Op=Op,
        reconstruction=reconstruction,
        P=P,
        mapper=mapper,
        tracer=tracer,
        transformer=transformer,
        regularization_matrix=regularization_matrix,
        shape=masked_dataset.data_shape,
    )


def redshifts_from_list_of_galaxies(galaxies):

    return [galaxy.redshift for galaxy in galaxies]


def optimize_(masked_dataset, transformer, lens_galaxies, source_redshift, pixelization, x0=0.0, regularization_class=al.reg.Constant):

    import autoarray_classes.autoarray_fit_classes as autoarray_fit_classes

    class optimization:

        def __init__(self, source_redshift, pixelization, lens_galaxies, regularization_class):

            self.source_redshift = source_redshift

            # NOTE:
            self.pixelization = pixelization

            if all(redshift < self.source_redshift
                for redshift in redshifts_from_list_of_galaxies(galaxies=lens_galaxies)
            ):
                self.lens_galaxies = lens_galaxies
            else:
                raise ValueError("...")

            self.regularization_class = regularization_class


        def source_galaxy(self, coefficient):

            return al.Galaxy(
                redshift=self.source_redshift,
                pixelization=self.pixelization,
                regularization=self.regularization_class(
                    coefficient=coefficient
                ),
            )

        def tracer_from(self, coefficient):

            return al.Tracer.from_galaxies(
                galaxies=[
                    *self.lens_galaxies, self.source_galaxy(coefficient=coefficient)
                ]
            )



    def func(log_coefficient, masked_dataset, grid, transformer, optimization_class):

        coefficient = 10.0**log_coefficient[0];print(log_coefficient)

        inversion = inversion_from_(
            masked_dataset=masked_dataset,
            grid=grid,
            transformer=transformer,
            tracer=optimization_class.tracer_from(
                coefficient=coefficient
            ),
        )

        fit = autoarray_fit_classes.FitInterferometer(
            masked_dataset=masked_dataset,
            model_data=inversion.model_data,
            inversion=inversion
        )

        return -fit.evidence

    # NOTE:
    # if isinstance(regularization_class, al.reg.Constant):
    #     print("OK")
    # else:
    #     raise TypeError(
    #         "This has not been implemented yet"
    #     )



    optimization_class = optimization(
        source_redshift=source_redshift,
        pixelization=pixelization,
        lens_galaxies=lens_galaxies,
        regularization_class=regularization_class
    )

    grid = al.Grid.from_mask(
        mask=transformer.real_space_mask
    )

    # log_coefficient_arr = np.linspace(-3.0, 8.0, 20)
    # evidences = np.zeros(shape=log_coefficient_arr.shape)
    # for i, log_coefficient in enumerate(log_coefficient_arr):
    #
    #     evidence = func(
    #         log_coefficient=log_coefficient,
    #         masked_dataset=masked_dataset,
    #         grid=grid,
    #         transformer=transformer,
    #         optimization_class=optimization_class
    #     )
    #
    #     evidences[i] = evidence
    #
    # plt.figure()
    # plt.plot(log_coefficient_arr, evidences)
    # plt.show()
    # exit()

    res = optimize.minimize(
        func,
        x0=x0,
        args=(masked_dataset, grid, transformer, optimization_class),
        method='COBYLA'
    )
    #print(res)

    return res.x
