from autolens.data.instrument import abstract_data
from autolens.array import grids
from autolens.lens import ray_tracing
from autolens.model.galaxy import galaxy as g
from autolens.model.profiles import light_profiles as lp
from autolens.model.profiles import mass_profiles as mp

import matplotlib.pyplot as plt
from workspace_jam.helping.qiuhan.deflection_map import deflection_multi_map
from astropy.cosmology import FlatLambdaCDM


pixel_scale_img = 0.05
npixel_img = int(10 / pixel_scale_img)

H0 = 100.0
OM0 = 0.3
cosmo = FlatLambdaCDM(H0=H0, Om0=OM0)

psf = abstract_data.PSF.from_gaussian(
    shape=(11, 11), sigma=0.1, pixel_scale=pixel_scale_img
)

image_plane_grid = grids.Grid.from_shape_pixel_scale_and_sub_grid_size(
    shape=(200, 200), pixel_scale=pixel_scale_img, psf_shape=psf.shape, sub_grid_size=1
)

lens_galaxy = g.Galaxy(
    mass=al.mass_profiles.SphericalIsothermal(centre=(3.0, 3.0), einstein_radius=1.6),
    redshift=0.5,
)

source_galaxy = g.Galaxy(
    light=al.light_profiles.EllipticalSersic(
        centre=(0.1, 0.0),
        axis_ratio=0.8,
        phi=0.0,
        intensity=1.0,
        effective_radius=1.0,
        sersic_index=2.5,
    ),
    redshift=1.0,
)

# "zero kappa" Mass Sheet
mass_sheet = g.Galaxy(
    sheet=al.mass_profiles.MassSheet(centre=(0.0, 0.0), kappa=0.0), redshift=0.2
)


# Below, all variables with "_test" means inpluding the mass sheet and all
# variables with "_ref" means they don't have a mass sheet

lens_list_ref = [lens_galaxy, source_galaxy]

tracer_ref = ray_tracing.Tracer(
    galaxies=lens_list_ref, image_plane_grid=image_plane_grid, cosmology=cosmo
)

################################################################################
# Here, I use deflection_multi_map function to calculate the deflection angles
# in multiplane situations. This function will return an N*N*2 array with
# [:, :, 1] to be deflections in x direction and [:, :, 0] to be deflections in
# y direction.

alpha_ref = deflection_multi_map(
    N=npixel_img, map_scale=pixel_scale_img, tracer=tracer_ref, z_lens=0.5, z_source=1.0
)

alphax_ref = alpha_ref[:, :, 1]
alphay_ref = alpha_ref[:, :, 0]

lens_list_test = [lens_galaxy, mass_sheet, source_galaxy]

tracer_test = ray_tracing.Tracer(
    galaxies=lens_list_test, image_plane_grid=image_plane_grid, cosmology=cosmo
)

alpha_test = deflection_multi_map(
    N=npixel_img,
    map_scale=pixel_scale_img,
    tracer=tracer_test,
    z_lens=0.5,
    z_source=1.0,
)

alphax_test = alpha_test[:, :, 1]
alphay_test = alpha_test[:, :, 0]


fig, axes = plt.subplots(1, 2, figsize=(8, 4))
axes[0].imshow(alphax_ref)
axes[0].set_title("no mass sheet")
axes[1].imshow(alphax_test)
axes[1].set_title("with mass sheet")
plt.show()
