
from matplotlib import pylab
import numpy as np
from math import ceil
import matplotlib.pyplot as plt

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


plt.rcParams.update({
    "text.usetex": True,
    "font.family": "Helvetica"
})



def fpairtest(delta):
    return 1./(delta+1.)*(1.-0.25**(delta+1.))

def getneffGAL(nFreqBins):

    
    if nFreqBins==9:
        path = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/gal_model/f9/ptmcmc_f{}_1/output/'.format(nFreqBins)
    elif nFreqBins==5: 
        path = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/gal_model/f5/ptmcmc_2/output/'.format(nFreqBins)
    elif nFreqBins==3:
        path = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/gal_model/f3/ptmcmc_f{}_1/output/'.format(nFreqBins)

    chains = np.genfromtxt('{}/chain_1.txt'.format(path))

    burnin = int(0.25*len(chains))


    n1 = chains[burnin:,0]-chains[burnin:,2]+np.log10(chains[burnin:,5]/chains[burnin:,9]/fpairtest(chains[burnin:,8]))+chains[burnin:,10]*np.log10(0.4/0.7)+(chains[burnin:,10]-chains[burnin:,6])*(11.-chains[burnin:,2])+11.


    priorChains = np.genfromtxt('/home/ADF/middlehr/repositories/PTAInference/runs/galaxyModel_ext/prior.dat')

    priorn1 = priorChains[:,0]-priorChains[:,2]+np.log10(priorChains[:,5]/priorChains[:,9]/fpairtest(priorChains[:,8]))+priorChains[:,10]*np.log10(0.4/0.7)+(priorChains[:,10]-priorChains[:,6])*(11.-priorChains[:,2])+11.
    
    return n1, priorn1
    
    
def getnAGN(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']))

    return logns


def histOutline(dataIn, *args, **kwargs):
    """
    code copied from http://www.scipy.org/Cookbook/Matplotlib/UnfilledHistograms?action=AttachFile&do=get&target=histNofill.py

    Make a histogram that can be plotted with plot() so that
    the histogram just has the outline rather than bars as it
    usually does.

    Example Usage:
    binsIn = numpy.arange(0, 1, 0.1)
    angle = pylab.rand(50)

    (bins, data) = histOutline(binsIn, angle)
    plot(bins, data, 'k-', linewidth=2)

    """

    (histIn, binsIn) = np.histogram(dataIn, *args, **kwargs)

    stepSize = binsIn[1] - binsIn[0]

    bins = np.zeros(len(binsIn)*2 + 2, dtype=float)
    data = np.zeros(len(binsIn)*2 + 2, dtype=float)    
    for bb in range(len(binsIn)):
        bins[2*bb + 1] = binsIn[bb]
        bins[2*bb + 2] = binsIn[bb] + stepSize
        if bb < len(histIn):
            data[2*bb + 1] = histIn[bb]
            data[2*bb + 2] = histIn[bb]

    bins[0] = bins[1]
    bins[-1] = bins[-2]
    data[0] = 0
    data[-1] = 0
    
    return (bins, data)

def plotOneDPosterior(data,paramLabel,setBins,colour='k-',label='',prior=False):

    #pathToRuns = "../../runs/simpleModelPosteriors/"
    #setBins = np.arange(int(min(data[paramName])), ceil(max(data[paramName])),0.5)
    #setBins = np.linspace(min(data[paramName]), max(data[paramName]),50)
    #setBins = np.linspace(-14.5,2.5,50)

    p05 = np.percentile(data,5)
    p50 = np.percentile(data,50)
    p95 = np.percentile(data,95)
    

    if prior==True: 
        plt.hist(data, bins=setBins, histtype='step', facecolor=colour, alpha=0.3, edgecolor=None,density=1,fill=True,label=label)

    else: 
        (bins,dat) = histOutline(data,setBins,density=1)
        pylab.plot(bins,dat,colour,linewidth=2,alpha=1.0,label=label)
        plt.axvline(p05,ls=':',color=colour)
        plt.axvline(p95,ls=':',color=colour)
    plt.xlabel(paramLabel)
    plt.ylabel(r'${\rm Normalised~counts}$')
    plt.tight_layout()
    #plt.savefig('{}lognoHist.pdf'.format(runLoc))
    #plt.savefig('logno.png')
    #pylab.show()


    #print("""{}
    #05%: {:.2f} ({})
    #50%: {:.2f} ({})
    #95%: {:.2f} ({})
    #""".format(label,p05,10**p05,p50,10.**p50,p95,10.**p95))
    #
    print("""{}
    {:.2f}^{{+{:.2f}}}_{{-{:.2f}}}""".format(label, p50, p95-p50, p50-p05))

    return 0



n = 9

print('Number of frequency bins = {}'.format(n))


neffGAL, neffPriorGAL = getneffGAL(n)
nAGN = getnAGN(n)

# bins
low  = int(min(nAGN))
high = ceil(max(nAGN))
bins = np.linspace(low,high,80)

orange = '#E69F00'
blue = '#0072B2'
green = '#009E73'
red = '#D55E00'
purple = '#CC79A7'

colSimp=orange
colGalExt=green
#colGal = 'k'
#colSimpNegAlpha = red

lognoLabel = r'$\log_{10}\frac{{\dot n}_0}{{\rm Mpc}^{3}{\rm Gyr}}$'

plotOneDPosterior(nAGN,lognoLabel,bins,colour=colSimp,label=r'$\textrm{Agnostic posterior}$')
plotOneDPosterior(neffGAL,lognoLabel,bins,colour=colGalExt,label=r'$\textrm{Astro-informed~posterior}$')
plotOneDPosterior(neffPriorGAL,lognoLabel,bins,colour=colGalExt,label=r'$\textrm{Astro-informed prior}$',prior=True)
plt.legend(fontsize=12)
plt.xlim(min(bins),max(bins))
plt.savefig('forPaper/nhistogram_nFreq{}.png'.format(n))
plt.savefig('forPaper/nhistogram_nFreq{}.pdf'.format(n))
plt.show()



