import sys
import math
import numpy as np
import matplotlib.pyplot as plt

sys.path.append("../")

from paper_plots import density_plot_tools
from autolens.analysis import galaxy
from autolens.model.profiles import mass_profiles


ltm_skip = 0
center_skip = 0

image_dir = (
    "/gpfs/data_vector/pdtw24/PL_Data/SL03_2/"
)  # Dir of Object to make evidence tables from

image_name = "SLACSJ0252+0039"
image_name = "SLACSJ1250+0523"
# image_name = 'SLACSJ1430+4105'

slacs = density_plot_tools.SLACS(image_dir, image_name)

# pipeline_folder = 'Pipeline_Total'
# phase_folder = 'PL_Phase_8'

# pipeline_folder = 'Pipeline_LMDM'
# phase_folder = 'PL_Phase_7'

pipeline_folder = "Pipeline_LTM2"
phase_folder = "PL_5_LMDM_LTM_R_MC"

sersic_bulge = mass_profiles.EllipticalSersicMass(
    centre=(0.02, 0.008),
    axis_ratio=0.87,
    phi=71.1,
    intensity=0.09458,
    effective_radius=1.0153,
    sersic_index=4.57468,
    mass_to_light_ratio=1.55265,
)

exponential_halo = mass_profiles.EllipticalSersicMassRadialGradient(
    centre=(0.043, 0.036),
    axis_ratio=0.7861,
    phi=161.34,
    intensity=0.02475,
    effective_radius=2.1389,
    mass_to_light_ratio=1.2556,
    mass_to_light_gradient=-0.2229,
)

dark_matter_halo = mass_profiles.SphericalNFW(
    kappa_s=0.1035, scale_radius=30.0 * slacs.arcsec_per_kpc
)

lens_galaxy = galaxy.Galaxy(
    redshift=slacs.redshift,
    mass_profiles=[sersic_bulge, exponential_halo, dark_matter_halo],
)
source_galaxy = galaxy.Galaxy(redshift=slacs.source_redshift)

number_bins = 30
radius_kpc = 30

radii = list(np.linspace(5e-3, radius_kpc * slacs.arcsec_per_kpc, number_bins + 1))

radii_plot = []
bulge_densities = []
halo_densities = []
dark_densities = []

for r in range(number_bins):

    annuli_area = math.pi * radii[r + 1] ** 2 - math.pi * radii[r] ** 2

    densities = (
        lens_galaxy.dimensionless_mass_within_circles_individual(radii[r + 1])
        - lens_galaxy.dimensionless_mass_within_circles_individual(radii[r])
    ) / annuli_area

    bulge_densities.append(densities[0])
    halo_densities.append(densities[1])
    dark_densities.append(densities[2])

    radii_plot.append(((radii[r + 1] + radii[r]) / 2.0) * slacs.kpc_per_arcsec)

print(bulge_densities)

plt.semilogy(radii_plot, bulge_densities, color="r", label="Sersic Bulge")
plt.semilogy(
    radii_plot, halo_densities, color="g", label="EllipticalExponentialMass Halo"
)
plt.semilogy(radii_plot, dark_densities, color="k", label="Dark Matter Halo")

plt.legend()
plt.show()

# slacs.weighted_densities_vs_radii(radius_kpc=30.0, weight_cut=1e-3, number_bins=20)

# slacs.plot_density(image_name=image_name, labels=['Sersic', 'Exponential', 'NFWSph'])
