import numpy as np
import cpnest.model
from scipy.stats import gaussian_kde as kde
import astropy.units as u
import hmodel


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 converthc(Ayr,freq):
    '''
    convert from hc at 1/1yr to whatever frequency we want
    '''
    # 1/1yr in Hz
    freq1yr = 1./(1.*u.year.to(u.s))
    
    # convert 
    freqRatio = (freq / freq1yr)
    hc = Ayr*(freqRatio**(-2./3.))

    log10hc = np.log10(float(hc))
    return log10hc

    

def priorBounds():
	
    logno = [-20.0,3.0]
    beta  = [-2.0,7.0]
    gamma = [0.2,5.0]
    alpha = [-3.0,3.0]
    delta = [-1.0,2.0]

    priors = [logno,beta,gamma,alpha,delta]

    return priors


class kdeModel(cpnest.model.Model):




    # parameter names and get flat prior bounds
    names=['logn','beta','gamma','alpha','delta']
    bounds=priorBounds()
    
    
    def log_likelihood(self,param):
    
        # integration limits
        MlowIntLimit  = 10.0**6.0	
        MhighIntLimit = 10.0**11.0
        zlowIntLimit  = 0.0 
        zhighIntLimit = 5.0

        # calculate h
        theta = param['logn'],param['beta'],param['gamma'],param['alpha'],param['delta']
        A = hmodel.hmodel(theta,\
                          MlowIntLimit,MhighIntLimit,\
                          zlowIntLimit,zhighIntLimit)
                                         
        #log10hModel = np.log10(float(hc))

        # convert hc
        
        log10hcs = [ converthc(A,f) for f in frequencies ] 

        logL = 0.
        for i in range(numberOfFrequencyBinsToUse):
            logL += KDEs[i].logpdf(log10hcs[i])

        return logL[0]

        
        

numberOfFrequencyBinsToUse = 5
frequencies, samples = readChains(numberOfFrequencyBinsToUse)

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

        
mymodel = kdeModel()


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

