# %%
"""
Tutorial 3: Inversions
======================

We've covered `Mapper``., which, if I haven`t emphasised it enough yet, map things. Now, we`re going to look at how we
can use these `Mapper`'s (which map things) to reconstruct the source galaxy - I hope you`re excited!
"""

# %%
#%matplotlib inline

from pyprojroot import here

workspace_path = str(here())
#%cd $workspace_path
print(f"Working Directory has been set to `{workspace_path}`")

from os import path
import autolens as al
import autolens.plot as aplt


# %%
"""
we'll use the same strong lensing data as the previous tutorial, where:

 - The lens `Galaxy`'s light is omitted.
 - The lens `Galaxy`'s total mass distribution is an `EllipticalIsothermal`.
 - The source `Galaxy`'s `LightProfile` is an `EllipticalSersic`.
"""

# %%
dataset_name = "mass_sie__source_sersic"
dataset_path = path.join("dataset", "howtolens", "chapter_4", dataset_name)

imaging = al.Imaging.from_fits(
    image_path=path.join(dataset_path, "image.fits"),
    noise_map_path=path.join(dataset_path, "noise_map.fits"),
    psf_path=path.join(dataset_path, "psf.fits"),
    pixel_scales=0.1,
)

# %%
"""
Lets create an annular `Mask2D` which traces the stongly lensed source ring.
"""

# %%
mask = al.Mask2D.circular_annular(
    shape_2d=imaging.shape_2d,
    pixel_scales=imaging.pixel_scales,
    sub_size=1,
    inner_radius=0.5,
    outer_radius=2.8,
)

aplt.Imaging.image(imaging=imaging, mask=mask)

# %%
"""
Next, lets set the `Imaging` and `Mask2D` up as a `MaskedImaging` object and setup a `Tracer` using the input lens 
galaxy model (we don't need to provide the source`s `LightProfile`, as we`re using a `Mapper` to reconstruct it).
"""

# %%
masked_imaging = al.MaskedImaging(
    imaging=imaging, mask=mask, settings=al.SettingsMaskedImaging(sub_size=2)
)

lens_galaxy = al.Galaxy(
    redshift=0.5,
    mass=al.mp.EllipticalIsothermal(
        centre=(0.0, 0.0), elliptical_comps=(0.1, 0.0), einstein_radius=1.6
    ),
)

tracer = al.Tracer.from_galaxies(galaxies=[lens_galaxy, al.Galaxy(redshift=1.0)])

source_plane_grid = tracer.traced_grids_of_planes_from_grid(grid=masked_imaging.grid)[1]

# %%
"""
we'll use another rectangular `Pixelization` and `Mapper` to perform the reconstruction.
"""

# %%
rectangular = al.pix.Rectangular(shape=(25, 25))

mapper = rectangular.mapper_from_grid_and_sparse_grid(grid=source_plane_grid)

aplt.Mapper.subplot_image_and_mapper(
    image=imaging.image,
    mapper=mapper,
    include=aplt.Include(mask=True, inversion_grid=True),
)

# %%
"""
And now, finally, we`re going to use our `Mapper` to invert the image using an `Inversion`. I`ll explain how this 
works in a second - but lets just go ahead and use the `Inversion` first. (Ignore the `regularization` input below for 
now, we'll cover this in the next tutorial).
"""

# %%
inversion = al.Inversion(
    masked_dataset=masked_imaging,
    mapper=mapper,
    regularization=al.reg.Constant(coefficient=1.0),
)

# %%
"""
Our `Inversion` has a reconstructed image and `Pixeilzation`, whcih we can plot using an `Inversion` plotter
"""

# %%
aplt.Inversion.reconstructed_image(inversion=inversion, include=aplt.Include(mask=True))

aplt.Inversion.reconstruction(
    inversion=inversion, include=aplt.Include(inversion_grid=True)
)

# %%
"""
And there we have it, we've successfully reconstructed, or, *inverted*, our source using the mapper`s rectangular 
grid. Whilst this source was simple (a blob of light in the centre of the source-plane), `Inversion`'s come into their 
own when fitting sources with complex morphologies. Infact, given we`re having so much fun inverting things, lets 
invert a really complex source!
"""

# %%
dataset_name = "mass_sie__source_sersic_x5"
dataset_path = path.join("dataset", "howtolens", "chapter_4", dataset_name)

imaging = al.Imaging.from_fits(
    image_path=path.join(dataset_path, "image.fits"),
    noise_map_path=path.join(dataset_path, "noise_map.fits"),
    psf_path=path.join(dataset_path, "psf.fits"),
    pixel_scales=0.1,
)

aplt.Imaging.image(imaging=imaging)

# %%
"""
This code is doing all the the same as above (setup the `Mask2D`, `Galaxy`'s `Tracer`, `Mapper`, ec.).
"""

# %%

mask = al.Mask2D.circular_annular(
    shape_2d=imaging.shape_2d,
    pixel_scales=imaging.pixel_scales,
    sub_size=1,
    inner_radius=0.1,
    outer_radius=3.2,
)

aplt.Imaging.image(imaging=imaging, mask=mask)

