""" 
plot distribution in chirp mass
used to make fugure 3 in paper
"""

import numpy as np

import matplotlib.pyplot as plt



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


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


def getPercentileLines(lineData,xs):
    print(np.shape(lineData))

    p01,p05,p25,p50,p75,p95,p99 = np.zeros(len(xs)), \
                          np.zeros(len(xs)), \
                          np.zeros(len(xs)), \
                          np.zeros(len(xs)), \
                          np.zeros(len(xs)), \
                          np.zeros(len(xs)), \
                          np.zeros(len(xs))
    
    for i in range(len(xs)):
        a = lineData[:,i]
        p01[i] = np.percentile(a,  0.5)        
        p05[i] = np.percentile(a,  5)
        p25[i] = np.percentile(a, 25)
        p50[i] = np.percentile(a, 50)
        p75[i] = np.percentile(a, 75)
        p95[i] = np.percentile(a, 95)
        p99[i] = np.percentile(a, 99.5)
    return p01, p05, p25, p50, p75, p95, p99
    
    

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


green2 = '#004D40'
blue2 = '#1E88E5'
orange2 = '#FFC107'
red2 = '#D81B60'

colSimp=orange
colGalExt=green
colGal = 'k'

nFreqs = 9
agnLinesDir = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/agn_model/f{0}/f{0}_plotting_things/'.format(nFreqs)
lines = np.genfromtxt('{}/mLines.dat'.format(agnLinesDir))

simpModMs = np.genfromtxt('{}/ms.dat'.format(agnLinesDir))


print(type(simpModMs),np.shape(simpModMs))
_, p05, p25, p50, p75, p95, _ = getPercentileLines(lines,simpModMs)


plt.fill_between(simpModMs,p05,p95,alpha=0.3,color=colSimp)
plt.fill_between(simpModMs,p25,p75,alpha=0.3,color=colSimp)
plt.plot(simpModMs,p50,color=colSimp,ls='--')







plt.ylim(1E-9,1E6)
plt.xlim(1E6,1E11)
plt.yscale('log')
plt.xscale('log')
plt.xlabel(r'$\mathcal{M}~~({\rm M_{\odot}})$')
plt.ylabel(r'${\rm d}n / {\rm d} \log_{10} \mathcal{M}/{\rm M_{\odot}}~~({\rm Mpc}^{-3})$')
plt.tight_layout()
#plt.savefig('dndlogM_3FreqBins.png')





galLinesDir = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/gal_model/f{0}/f{0}_plotting_things/'.format(nFreqs)
galExtModData = np.genfromtxt('{}/m_lines.dat'.format(galLinesDir))
galExtModLogMs = np.linspace(5,11,30)

print(type(galExtModData),np.shape(galExtModData))

galExtModMs = [ 10.0**logM for logM in galExtModLogMs ]
galExtModMs = np.atleast_1d(galExtModMs)

_, p05,p25,p50,p75,p95, _ = getPercentileLines(galExtModData,galExtModMs)

plt.fill_between(galExtModMs,p05,p95,alpha=0.3,color=colGalExt)
plt.fill_between(galExtModMs,p25,p75,alpha=0.3,color=colGalExt,label='new posterior')
plt.plot(galExtModMs,p50,color=colGalExt,ls='--')

#plt.savefig('forPaper/mlines_nFreq{}.pdf'.format(nFreqs))
#plt.savefig('forPaper/mlines_nFreq{}.png'.format(nFreqs))

#plt.show()
#exit()

galPriorLinesDir = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/gal_model/f{0}/f{0}_plotting_things_prior/'.format(nFreqs)
galExtModPrior = np.genfromtxt('{}/m_linesprior.dat'.format(galPriorLinesDir))
p01,p05,p25,p50,p75,p95,p99= getPercentileLines(galExtModPrior,galExtModMs)
plt.plot(galExtModMs, p01, color='k', alpha=0.7, ls=':',label='newly computed prior')
plt.plot(galExtModMs, p99, color='k', alpha=0.7, ls=':')



"""
galModRunLoc = '../../runs/galaxyModel/'
galModData = np.genfromtxt('{}/m_lines.dat'.format(galModRunLoc))
galModLogMs = np.genfromtxt('{}/mc.txt'.format(galModRunLoc))
galModMs = [ 10.0**logM for logM in galModLogMs ]

p05,p25,p50,p75,p95 = confBands.getPercentiles(galModMs,galModData)

plt.fill_between(galModMs,p05,p95,alpha=0.3,color=colGal)
plt.fill_between(galModMs,p25,p75,alpha=0.3,color=colGal)
plt.plot(galModMs,p50,color=colGal,ls='--')
"""



galPriorLinesDir = '/home/ADF/middlehr/repositories/PTAInference/runs/galaxyModel_ext/'
galExtModPrior = np.genfromtxt('{}/m_linesprior.dat'.format(galPriorLinesDir))
p01,p05,p25,p50,p75,p95,p99= getPercentileLines(galExtModPrior,galExtModMs)
plt.plot(galExtModMs, p01, color='k', alpha=0.7, ls='-.',label='M21 prior')
plt.plot(galExtModMs, p99, color='k', alpha=0.7, ls='-.')




galLinesDir = '/home/ADF/middlehr/repositories/PTAInference/runs/galaxyModel_ext/'
galExtModData = np.genfromtxt('{}/m_lines.dat'.format(galLinesDir))

_, p05,p25,p50,p75,p95, _ = getPercentileLines(galExtModData,galExtModMs)

plt.fill_between(galExtModMs,p05,p95,alpha=0.2,color='b',label='M21 posterior')
plt.fill_between(galExtModMs,p25,p75,alpha=0.2,color='b')
plt.plot(galExtModMs,p50,color='b',ls='--')

plt.legend()

plt.ylim(1E-9,1E6)
plt.xlim(1E6,1E11)
plt.yscale('log')
plt.xscale('log')
plt.xlabel(r'$\mathcal{M}~~({\rm M_{\odot}})$')
plt.ylabel(r'${\rm d}n / {\rm d} \log_{10} \mathcal{M}/{\rm M_{\odot}}~~({\rm Mpc}^{-3})$')
plt.tight_layout()
plt.savefig('forPaper/mlines_nFreq{}.pdf'.format(nFreqs))
plt.savefig('forPaper/mlines_nFreq{}.png'.format(nFreqs))

plt.show()
