import numpy as np
import matplotlib.pyplot as plt
import corner

import matplotlib
matplotlib.rcParams.update({'font.size': 16})


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}

    return posteriorCombined



def fiveDCorner(data,nFBins):

    dataMstar = [ np.log10(10.0**(delta+7.0)) for delta in data['delta'] ]


    d2 =  np.atleast_1d([data['logn'],
                         data['beta'],
                         data['gamma'],
                         data['alpha'],
                         dataMstar]).T
    parameterLabels = ([r'$\log_{10}\frac{{\dot n}_0}{{\rm Mpc}^{3}{\rm Gyr}}$',\
                        r'$\beta_z$', r'$z_0$', r'$\alpha_{\mathcal{M}}$',\
                        r'$\log_{10} \frac{\mathcal{M}_*}{ M_{\odot}}$'])

    figure = corner.corner(d2,labels=parameterLabels,
                           quantiles=[0.05, 0.5, 0.95])
    plt.savefig('forPaper/agn_fullCorner_fBins{}.pdf'.format(nFBins))
    plt.savefig('forPaper/agn_fullCorner_fBins{}.png'.format(nFBins))
    return 


n=9
posterior = getAGNPost(n)#np.genfromtxt('posterior.dat', names=True)
fiveDCorner(posterior,n)


exit()
names = ['logn', 'beta', 'gamma', 'alpha', 'delta']
for i in range(5):

    oneDData = posterior[names[i]]
    
    # percentile:
    
    
    plt.hist(oneDData)
    plt.show()
