import numpy as np
from cosmosis.datablock import names, option_section

cosmo = names.cosmological_parameters

def setup(options):
    direc='distance_result/'
    data             = np.loadtxt(direc+'mean.dat', skiprows=1)
    z_data           = data[:,0]
    lnDA_data        = data[:,1]

    cov_data         = np.loadtxt(direc+'cov_matrix.dat', skiprows=1)
    
    return z_data, lnDA_data, cov_data

def execute(block, config):

    z_data, lnDA_data, cov_data = config     
    #inversion of the covariance matrix
    precision = np.linalg.inv(cov_data)

    z        = block[names.distances, "z"]
    DA       = block[names.distances, "D_A"] 
    S_D0     = block[cosmo, "S_D0"]
    
   # Note that I add S_D0 to change the *theory*, not to
   # change the *data*. Mathematically it is the same, but
   # conceptually it is different: cov_data then represents *data*
   # exactly.

    DA_model            = np.interp(z_data, z, DA)
    lnDA_model          = np.log(DA_model) - np.log(S_D0)
 
    d    = lnDA_model - lnDA_data
    chi2 = np.dot(d, np.dot(precision,d))
    
    like = -chi2/2.0

    #ln(likehood)
    block[names.likelihoods, "massE_like"]= like

    #signal success
    return 0

def cleanup(config):
    pass
