# This is based on the ptmcmcrun.py here:
# https://github.com/hannahm8/PTAInference/tree/master/code/galaxy_model 
# 

import numpy as np
from scipy.stats import gaussian_kde as kde
import matplotlib.pyplot as plt


def readChains(nFreqBins):

    ''' 
    read in the data and return the lowest nFreqBins
    chains and frequencies
    '''
    frequencies = np.genfromtxt('../../chains/dr2new_fs/freqs.txt')
    samples = np.genfromtxt('../../chains/dr2new_fs/chains_h.txt')

    chainsToUse = samples[:,0:nFreqBins]
    frequenciesToUse = frequencies[0:nFreqBins]
    
    return frequenciesToUse, chainsToUse


def makeKDEs(samples,nFreqBins):
    '''
    make the KDEs to use in the likelihood
    '''
   
    KDE = {i: kde(samples[:,i]) for i in range(nFrequencies)}

    return KDE


def in_bounds(par):
    return all(bounds[i][0] < par[i] < bounds[i][1] for i in range(len(par)))


def log_prior(par):
    if in_bounds(par):
        param = dict(Phi0 = par[0], PhiI = par[1], M0 = par[2], alpha = par[3], alphaI = par[4], f0 = par[5], beta = par[6], gamma = par[7], delta = par[8], t0 = par[9], epsilon = par[10], zeta = par[11], eta = par[12], Ms = par[13], theta = par[14], sigma = par[15], e0 = par[16], rho = par[17])        
        return prior_check.check_p(param)
    else:
        return -np.inf


def log_likelihood(par):
    
    M1 = np.linspace(9,12,25)
    q = np.linspace(0.25,1,10)
    z = np.linspace(0.,1.5,5)

    initpar = dict(Phi0 = par[0], 
                   PhiI = par[1], 
                   M0 = par[2], 
                   alpha = par[3], 
                   alphaI = par[4], 
                   f0 = par[5], 
                   beta = par[6], 
                   gamma = par[7], 
                   delta = par[8], 
                   t0 = par[9], 
                   epsilon = par[10], 
                   zeta = par[11], 
                   eta = par[12], 
                   Ms = par[13], 
                   theta = par[14], 
                   sigma = par[15], 
                   e0 = par[16], 
                   rho = par[17])
    model = mr.mergerrate(M1,q,z,f,**initpar).hmodelt(fbin=None)[0]

    logL = 0.
    for i in range(nFrequencies):
        logL += KDE[i].logpdf(model)
    return logL




numberOfFrequencyBinsToUse = 5
frequencies, samples = readChains(numberOfFrequencyBinsToUse)
KDEs = makeKDEs(samples, numberOfFrequencyBinsToUse)



to do - add sampler



'''
some testing...
KDE = {i: kde(samples[:,i]) for i in range(nFrequencies)}
print(KDE)


for i in range(nFrequencies):

    bins = np.linspace(min(samples[:,i]), max(samples[:,i]), 100)
    print('bins' , bins)
    toPlot = np.linspace(min(samples[:,i]), max(samples[:,i]), 100)
    print(KDE[i].logpdf(-14))
    
    plt.plot(toPlot,KDE[i].logpdf(toPlot))
    plt.hist(samples[:,i],bins=bins,density=True,histtype='step')
    plt.show()

    
'''
