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

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


def getAGNPost(nFreqBins):

    # read in the posterior files for the chosen number of frequency bins 
    path0 = '../results/agnosticPosterior/nFreq{}/posterior_0.dat'.format(nFreqBins)
    path1 = '../results/agnosticPosterior/nFreq{}/posterior_1.dat'.format(nFreqBins)
    path2 = '../results/agnosticPosterior/nFreq{}/posterior_2.dat'.format(nFreqBins)
    post0 = np.genfromtxt(path0,names=True)
    post1 = np.genfromtxt(path1,names=True)
    post2 = np.genfromtxt(path2,names=True)

    # combine the posteriors
    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']))

    # make a dictionary with the full results 
    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}}$'])

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


def oneDimensionalHistograms(data,nFBins):

    # individual one dimensional plots
    names = ['logn', 'beta', 'gamma', 'alpha', 'delta']
    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}}$'])

    for i in range(len(names)):

        oneDData = data[names[i]]

        plt.clf()
        plt.hist(oneDData,histtype='step')
        plt.xlabel(parameterLabels[i])
        plt.ylabel('count')
        plt.tight_layout()
        plt.show()


def main():

    # choose the number of frequency bins to use
    n=9

    # full corner plot
    posterior = getAGNPost(n)
    fiveDCorner(posterior,n)

    # one dimensional histograms
    oneDimensionalHistograms(posterior,n)


    
if __name__=='__main__':
    main()





