# trying to convert to cpnest

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

import mergerratemodel as mr
import prior_check 

import cpnest.model


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,names):


    #return all(bounds[i][0] < par[i] < bounds[i][1] for i in range(len(par)))
    return all(bounds[i][0] < par[i] < bounds[i][1] for i in names)
'''

class galModelKDEs(cpnest.model.Model):

    def __init__(self):
        self.names = ['Phi0', 'PhiI', 'M0', 'alpha', 'alphaI', \
                 'f0', 'beta', 'gamma', 'delta', 't0', \
                 'epsilon', 'zeta', 'eta', 'Ms', 'theta',\
                 'sigma','e0','rho']

        self.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.]]
        data=None

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

    def log_likelihood(self,par):
        M1 = np.linspace(9,12,25)
        q = np.linspace(0.25,1,10)
        z = np.linspace(0.,1.5,5)
        f = np.atleast_1d([ np.log10(fi) for fi in frequencies ])
        initpar = dict(Phi0    = par[self.names[0]], 
                       PhiI    = par[self.names[1]], 
                       M0      = par[self.names[2]], 
                       alpha   = par[self.names[3]], 
                       alphaI  = par[self.names[4]], 
                       f0      = par[self.names[5]], 
                       beta    = par[self.names[6]], 
                       gamma   = par[self.names[7]], 
                       delta   = par[self.names[8]], 
                       t0      = par[self.names[9]], 
                       epsilon = par[self.names[10]], 
                       zeta    = par[self.names[11]], 
                       eta     = par[self.names[12]], 
                       Ms      = par[self.names[13]], 
                       theta   = par[self.names[14]], 
                       sigma   = par[self.names[15]], 
                       e0      = par[self.names[16]], 
                       rho     = par[self.names[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])

        return logL[0]






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

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

# 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.]]


mymodel = galModelKDEs()

nest = cpnest.CPNest(mymodel,maxmcmc=1000,nlive=128,verbose=3)
nest.run()
cpnest.CPNest.get_posterior_samples(nest) 

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


