import sys
import numpy as np

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

from workspace_jam.plotting.density_rj_lens_plots import density_plot_tools

slacs = density_plot_tools.SLACS()

image_dir = (
    "/gpfs/data/pdtw24/PL_Data/RJLens3/"
)  # Dir of Object to make evidence tables from
image_name = "Abell 1201"
file_name = "density_decomposed_aligned"
title = "SIE + Decomposed (Aligned) Density Profiles of "

number_bins = 50

## SIE ##

values = []
values.append(0.05247720)  # Lens x
values.append(0.03497244)  # Lens y
values.append(1.91716393)  # Lens einR
values.append(0.68532604)  # Lens axis ratio
values.append(66.83060729)  # Lens phi

sie_mass_ein_r = slacs.masses_of_all_samples_sie(
    values=values, radius_kpc=slacs.einstein_radius * slacs.kpc_per_arcsec
)
sie_mass_475 = slacs.masses_of_all_samples_sie(values=values, radius_kpc=4.75)
slacs.density_vs_radii_sie(radius_kpc=15.0, values=values, number_bins=number_bins)

### Decomposed Aligned model ###

center_skip = 0
ltm_skip = 0

values = []
values.append(0.02191929)  # fg 1 x
values.append(0.02982811)  # fg 1 y
values.append(0.19665282)  # fg 1 intensity
values.append(0.55871552)  # fg 1 eff rad
values.append(1.39390583)  # fg 1 sersic inex
values.append(0.83443597)  # fg 1 axis ratio
values.append(66.48069170)  # fg 1 phi
values.append(0.03447566)  # fg 2 intensity
values.append(4.52134788)  # fg 2 effective radius
values.append(0.61310712)  # fg 2 axis ratio
values.append(66.48069170)  # fg 2 phi
values.append(0.04173277)  # Lens Kappa s
values.append(2.71727396)  # Lens MLR 1
values.append(5.67223821)  # Lens MLR 2

total_mass, stellar_mass, bulge_mass, envelope_mass, dark_mass = slacs.masses_of_all_samples_ltm2(
    radius_kpc=slacs.einstein_radius * slacs.kpc_per_arcsec,
    values=values,
    center_skip=center_skip,
    ltm_skip=ltm_skip,
)

print("Masses as R_Ein (kpc) = ", slacs.einstein_radius * slacs.kpc_per_arcsec)
print()
print("SIE Mass = {:.3E}".format(sie_mass_ein_r))
print("Total Decomposed Mass = {:.3E}".format(total_mass))
print("Stellar Mass = {:.3E}".format(stellar_mass))
print("Dark Mass = {:.3E}".format(dark_mass))
print("Bulge Mass = {:.3E}".format(bulge_mass))
print("Envelope Mass = {:.3E}".format(envelope_mass))

total_mass, stellar_mass, bulge_mass, envelope_mass, dark_mass = slacs.masses_of_all_samples_ltm2(
    radius_kpc=4.75, values=values, center_skip=center_skip, ltm_skip=ltm_skip
)

print("Masses as 4.75 kpc")
print()
print("SIE Mass = {:.3E}".format(sie_mass_475))
print("Total Decomposed Mass = {:.3E}".format(total_mass))
print("Stellar Mass = {:.3E}".format(stellar_mass))
print("Dark Mass = {:.3E}".format(dark_mass))
print("Bulge Mass = {:.3E}".format(bulge_mass))
print("Envelope Mass = {:.3E}".format(envelope_mass))

slacs.density_vs_radii_ltm2(
    radius_kpc=10.0, values=values, center_skip=0, ltm_skip=1, number_bins=number_bins
)
slacs.plot_density(
    image_name=image_name,
    file_name=file_name,
    title=title,
    labels=["SIE", "Sersic", "Exponential", "NFWSph"],
)
