import numpy as np

from hierarc.Likelihood.SneLikelihood.sne_likelihood_from_file import SneLikelihoodFromFile
from hierarc.Likelihood.SneLikelihood.sne_likelihood_custom import CustomSneLikelihood


class SneLikelihood(object):
    """
    Supernovae likelihood
    This class supports custom likelihoods as well as likelihoods from the Pantheon sample from file
    """
    def __init__(self, sample_name='CUSTOM', **kwargs_sne_likelihood):
        """

        :param sample_name: string, either 'CUSTOM' or a specific name supported by SneLikelihoodFromFile() class
        :param kwargs_sne_likelihood: keyword arguments to initiate likelihood class
        """
        if sample_name == 'CUSTOM':
            self._likelihood = CustomSneLikelihood(**kwargs_sne_likelihood)
        else:
            self._likelihood = SneLikelihoodFromFile(sample_name=sample_name, **kwargs_sne_likelihood)
        self.zhel = self._likelihood.zhel
        self.zcmb = self._likelihood.zcmb

    def log_likelihood(self, cosmo, apparent_m_z=None, sigma_m_z=None, z_anchor=0.1):
        """

        :param cosmo: instance of a class to compute angular diameter distances on arrays
        :param apparent_m_z: mean apparent magnitude of SN Ia at z=z_anchor (optional)
        :param z_anchor: redshift where definition of apparent_m_z is set (only applicable when apparent_m_z != None)
        :param sigma_m_z: 1-sigma scatter in magnitude in the intrinsic SNe brightness distribution not accounted-for
         by the covariance matrix
        :return: log likelihood of the data given the specified cosmology
        """
        angular_diameter_distances = cosmo.angular_diameter_distance(self.zcmb).value
        lum_dists = (5 * np.log10((1 + self.zhel) * (1 + self.zcmb) * angular_diameter_distances))

        ang_dist_anchor = cosmo.angular_diameter_distance(z_anchor).value
        lum_dist_anchor = (5 * np.log10((1 + z_anchor) * (1 + z_anchor) * ang_dist_anchor))

        return self._likelihood.log_likelihood_lum_dist(lum_dists - lum_dist_anchor, apparent_m_z, sigma_m_z)

import numpy as np

from hierarc.Likelihood.SneLikelihood.sne_likelihood_from_file import SneLikelihoodFromFile
from hierarc.Likelihood.SneLikelihood.sne_likelihood_custom import CustomSneLikelihood


class SneLikelihood(object):
    """
    Supernovae likelihood
    This class supports custom likelihoods as well as likelihoods from the Pantheon sample from file
    """
    def __init__(self, sample_name='CUSTOM', **kwargs_sne_likelihood):
        """

        :param sample_name: string, either 'CUSTOM' or a specific name supported by SneLikelihoodFromFile() class
        :param kwargs_sne_likelihood: keyword arguments to initiate likelihood class
        """
        if sample_name == 'CUSTOM':
            self._likelihood = CustomSneLikelihood(**kwargs_sne_likelihood)
        else:
            self._likelihood = SneLikelihoodFromFile(sample_name=sample_name, **kwargs_sne_likelihood)
        self.zhel = self._likelihood.zhel
        self.zcmb = self._likelihood.zcmb

    def log_likelihood(self, cosmo, apparent_m_z=None, sigma_m_z=None, z_anchor=0.1):
        """

        :param cosmo: instance of a class to compute angular diameter distances on arrays
        :param apparent_m_z: mean apparent magnitude of SN Ia at z=z_anchor (optional)
        :param z_anchor: redshift where definition of apparent_m_z is set (only applicable when apparent_m_z != None)
        :param sigma_m_z: 1-sigma scatter in magnitude in the intrinsic SNe brightness distribution not accounted-for
         by the covariance matrix
        :return: log likelihood of the data given the specified cosmology
        """
        angular_diameter_distances = cosmo.angular_diameter_distance(self.zcmb).value
        lum_dists = (5 * np.log10((1 + self.zhel) * (1 + self.zcmb) * angular_diameter_distances))

        ang_dist_anchor = cosmo.angular_diameter_distance(z_anchor).value
        lum_dist_anchor = (5 * np.log10((1 + z_anchor) * (1 + z_anchor) * ang_dist_anchor))

        return self._likelihood.log_likelihood_lum_dist(lum_dists - lum_dist_anchor, apparent_m_z, sigma_m_z)

Supernovae likelihood     This class supports custom likelihoods as well as likelihoods from the Pantheon sample from file :param sample_name: string, either 'CUSTOM' or a specific name supported by SneLikelihoodFromFile() class         :param kwargs_sne_likelihood: keyword arguments to initiate likelihood class :param cosmo: instance of a class to compute angular diameter distances on arrays         :param apparent_m_z: mean apparent magnitude of SN Ia at z=z_anchor (optional)         :param z_anchor: redshift where definition of apparent_m_z is set (only applicable when apparent_m_z != None)         :param sigma_m_z: 1-sigma scatter in magnitude in the intrinsic SNe brightness distribution not accounted-for          by the covariance matrix         :return: log likelihood of the data given the specified cosmology