import getdist
import re
import os
import math
import numpy as np
from astropy import cosmology as cosmo
import matplotlib.pyplot as plt

from autolens.model.galaxy import galaxy
from autolens.model.profiles import light_profiles
from autolens.model.profiles import mass_profiles


def weighted_avg_and_std(values, weights):
    """
    Return the weighted average and standard deviation.

    values, weights -- Numpy ndarrays with the same shape.
    """
    average = np.average(values, weights=weights)
    # Fast and numerically precise:
    variance = np.average((values - average) ** 2, weights=weights)
    return (average, math.sqrt(variance))


def getModelFolders(dir):
    regex = re.compile(".*Hyper")
    folders = [
        name
        for name in os.listdir(dir)
        if os.path.isdir(dir + name) and re.match(regex, name) == None
    ]
    return folders


class SLACS(object):
    def __init__(self, image_dir, image_name):

        self.image_dir = image_dir

        self.image_name = image_name

        cosmological_model = cosmo.LambdaCDM(H0=70, Om0=0.3, Ode0=0.7)

        if image_name == "SLACSJ0252+0039":

            self.redshift = 0.2803
            self.source_redshift = 0.9818
            self.einstein_radius = 1.02
            self.source_light_min = 0.8
            self.source_light_max = 1.3
            self.ein_r_kpc = 4.18
            self.mass = "{0:e}".format(10 ** 11.25)
            self.radius = 4.4
            self.chab_stellar_mass = "{0:e}".format(10 ** 11.21)
            self.chab_stellar_mass_lower = "{0:e}".format(10 ** 11.08)
            self.chab_stellar_mass_upper = "{0:e}".format(10 ** 11.34)
            self.sal_stellar_mass = "{0:e}".format(10 ** 11.46)
            self.sal_stellar_mass_lower = "{0:e}".format(10 ** 11.33)
            self.sal_stellar_mass_upper = "{0:e}".format(10 ** 11.59)
            self.chab_stellar_frac = 0.4
            self.chab_stellar_frac_error = 0.12
            self.sal_stellar_frac = 0.71
            self.sal_stellar_frac_error = 0.21
            self.critical_density_arcsec = 53947971687.96673

        elif image_name == "SLACSJ1250+0523":

            self.redshift = 0.2318
            self.source_redshift = 0.7953
            self.einstein_radius = 1.120
            self.source_light_min = 0.5
            self.source_light_max = 1.5
            self.ein_r_kpc = 4.18
            self.mass = "{0:e}".format(10 ** 11.26)
            self.chab_stellar_mass = "{0:e}".format(10 ** 11.53)
            self.chab_stellar_mass_lower = "{0:e}".format(10 ** 11.46)
            self.chab_stellar_mass_upper = "{0:e}".format(10 ** 11.60)
            self.sal_stellar_mass = "{0:e}".format(10 ** 11.77)
            self.sal_stellar_mass_lower = "{0:e}".format(10 ** 11.7)
            self.sal_stellar_mass_upper = "{0:e}".format(10 ** 12.84)
            self.chab_stellar_frac = 0.68
            self.chab_stellar_frac_error = 0.11
            self.sal_stellar_frac = 1.2
            self.sal_stellar_frac_error = 0.19
            self.critical_density_arcsec = 46542666435.694626

        elif image_name == "SLACSJ1430+4105":

            self.redshift = 0.2850
            self.source_redshift = 0.5753
            self.einstein_radius = 1.468
            self.source_light_min = 0.2
            self.source_light_max = 2.3
            self.our_einr_arc = 1.4668492852
            self.ein_r_kpc = 6.53
            self.mass = "{0:e}".format(10 ** 11.73)
            self.chab_stellar_mass = "{0:e}".format(10 ** 11.68)
            self.chab_stellar_mass_lower = "{0:e}".format(10 ** 11.56)
            self.chab_stellar_mass_upper = "{0:e}".format(10 ** 11.8)
            self.sal_stellar_mass = "{0:e}".format(10 ** 11.93)
            self.sal_stellar_mass_lower = "{0:e}".format(10 ** 11.82)
            self.sal_stellar_mass_upper = "{0:e}".format(10 ** 12.04)
            self.chab_stellar_frac = 0.33
            self.chab_stellar_frac_error = 0.09
            self.sal_stellar_frac = 0.59
            self.sal_stellar_frac_error = 0.16
            self.critical_density_arcsec = 76974206518.66878

        self.arcsec_per_kpc = cosmological_model.arcsec_per_kpc_proper(
            self.redshift
        ).value
        self.kpc_per_arcsec = 1.0 / self.arcsec_per_kpc

    def load_samples(self, pdf_file, center_skip, ltm_skip):

        self.pdf = getdist.mcsamples.loadMCSamples(pdf_file)
        self.params = self.pdf.paramNames

        self.center_skip = center_skip
        self.ltm_skip = ltm_skip

    def setup_lensing_plane(self, values):

        sersic_bulge_light = light_profiles.EllipticalSersic(
            centre=(values[0], values[1]),
            axis_ratio=values[5],
            phi=values[6],
            intensity=values[2],
            effective_radius=values[3],
            sersic_index=values[4],
        )

        exponential_halo_light = light_profiles.EllipticalExponential(
            axis_ratio=values[9 + self.center_skip],
            phi=values[10 + self.center_skip],
            intensity=values[7 + self.center_skip],
            effective_radius=values[8 + self.center_skip],
        )

        sersic_bulge = mass_profiles.EllipticalSersic(
            centre=(values[0], values[1]),
            axis_ratio=values[5],
            phi=values[6],
            intensity=values[2],
            effective_radius=values[3],
            sersic_index=values[4],
            mass_to_light_ratio=values[12 + self.center_skip],
        )

        exponential_halo = mass_profiles.EllipticalExponential(
            axis_ratio=values[9 + self.center_skip],
            phi=values[10 + self.center_skip],
            intensity=values[7 + self.center_skip],
            effective_radius=values[8 + self.center_skip],
            mass_to_light_ratio=values[12 + self.ltm_skip + self.center_skip],
        )

        dark_matter_halo = mass_profiles.SphericalNFW(
            kappa_s=values[11 + self.center_skip],
            scale_radius=30.0 * self.arcsec_per_kpc,
        )

        self.lens_galaxy = galaxy.Galaxy(
            redshift=self.redshift,
            light1=sersic_bulge_light,
            light2=exponential_halo_light,
            bulge=sersic_bulge,
            disk=exponential_halo,
            halo=dark_matter_halo,
        )

        self.source_galaxy = galaxy.Galaxy(redshift=self.source_redshift)

    def setup_lensing_plane_ltm2r(self, values):

        sersic_bulge_light = light_profiles.EllipticalSersic(
            centre=(values[0], values[1]),
            axis_ratio=values[5],
            phi=values[6],
            intensity=values[2],
            effective_radius=values[3],
            sersic_index=values[4],
        )

        exponential_halo_light = light_profiles.EllipticalExponential(
            axis_ratio=values[9 + self.center_skip],
            phi=values[10 + self.center_skip],
            intensity=values[7 + self.center_skip],
            effective_radius=values[8 + self.center_skip],
        )

        sersic_bulge = mass_profiles.EllipticalSersic(
            centre=(values[0], values[1]),
            axis_ratio=values[5],
            phi=values[6],
            intensity=values[2],
            effective_radius=values[3],
            sersic_index=values[4],
            mass_to_light_ratio=values[12 + self.center_skip],
        )

        exponential_halo = mass_profiles.EllipticalSersicRadialGradient(
            axis_ratio=values[9 + self.center_skip],
            phi=values[10 + self.center_skip],
            intensity=values[7 + self.center_skip],
            effective_radius=values[8 + self.center_skip],
            sersic_index=1.0,
            mass_to_light_ratio=values[12 + self.ltm_skip + self.center_skip],
            mass_to_light_gradient=values[13 + self.ltm_skip + self.center_skip],
        )

        dark_matter_halo = mass_profiles.SphericalNFW(
            kappa_s=values[11 + self.center_skip],
            scale_radius=30.0 * self.arcsec_per_kpc,
        )

        self.lens_galaxy = galaxy.Galaxy(
            redshift=self.redshift,
            light1=sersic_bulge_light,
            light2=exponential_halo_light,
            bulge=sersic_bulge,
            disk=exponential_halo,
            halo=dark_matter_halo,
        )

        self.source_galaxy = galaxy.Galaxy(redshift=self.source_redshift)

    def setup_model(self, model_index, ltmr2=False):

        values = self.pdf.samples[model_index]

    def masses_of_all_samples(self, radius_kpc):

        model_indexes = []
        sample_weights = []
        total_masses = []
        stellar_masses = []
        bulge_masses = []
        disk_masses = []
        dark_masses = []
        stellar_fractions = []

        for i in range(len(self.pdf.samples)):

            if self.pdf.weights[i] > 1e-6:

                model_indexes.append(i)

                values = self.pdf.samples[i]

                self.setup_lensing_plane_ltm2r(values)

                sample_weights.append(self.pdf.weights[i])

                radius_arc = radius_kpc * self.arcsec_per_kpc
                bulge_mass = (
                    self.critical_density_arcsec
                    * self.lens_galaxy.bulge.mass_within_circle_in_units(
                        radius=radius_arc
                    )
                )
                disk_mass = (
                    self.critical_density_arcsec
                    * self.lens_galaxy.disk.mass_within_circle_in_units(
                        radius=radius_arc
                    )
                )
                halo_mass = (
                    self.critical_density_arcsec
                    * self.lens_galaxy.halo.mass_within_circle_in_units(
                        radius=radius_arc
                    )
                )

                total_mass = bulge_mass + disk_mass + halo_mass
                stellar_mass = bulge_mass + disk_mass
                dark_mass = halo_mass

                total_masses.append(total_mass)
                stellar_masses.append(stellar_mass)
                bulge_masses.append(bulge_mass)
                disk_masses.append(disk_mass)
                dark_masses.append(dark_mass)
                stellar_fractions.append(stellar_mass / total_mass)

        return (
            model_indexes,
            sample_weights,
            total_masses,
            stellar_masses,
            bulge_masses,
            disk_masses,
            dark_masses,
            stellar_fractions,
        )

    def weighted_densities_vs_radii(self, radius_kpc, weight_cut=1e-6, number_bins=60):

        values = self.pdf.samples[1]
        self.setup_lensing_plane_ltm2r(values)
        radii = list(
            np.linspace(5e-3, radius_kpc * self.arcsec_per_kpc, number_bins + 1)
        )

        self.bulge_density_plot = []
        self.bulge_density_upper_plot = []
        self.bulge_density_lower_plot = []

        self.halo_density_plot = []
        self.halo_density_upper_plot = []
        self.halo_density_lower_plot = []

        self.dark_density_plot = []
        self.dark_density_upper_plot = []
        self.dark_density_lower_plot = []

        self.radii_plot = []

        for r in range(number_bins):

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

            weights = []
            bulge_densities = []
            halo_densities = []
            dark_densities = []

            for i in range(len(self.pdf.samples)):

                if self.pdf.weights[i] > weight_cut:

                    weights.append(self.pdf.weights[i])
                    values = self.pdf.samples[i]
                    self.setup_lensing_plane_ltm2r(values)

                    bulge_density = (
                        self.critical_density_arcsec
                        * (
                            self.lens_galaxy.bulge.mass_within_circle_in_units(
                                radii[r + 1]
                            )
                            - self.lens_galaxy.bulge.mass_within_circle_in_units(
                                radii[r]
                            )
                        )
                        / annuli_area
                    )

                    disk_density = (
                        self.critical_density_arcsec
                        * (
                            self.lens_galaxy.disk.mass_within_circle_in_units(
                                radii[r + 1]
                            )
                            - self.lens_galaxy.disk.mass_within_circle_in_units(
                                radii[r]
                            )
                        )
                        / annuli_area
                    )

                    halo_density = (
                        self.critical_density_arcsec
                        * (
                            self.lens_galaxy.halo.mass_within_circle_in_units(
                                radii[r + 1]
                            )
                            - self.lens_galaxy.halo.mass_within_circle_in_units(
                                radii[r]
                            )
                        )
                        / annuli_area
                    )

                    bulge_densities.append(bulge_density)
                    halo_densities.append(disk_density)
                    dark_densities.append(halo_density)

            bulge_avg, bulge_std = weighted_avg_and_std(
                np.asarray(bulge_densities), np.asarray(weights)
            )
            halo_avg, halo_std = weighted_avg_and_std(
                np.asarray(halo_densities), np.asarray(weights)
            )
            dark_avg, dark_std = weighted_avg_and_std(
                np.asarray(dark_densities), np.asarray(weights)
            )

            self.bulge_density_plot.append(bulge_avg)
            self.bulge_density_upper_plot.append(bulge_avg + bulge_std)
            self.bulge_density_lower_plot.append(bulge_avg - bulge_std)

            self.halo_density_plot.append(halo_avg)
            self.halo_density_upper_plot.append(halo_avg + halo_std)
            self.halo_density_lower_plot.append(halo_avg - halo_std)

            self.dark_density_plot.append(dark_avg)
            self.dark_density_upper_plot.append(dark_avg + dark_std)
            self.dark_density_lower_plot.append(dark_avg - dark_std)

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

    def plot_density(
        self, image_name="", labels=None, xaxis_is_physical=True, yaxis_is_physical=True
    ):

        plt.figure(figsize=(10, 10))
        plt.title("Decomposed convergence profile of " + image_name, size=20)

        if xaxis_is_physical:
            plt.xlabel("Distance From Galaxy Center (kpc)", size=20)
        else:
            plt.xlabel('Distance From Galaxy Center (")', size=20)

        if yaxis_is_physical:
            plt.ylabel(
                r"Surface Mass Density $\Sigma$ ($M_{\odot}$ / kpc $^2$)", size=20
            )
        else:
            pass

        plt.semilogy(
            self.radii_plot, self.bulge_density_plot, color="r", label="Sersic"
        )
        plt.semilogy(
            self.radii_plot, self.halo_density_plot, color="g", label="Exponential"
        )
        plt.semilogy(
            self.radii_plot, self.dark_density_plot, color="k", label="Dark Matter"
        )

        plt.semilogy(
            self.radii_plot, self.bulge_density_upper_plot, color="r", linestyle="--"
        )
        plt.semilogy(
            self.radii_plot, self.halo_density_upper_plot, color="g", linestyle="--"
        )
        plt.semilogy(
            self.radii_plot, self.dark_density_upper_plot, color="k", linestyle="--"
        )

        plt.semilogy(
            self.radii_plot, self.bulge_density_lower_plot, color="r", linestyle="--"
        )
        plt.semilogy(
            self.radii_plot, self.halo_density_lower_plot, color="g", linestyle="--"
        )
        plt.semilogy(
            self.radii_plot, self.dark_density_lower_plot, color="k", linestyle="--"
        )

        # plt.loglog(radii_plot, density_0, color='r', label='Sersic Bulge')
        # plt.loglog(radii_plot, density_1, color='g', label='EllipticalExponentialMass Halo')
        # plt.loglog(radii_plot, density_2, color='k', label='Dark Matter Halo')

        # plt.semilogy(radii_plot, density_lower[:, 0], color='r', linestyle='--')
        # plt.semilogy(radii_plot, density_upper[:, 0], color='r', linestyle='--')
        # plt.semilogy(radii_plot, density_lower[:, 1], color='g', linestyle='--')
        # plt.semilogy(radii_plot, density_upper[:, 1], color='g', linestyle='--')
        # plt.semilogy(radii_plot, density_lower[:, 2], color='k', linestyle='--')
        # plt.semilogy(radii_plot, density_upper[:, 2], color='k', linestyle='--')

        plt.axvline(x=self.einstein_radius * self.kpc_per_arcsec, linestyle="--")
        plt.axvline(x=self.source_light_min * self.kpc_per_arcsec, linestyle="-")
        plt.axvline(x=self.source_light_max * self.kpc_per_arcsec, linestyle="-")

        plt.tick_params(labelsize=16)
        plt.legend(fontsize=16)

        import os

        path = "{}/".format(os.path.dirname(os.path.realpath(__file__)))
        plt.savefig(path + "density/slacs1250_ltmr2.png", bbox_inches="tight")
        plt.show()
