import numpy as np
from scipy.integrate import quad
import astropy.units as u
from astropy.cosmology import Planck18 as cosmo


def getAGNPost(nFreqBins):

    if nFreqBins==9: 
        path0 = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{}/nFreq9/posterior.dat'.format(nFreqBins)
        path1 = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{}/nFreq9_1/posterior.dat'.format(nFreqBins)
        path2 = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{}/nFreq9_2/posterior.dat'.format(nFreqBins)
        post0 = np.genfromtxt(path0,names=True)
        post1 = np.genfromtxt(path1,names=True)
        post2 = np.genfromtxt(path2,names=True)
    elif nFreqBins==5:
        path0 = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{}/nFreq5/posterior.dat'.format(nFreqBins)
        path1 = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{}/nFreq5_1/posterior.dat'.format(nFreqBins)
        path2 = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{}/nFreq5_2/posterior.dat'.format(nFreqBins)
        post0 = np.genfromtxt(path0,names=True)
        post1 = np.genfromtxt(path1,names=True)
        post2 = np.genfromtxt(path2,names=True)
    elif nFreqBins==3:
        path0 = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{}/nFreq3/posterior.dat'.format(nFreqBins)
        path1 = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{}/nFreq3_1/posterior.dat'.format(nFreqBins)
        path2 = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{}/nFreq3_2/posterior.dat'.format(nFreqBins)
        post0 = np.genfromtxt(path0,names=True)
        post1 = np.genfromtxt(path1,names=True)
        post2 = np.genfromtxt(path2,names=True)

    #combine the ns 
    post0And1 = np.concatenate((post0['logn'], post1['logn']))
    logns = np.concatenate((post0And1, post2['logn']))

    post0And1 = np.concatenate((post0['alpha'], post1['alpha']))
    alphas = np.concatenate((post0And1, post2['alpha']))

    post0And1 = np.concatenate((post0['delta'], post1['delta']))
    deltas = np.concatenate((post0And1, post2['delta']))

    post0And1 = np.concatenate((post0['beta'], post1['beta']))
    betas = np.concatenate((post0And1, post2['beta']))

    post0And1 = np.concatenate((post0['gamma'], post1['gamma']))
    gammas = np.concatenate((post0And1, post2['gamma']))

    posteriorCombined = {'logn': logns, 
                         'alpha': alphas, 
                         'delta': deltas,
                         'beta': betas,
                         'gamma': gammas}

    saveLinesHere = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{0}/f{0}_plotting_things/'.format(nFreqBins)
    return posteriorCombined, saveLinesHere


def zInt(beta,gamma,zlow,zhigh):

    OmegaM = cosmo.Om0
    OmegaL = cosmo.Ode0

    # Ho from Planck18, which we convert to units of 1/Gyrs
    Gyr_in_s = 1.e9 * 1.*u.year.to(u.s)
    Mpc_in_m = 1.e6 * 1.*u.parsec.to(u.m)
    H0_in_km_per_s_per_Mpc = cosmo.H(0).value
    H0_in_per_Gyr =  H0_in_km_per_s_per_Mpc * Gyr_in_s / (Mpc_in_m / 1000)
    Ho = H0_in_per_Gyr

    zfn = lambda zi: ( (1.0 + zi)**(beta-1.0) * np.exp(zi/gamma) ) / \
                     ( Ho * ( OmegaM*(1.0+zi)**3. + OmegaL )**(1./2.) )
    zInt = quad(zfn,zlow,zhigh, epsabs=1.49e-12, epsrel=1.49e-12)
    
    return zInt[0]





def dndlog10M(theta,log10Ms):

    logno, beta, gamma, alpha, delta = theta

    no = 10.0**(logno)

    zIntegral = zInt(beta,gamma,0.0,5.0)

    Ms = [ 10.0**l10Ms for l10Ms in log10Ms ] 

    dndlogMs = [ no * zIntegral \
                 * (Mi / (10.**7.))**(-alpha) \
                 * np.exp(-Mi / 10.**(delta+7.0)) \
                 for Mi in Ms ]

    return dndlogMs
    






def main():

    #pathToRuns = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/nFreq3/posterior.dat'
    #pathToRuns = '../models/agnostic_model/posterior.dat'
    
    nFreqBins=9
    simpleModelData,resultsDir = getAGNPost(nFreqBins)


    # set up M range and write to file
    log10MRange = np.linspace(6.0, 11.0, 40)
    writeMs = open('{}/ms.dat'.format(resultsDir),'w')
    for logM in log10MRange:
        writeMs.write('{}\n'.format(10.0**logM))
    writeMs.close()

    # place to write results 
    writeLines = open('{}/mLines.dat'.format(resultsDir),'w')    

    # compute lines
    for i in range(len(simpleModelData['logn'])):
        thetas = simpleModelData['logn'][i],  \
                 simpleModelData['beta'][i],  \
                 simpleModelData['gamma'][i], \
                 simpleModelData['alpha'][i], \
                 simpleModelData['delta'][i]
        ys = dndlog10M(thetas,log10MRange)   
        
        for yi in ys:        
            writeLines.write('{}\t'.format(yi))
        writeLines.write('\n')
    
    writeLines.close()
    

 
if __name__ == "__main__":
    main()

