# 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

import mergerratemodel as mr
import prior_check 

from PTMCMCSampler.PTMCMCSampler import PTSampler as ptmcmc


def readChains(nFreqBins):

    ''' 
    read in the data and return the lowest nFreqBins
    chains and frequencies
    '''
    try: 
        samples = np.genfromtxt('/rds/projects/v/vecchioa-gw-pta/hannah/repositories/PTAInterpretation/chains/dr2new_fs/chains_h.txt')
        frequencies = np.genfromtxt('/rds/projects/v/vecchioa-gw-pta/hannah/repositories/PTAInterpretation/chains/dr2new_fs/freqs.txt')
    except: 
        samples = np.genfromtxt('/home/ADF/middlehr/repositories/PTAInterpretation/chains/dr2new_fs/chains_h.txt')        
        frequencies = np.genfromtxt('/home/ADF/middlehr/repositories/PTAInterpretation/chains//dr2new_fs/freqs.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(nFreqBins)}

    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.,5.,15)
    f = np.atleast_1d([ np.log10(fi) for fi in frequencies ])
    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(numberOfFrequencyBinsToUse):
        logL += KDEs[i].logpdf(model[i]) - 0.5*np.log10(2.*np.pi)
    #print(model,logL)
    return logL[0]



# prior bounds
bounds=[(-3.4,-2.4),
        (-0.6,0.2),
        (11,11.5),
        (-1.5,-1.),
        (-0.2,0.2),
        (0.01,0.05),
        (0.,2.),
        (-0.5,0.5),
        (-0.2,0.2),
        (0.1,10.),
        (-0.5,0.5),
        (-3.,1.),
        (-0.2,0.2),
        (7.75,8.75),
        (0.9,1.1),
        (0.2,0.5),
        (0.01,0.99),
        (-2.,2.)]






# read in the data
numberOfFrequencyBinsToUse = 9
frequencies, samples = readChains(numberOfFrequencyBinsToUse)

# make the kdes
KDEs = makeKDEs(samples, numberOfFrequencyBinsToUse)


#parTest = np.array([-2.78,-0.28,11.3,-1.24,-0.04,0.04,0.10,1.27,0.01,0.39,-0.2,-1.91,0.01,8.42,1.03,0.38,0.51,-0.7])
#log_likelihood(parTest)
#exit()

# ptmcmc things
x0 = np.array([-2.8,-0.2,11.25,-1.25,0.,0.025,0.8,0.,0.,1.,0.,-0.5,0.,8.25,1.,0.4,0.5,0.])
ndim = len(x0)
cov = np.diag(np.ones(ndim) * 0.01**2)
N = 1000000

sampler = ptmcmc(ndim, log_likelihood, log_prior, cov, outDir='./output', resume=False)
sampler.sample(x0, N, SCAMweight=30, AMweight=15, DEweight=50, isave=100)




'''
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()

    
'''