masked_imaging = al.MaskedImaging(imaging=imaging, mask=mask)

tracer = al.Tracer.from_galaxies(galaxies=[lens_galaxy, al.Galaxy(redshift=1.0)])

source_plane_grid = tracer.traced_grids_of_planes_from_grid(grid=masked_imaging.grid)[1]

mapper = rectangular.mapper_from_grid_and_sparse_grid(grid=source_plane_grid)

inversion = al.Inversion(
    masked_dataset=masked_imaging,
    mapper=mapper,
    regularization=al.reg.Constant(coefficient=1.0),
)

# %%
"""
Lets inspect the complex source reconstruction.
"""

# %%
aplt.Inversion.reconstructed_image(inversion=inversion, include=aplt.Include(mask=True))

aplt.Inversion.reconstruction(
    inversion=inversion, include=aplt.Include(inversion_grid=True)
)


# %%
"""
Pretty great, huh? If you ran the complex source pipeline, you`ll remember that getting a model image that looked that 
good simply *was not possible*. With an `Inversion`, we can do it with ease and without fitting 30+ parameters!

Lets discuss how an `Inversion` actually works. The explanation I give below is overly-simplified. I'm avoiding the 
technical details of how an `Inversion` *actually* works. To be good at lens modeling you don't need to understand the 
nitty-gritty details of linear inversions, you just need an instinct for how to use them as a tool to model lenses.

Nevertheless, I know a lot of you hate `black-boxes`, or have an interest in linear algrebra. If you`re that way 
inclined, then checkout the documentation of the autolens source code for more information. In particular, you should 
look at the following functions in the project PyAutoArray:

autoarray.inversions.mappers.mapping_matrix
autoarray.opterators.convolution.convolve_mapping_matrix
autoarray.opterators.inversions.regularization.Regularization
autoarray.opterators.inversions.inversions.Inversion

To begin, lets consider some random mappings between our mapper`s source-pixels and the image.
"""

# %%
aplt.Mapper.subplot_image_and_mapper(
    image=masked_imaging.image,
    mapper=mapper,
    include=aplt.Include(mask=True, inversion_grid=True),
    source_pixel_indexes=[[445], [285], [313], [132], [11]],
)

# %%
"""
These mappings are known before the `Inversion`, which means pre-inversion we know two key pieces of information:

 1) The mappings between every source-pixel and sets of image-pixels.
 2) The flux values in every observed image-pixel, which are the values we want to fit successfully.

It turns out that with these two pieces of information we can linearly solve for the set of source-pixel fluxes that 
best-fit (e.g. maximize the log likelihood of) our observed image. Essentially, we set up the mapping between source and 
image pixels as a large matrix and solve for the source-pixel fluxes in an analogous fashion to how you would solve a 
set of simultaneous linear equations. This process is called a `linear inversion`.

There are three more things about a linear `Inversion` that are worth knowing:

 1) We've discussed the image sub-grid before, which splits each image-pixel into a sub-pixel. If a sub-grid is 
 used, it is the mapping between every sub-pixel and source-pixel that is computed and used to perform the 
 `Inversion`. This prevents aliasing effects degrading the image reconstruction, and, as a rule of thumb, I 
 would suggest you use sub-gridding of degree 2x2.

 2) When fitting using `LightProfile`'s we discussed how a `model_image` was generated by blurring them with the 
 data`s PSF. A similar blurring operation is incorporated into the `Inversion`, such that the reconstructed 
 image and source fully account for the telescope optics and effect of the PSF.

 3) The inversion`s solution is regularized. But wait, that`s what we'll cover in the next tutorial!

Finally, let me show you how easy it is to fit an image with an `Inversion` using a `FitImaging` oboject. Instead of 
giving the source galaxy a `LightProfile`, we give it a `Pixelization`.nd `Regularization`, and pass it to a `Tracer`.
"""

# %%
source_galaxy = al.Galaxy(
    redshift=1.0,
    pixelization=al.pix.Rectangular(shape=(40, 40)),
    regularization=al.reg.Constant(coefficient=1.0),
)

tracer = al.Tracer.from_galaxies(galaxies=[lens_galaxy, source_galaxy])

# %%
"""
Then, like before, we pass the `MaskedImaging` and `Tracer` to a `FitImaging` object. Indeed, we see some 
pretty good looking residuals - we`re certainly fitting the lensed source accurately!
"""

# %%
fit = al.FitImaging(masked_imaging=masked_imaging, tracer=tracer)

aplt.FitImaging.subplot_fit_imaging(fit=fit, include=aplt.Include(mask=True))

# %%
"""
And, we`re done, here are a few questions to get you thinking about `Inversion``.:

 1) The `Inversion` provides the maximum log likelihood solution to the observed image. Is there a problem with seeking 
 the `best-fit`? Is there a risk that we`re going to fit other things in the image than just the lensed source 
 galaxy? What happens if you reduce the `regularization_coefficient` above to zero?

 2) The exterior pixels in the `Rectangular` `Grid`.have no image-pixels in them. However, they are still given a 
 reconstructed flux. If this value isn't` coming from a util to an image-pixel, where is it be coming from?
"""
