import numpy as np
from lenstronomy.Analysis.td_cosmography import TDCosmography
from .diagonal_bnn_prior import DiagonalBNNPrior
from .base_cosmo_bnn_prior import BaseCosmoBNNPrior
import baobab.sim_utils.kinematics_utils as kinematics_utils

class DiagonalCosmoBNNPrior(DiagonalBNNPrior, BaseCosmoBNNPrior):
    """BNN prior with independent parameters

    Note
    ----
    This BNNPrior is cosmology-aware. For a version that's not tailored to H0 inference, see `DiagonalBNNPrior`.

    """
    def __init__(self, bnn_omega, components, kwargs_lens_eqn_solver={}):
        """
        Note
        ----
        The dictionary attributes are copies of the config corresponding to each component.
        The number of attributes depends on the number of components.

        Attributes
        ----------
        bnn_omega : dict
            copy of `cfg.bnn_omega`
        components : list
            list of components, e.g. `lens_mass`

        """
        DiagonalBNNPrior.__init__(self, bnn_omega, components)
        BaseCosmoBNNPrior.__init__(self, bnn_omega)
        self.params_to_exclude = []
        self.set_params_list(self.params_to_exclude)
        self.set_comps_qphi_to_e1e2()
        if self.kinematics.calculate_vel_disp or self.time_delays.calculate_time_delays:
            self.get_cosmography_observables = True
        else:
            self.get_cosmography_observables = False
        self.get_velocity_dispersion = getattr(kinematics_utils, 'velocity_dispersion_analytic') if self.kinematics.anisotropy_model == 'analytic' else getattr(kinematics_utils, 'velocity_dispersion_numerical')
        self.kwargs_lens_eqn_solver = kwargs_lens_eqn_solver

    def get_cosmo_observables(self, kwargs, z_lens, z_src, kappa_ext):
        """Calculate the central estimates of the observables for cosmography, i.e. the velocity dispersion and time delays, with and without noise realization

        Parameters
        ----------
        kwargs : dict
            the realized kwargs
        z_lens : float
        z_src : float
        kappa_ext : float

        Returns
        -------
        dict
            the computed central estimates of velocity dispersion and time delays with and without noise realization

        """
        td_cosmo = TDCosmography(z_lens, z_src, self.kwargs_model, cosmo_fiducial=self.cosmo, kwargs_lens_eqn_solver=self.kwargs_lens_eqn_solver)
        kwargs_lens_mass = dict(
                                theta_E=kwargs['lens_mass']['theta_E'],
                                gamma=kwargs['lens_mass']['gamma'],
                                center_x=kwargs['lens_mass']['center_x'],
                                center_y=kwargs['lens_mass']['center_y'],
                                e1=kwargs['lens_mass']['e1'],
                                e2=kwargs['lens_mass']['e2'],
                                )
        kwargs_ext_shear = dict(
                                gamma_ext=kwargs['external_shear']['gamma_ext'],
                                psi_ext=kwargs['external_shear']['psi_ext'],
                                ra_0=kwargs['external_shear']['ra_0'],
                                dec_0=kwargs['external_shear']['dec_0'],
                                )
        kwargs_lens = [kwargs_lens_mass, kwargs_ext_shear] # FIXME: hardcoded for SPEMD
        kwargs_lens_light = [kwargs['lens_light']]
        kwargs_ps = [dict(ra_source=kwargs['src_light']['center_x'],
                          dec_source=kwargs['src_light']['center_y'])]
        # Time delays
        if self.time_delays.calculate_time_delays:
            true_td, x_image, y_image = td_cosmo.time_delays(kwargs_lens, kwargs_ps, kappa_ext=kappa_ext)

        else:
            true_td = -1
        # Velocity dispersion
        if self.kinematics.calculate_vel_disp:
            true_vd = self.get_velocity_dispersion(
                                                   td_cosmo, 
                                                   kwargs_lens, 
                                                   kwargs_lens_light, 
                                                   self.kinematics.kwargs_anisotropy, 
                                                   self.kinematics.kwargs_aperture, 
                                                   self.kinematics.kwargs_psf, 
                                                   self.kinematics.anisotropy_model, 
                                                   kwargs['lens_light']['R_sersic'],
                                                   self.kinematics.kwargs_numerics,
                                                   kappa_ext,
                                                   )
        else:
            true_vd = -1
        obs = dict(true_td=true_td.tolist(),
                   true_vd=true_vd,
                   x_image=x_image,
                   y_image=y_image
                   )
        return obs

    def sample(self):
        kwargs = DiagonalBNNPrior.sample(self)
        H0 = self.cosmology.H0
        z_lens, z_src = self.sample_redshifts(self.redshift.copy())
        kappa_ext = self.sample_param(self.LOS.kappa_ext.copy())
        kwargs['misc'] = dict(
                             z_lens=z_lens,
                             z_src=z_src,
                             kappa_ext=kappa_ext,
                             H0=H0,
                             )
        if self.get_cosmography_observables:
            cosmo_obs = self.get_cosmo_observables(kwargs, z_lens, z_src, kappa_ext)
            kwargs['misc'].update(cosmo_obs)
        return kwargs