from autolens.array import grids
from autolens.model.profiles import mass_profiles as mp
from autolens.model.profiles.plotters import profile_plotters

import numpy as np

grid = grids.Grid.from_shape_pixel_scale_and_sub_grid_size(
    shape=(800, 800), pixel_scale=0.1
)

nfw_challenge = al.mass_profiles.SphericalTruncatedNFWChallenge(
    kappa_s=0.009, scale_radius=0.5
)
mass_at_200 = nfw_challenge.mass_at_200_for_units(
    critical_surface_density=1940654909.4133248,
    cosmic_average_density=262.30319684750657,
)
concentration = nfw_challenge.concentration_for_units(
    critical_surface_density=1940654909.4133248,
    cosmic_average_density=262.30319684750657,
)


print("{:.4e}".format(mass_at_200))
print(concentration)

deflections = nfw_challenge.deflections_from_grid(grid=grid)
print("defl y max = " + str(np.max(deflections[:, 0])))
print("defl x max = " + str(np.max(deflections[:, 1])))
print(
    "defl r max = "
    + str(np.max(np.sqrt(deflections[:, 1] ** 2.0 + deflections[:, 0] ** 2.0)))
)

profile_plotters.plot_deflections_x(mass_profile=nfw_challenge, grid=grid)
profile_plotters.plot_deflections_y(mass_profile=nfw_challenge, grid=grid)
