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


import sys

sys.path.insert(1, '../models/agnostic_model/')
import hmodel

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 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
    
    
    


posteriors, directory = getAGNPost(9)
nSamples = np.shape(posteriors['delta'])[0]


writeLines = open('./agn_h_nF{}.dat'.format(9),'w')

for i in range(nSamples):
    theta = posteriors['logn'][i],\
            posteriors['alpha'][i],\
            posteriors['delta'][i],\
            posteriors['beta'][i],\
            posteriors['gamma'][i]
    # integration limits
    MlowIntLimit  = 10.0**6.0	
    MhighIntLimit = 10.0**11.0
    zlowIntLimit  = 0.0 
    zhighIntLimit = 5.0 
    h1yr = hmodel.hmodel(theta,MlowIntLimit,MhighIntLimit,zlowIntLimit,zhighIntLimit)
    
    logfs = np.linspace(-9.,-6.,15)
    
    hLine = [ converthc(h1yr,10.**logf) for logf in logfs]

    for yi in hLine:        
        writeLines.write('{}\t'.format(yi))
    writeLines.write('\n')
    
