import corner
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 burninAndThin(samples,burnin,thin):
    burnItIn = samples[burnin:]
    thinIt = burnItIn[0::thin]
    return thinIt


def allCorner(burnin=500,thin=2):

    burninP, thinP = 0, 1
    nFreqBins = 9
    path = '/home/ADF/middlehr/Documents/PTA/EPTAInterp/results/gal_model/f9/ptmcmc_f{}_1/output/'.format(nFreqBins)
    galExt = np.genfromtxt('{}/chain_1.txt'.format(path))
    print(galExt)
    galParNames = np.loadtxt('/home/ADF/middlehr/repositories/PTAInference/runs/galaxyModel_ext/pars.txt',dtype=str)
    
    prior = np.genfromtxt('/home/ADF/middlehr/repositories/PTAInference/runs/galaxyModel_ext/prior.dat')

    burnin = int(0.25*len(galExt))

    Phi0Index = list(galParNames).index('Phi0')
    Phi0      = burninAndThin(galExt[:,Phi0Index],burnin,thin)   
    Phi0Prior = burninAndThin(prior[:,Phi0Index],burninP,thinP)
    
    PhiIIndex = list(galParNames).index('PhiI')
    PhiI      = burninAndThin(galExt[:,PhiIIndex],burnin,thin)
    PhiIPrior = burninAndThin(prior[:,PhiIIndex],burninP,thinP)   
    
    M0Index = list(galParNames).index('M0')
    M0      = burninAndThin(galExt[:,M0Index],burnin,thin)   
    M0Prior = burninAndThin(prior[:,M0Index],burninP,thinP)   

    alpha0Index = list(galParNames).index('alpha0')
    alpha0      = burninAndThin(galExt[:,alpha0Index],burnin,thin)
    alpha0Prior = burninAndThin(prior[:,alpha0Index],burninP,thinP)

    alphaIIndex = list(galParNames).index('alphaI')
    alphaI      = burninAndThin(galExt[:,alphaIIndex],burnin,thin)
    alphaIPrior = burninAndThin(prior[:,alphaIIndex],burninP,thinP)
    
    f0Index = list(galParNames).index('f0')
    f0      = burninAndThin(galExt[:,f0Index],burnin,thin)
    f0Prior = burninAndThin(prior[:,f0Index],burninP,thinP)

    betafIndex = list(galParNames).index('betaf')
    betaf      = burninAndThin(galExt[:,betafIndex],burnin,thin)
    betafPrior = burninAndThin(prior[:,betafIndex],burninP,thinP)

    alphafIndex = list(galParNames).index('alphaf')
    alphaf      = burninAndThin(galExt[:,alphafIndex],burnin,thin)
    alphafPrior = burninAndThin(prior[:,alphafIndex],burninP,thinP)
    
    gammafIndex = list(galParNames).index('gammaf')
    gammaf      = burninAndThin(galExt[:,gammafIndex],burnin,thin)
    gammafPrior = burninAndThin(prior[:,gammafIndex],burninP,thinP)

    tau0Index = list(galParNames).index('tau0')
    tau0      = burninAndThin(galExt[:,tau0Index],burnin,thin)
    tau0Prior = burninAndThin(prior[:,tau0Index],burninP,thinP)
    
    alphatauIndex = list(galParNames).index('alphatau')
    alphatau      = burninAndThin(galExt[:,alphatauIndex],burnin,thin)
    alphatauPrior = burninAndThin(prior[:,alphatauIndex],burninP,thinP)
    
    betatauIndex = list(galParNames).index('betatau')
    betatau      = burninAndThin(galExt[:,betatauIndex],burnin,thin)
    betatauPrior = burninAndThin(prior[:,betatauIndex],burninP,thinP)
    
    gammatauIndex = list(galParNames).index('gammatau')
    gammatau      = burninAndThin(galExt[:,gammatauIndex],burnin,thin)
    gammatauPrior = burninAndThin(prior[:,gammatauIndex],burninP,thinP)

    MstarIndex = list(galParNames).index('Mstar')
    Mstar      = burninAndThin(galExt[:,MstarIndex],burnin,thin)
    MstarPrior = burninAndThin(prior[:,MstarIndex],burninP,thinP)

    alphastarIndex = list(galParNames).index('alphastar')
    alphastar      = burninAndThin(galExt[:,alphastarIndex],burnin,thin)
    alphastarPrior = burninAndThin(prior[:,alphastarIndex],burninP,thinP)

    epsilonIndex = list(galParNames).index('epsilon')
    epsilon      = burninAndThin(galExt[:,epsilonIndex],burnin,thin)
    epsilonPrior = burninAndThin(prior[:,epsilonIndex],burninP,thinP)
    
    e0Index = list(galParNames).index('e0')
    e0      = burninAndThin(galExt[:,e0Index],burnin,thin)
    e0Prior = burninAndThin(prior[:,e0Index],burninP,thinP)
    
    zeta0Index = list(galParNames).index('zeta0')
    zeta0      = burninAndThin(galExt[:,zeta0Index],burnin,thin)
    zeta0Prior = burninAndThin(prior[:,zeta0Index],burninP,thinP)


    # do parameter swap
    # swap betaf and alphaf
    data = np.atleast_1d([Phi0,
                          PhiI,
                          M0,
                          alpha0,
                          alphaI,
                          f0,
                          alphaf,
                          betaf,
                          gammaf,
                          tau0,
                          alphatau,
                          betatau,
                          gammatau,
                          Mstar,
                          alphastar,
                          epsilon,
                          e0,
                          zeta0]).T

    priorSamples = np.atleast_1d([Phi0Prior,
                                  PhiIPrior,
                                  M0Prior,
                                  alpha0Prior,
                                  alphaIPrior,
                                  f0Prior,
                                  alphafPrior,
                                  betafPrior,
                                  gammafPrior,
                                  tau0Prior,
                                  alphatauPrior,
                                  betatauPrior,
                                  gammatauPrior,
                                  MstarPrior,
                                  alphastarPrior,
                                  epsilonPrior,
                                  e0Prior,
                                  zeta0Prior]).T

    parameterLabels = ([r'$\Phi_0$',
                        r'$\Phi_I$',
                        r'$\log_{10} M_0$', 
                        r'$\alpha_0$',
                        r'$\alpha_I$',
                        r'$f_0$',
                        r'$\alpha_f$',
                        r'$\beta_f$',
                        r'$\gamma_f$',
                        r'$\tau_0$',
                        r'$\alpha_{\tau}$',
                        r'$\beta_{\tau}$',
                        r'$\gamma_{\tau}$',
                        r'$\log_{10} M_{*}$',
                        r'$\alpha_{*}$',
                        r'$\epsilon$',
                        r'$e_0$',
                        r'$\log_{10}\zeta_0$'])


    bounds=[(-3.4,-2.4), #Phi0
            (-0.6,0.2),  #PhiI
            (11,11.5),   #M0
            (-1.5,-1.),  # alpha0
            (-0.2,0.2),  # alphaI
            (0.01,0.05), #f0
            (-0.5,0.5),  #alphaf
            (0.,2.),     #betaf
            (-0.2,0.2),  #gammaf
            (0.1,5.),   #tau0 full range is [0,10] 
            (-0.5,0.5),  #alphatau
            (-3.,1.),    #betatau
            (-0.2,0.2),  #gammatau
            (7.75,8.75), # M*
            (0.9,1.1),   #alpha*
            (0.2,0.5),   #epsilon
            (0.01,0.99), #e0
            (-2.,2.)]    #zeta

    figure = corner.corner(data,
                           levels=(0.5,0.9),
                           labels=parameterLabels,
                           quantiles=(0.05,0.50,0.95),
                           show_titles=True,
                           title_fmt='.2f',
                           plot_density=True,
                           range=bounds,normed=True)
    corner.corner(priorSamples,
                  fig=figure,
                  no_fill_contours=True,
                  plot_datapoints=False,
                  plot_density=False,
                  levels=([1.]),
                  color='g',
                  linewidths=100,
                  range=bounds,
                  normed=True)
    corner.corner(data, 
                  fig=figure,
                           levels=(0.5,0.9),
                           labels=parameterLabels,
                           quantiles=(0.05,0.50,0.95),
                           show_titles=True,
                           title_fmt='.2f',
                           plot_density=True,
                           range=bounds,normed=True)
                                        
    '''
    prior = np.genfromtxt('/home/ADF/middlehr/repositories/PTAInference/runs/galaxyModel_ext/prior.dat')

    thin=2
    mStarBurninPrior = prior[burnin:,mStarIndex]
    mStarPrior       = mStarBurninPrior[0::thin]

    tauBurninPrior = prior[burnin:,tauIndex]
    tauPrior       = tauBurninPrior[0::thin] 

    alphaTauBurninPrior = prior[burnin:,alphaTauIndex]
    alphaTauPrior       = alphaTauBurninPrior[0::thin]
    
    betaTauBurninPrior = prior[burnin:,betaTauIndex]
    betaTauPrior       = betaTauBurninPrior[0::thin]

    priorSamples = np.atleast_1d([tauPrior,alphaTauPrior,betaTauPrior,mStarPrior]).T
    print(np.shape(priorSamples))

    '''
    plt.savefig('forPaper/gal_fullCorner_fBins{}.pdf'.format(nFreqBins))
    plt.savefig('forPaper/gal_fullCorner_fBins{}.png'.format(nFreqBins))
    plt.show()

    return 


def selectedCorner(burnin=500,thin=2):

    burninP, thinP = 0, 1    
    
    nFreqBins = 9
    
    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)

    galExt = np.genfromtxt('{}/chain_1.txt'.format(path))
    print(len(galExt))
    burnin = int(0.25*len(galExt))
    print(burnin)

    
    
    galParNames = np.loadtxt('/home/ADF/middlehr/repositories/PTAInference/runs/galaxyModel_ext/pars.txt',dtype=str)


    prior = np.genfromtxt('/home/ADF/middlehr/repositories/PTAInference/runs/galaxyModel_ext/prior.dat')
    print(len(prior))
    #tau_0, f_0, M*, e_0, chi, beta_f, alpha_tau, beta_tau

    tau0Index = list(galParNames).index('tau0')
    tau0      = burninAndThin(galExt[:,tau0Index],burnin,thin)
    tau0Prior = burninAndThin(prior[:,tau0Index],burninP,thinP)

    f0Index = list(galParNames).index('f0')
    f0      = burninAndThin(galExt[:,f0Index],burnin,thin)
    f0Prior = burninAndThin(prior[:,f0Index],burninP,thinP)

    MstarIndex = list(galParNames).index('Mstar')
    Mstar      = burninAndThin(galExt[:,MstarIndex],burnin,thin)
    MstarPrior = burninAndThin(prior[:,MstarIndex],burninP,thinP)

    e0Index = list(galParNames).index('e0')
    e0      = burninAndThin(galExt[:,e0Index],burnin,thin)
    e0Prior = burninAndThin(prior[:,e0Index],burninP,thinP)
    
    betafIndex = list(galParNames).index('betaf')
    betaf      = burninAndThin(galExt[:,betafIndex],burnin,thin)
    betafPrior = burninAndThin(prior[:,betafIndex],burninP,thinP)
  
    alphatauIndex = list(galParNames).index('alphatau')
    alphatau      = burninAndThin(galExt[:,alphatauIndex],burnin,thin)
    alphatauPrior = burninAndThin(prior[:,alphatauIndex],burninP,thinP)
    
    betatauIndex = list(galParNames).index('betatau')
    betatau      = burninAndThin(galExt[:,betatauIndex],burnin,thin)
    betatauPrior = burninAndThin(prior[:,betatauIndex],burninP,thinP)
    
    


    M0Index = list(galParNames).index('M0')
    M0      = burninAndThin(galExt[:,M0Index],burnin,thin)   
    M0Prior = burninAndThin(prior[:,M0Index],burninP,thinP)   

    alpha0Index = list(galParNames).index('alpha0')
    alpha0      = burninAndThin(galExt[:,alpha0Index],burnin,thin)
    alpha0Prior = burninAndThin(prior[:,alpha0Index],burninP,thinP)

    alphaIIndex = list(galParNames).index('alphaI')
    alphaI      = burninAndThin(galExt[:,alphaIIndex],burnin,thin)
    alphaIPrior = burninAndThin(prior[:,alphaIIndex],burninP,thinP)
    



    alphafIndex = list(galParNames).index('alphaf')
    alphaf      = burninAndThin(galExt[:,alphafIndex],burnin,thin)
    alphafPrior = burninAndThin(prior[:,alphafIndex],burninP,thinP)
    
    gammafIndex = list(galParNames).index('gammaf')
    gammaf      = burninAndThin(galExt[:,gammafIndex],burnin,thin)
    gammafPrior = burninAndThin(prior[:,gammafIndex],burninP,thinP)


    

    gammatauIndex = list(galParNames).index('gammatau')
    gammatau      = burninAndThin(galExt[:,gammatauIndex],burnin,thin)
    gammatauPrior = burninAndThin(prior[:,gammatauIndex],burninP,thinP)


    alphastarIndex = list(galParNames).index('alphastar')
    alphastar      = burninAndThin(galExt[:,alphastarIndex],burnin,thin)
    alphastarPrior = burninAndThin(prior[:,alphastarIndex],burninP,thinP)

    epsilonIndex = list(galParNames).index('epsilon')
    epsilon      = burninAndThin(galExt[:,epsilonIndex],burnin,thin)
    epsilonPrior = burninAndThin(prior[:,epsilonIndex],burninP,thinP)
    

    zeta0Index = list(galParNames).index('zeta0')
    zeta0      = burninAndThin(galExt[:,zeta0Index],burnin,thin)
    zeta0Prior = burninAndThin(prior[:,zeta0Index],burninP,thinP)


    # do parameter swap
    # swap betaf and alphaf
    data = np.atleast_1d([tau0,    
                          alphatau,
                          betatau,
                          f0,
                          betaf,
                          Mstar,
                          e0,
                          zeta0]).T
    '''
                              M0,
                          alpha0,
                          alphaI,                   
                          alphaf,
                          gammaf,
                          gammatau,
                          alphastar,
                          epsilon,
                          zeta0]).T
    '''
    priorSamples = np.atleast_1d([tau0Prior,
                                  alphatauPrior,
                                  betatauPrior,
                                  f0Prior,
                                  betafPrior,
                                  MstarPrior,
                                  e0Prior,
                                  zeta0Prior]).T
    '''
                                  M0Prior,
                                  alpha0Prior,
                                  alphaIPrior,
                                  alphafPrior,
                                  gammafPrior,
                                  gammatauPrior,
                                  alphastarPrior,
                                  epsilonPrior,
                                  zeta0Prior]).T
    '''
    parameterLabels = ([r'$\tau_0$',
                        r'$\alpha_{\tau}$',
                        r'$\beta_{\tau}$',
                        r'$f_0$',
                        r'$\beta_f$',
                        r'$\log_{10} M_{*}$',
                        r'$e_0$',
                        r'$\log_{10}\zeta_0$'])
    '''
                        r'$\log_{10} M_0$', 
                        r'$\alpha_0$',
                        r'$\alpha_I$',
                        r'$\alpha_f$',
                        r'$\gamma_f$',
                        r'$\gamma_{\tau}$',
                        r'$\alpha_{*}$',
                        r'$\epsilon$',
                        r'$\log_{10}\zeta_0$'])
    '''


    bounds=[(0.1,5.),   #tau0
            (-0.5,0.5),  #alphatau
            (-3.,1.),    #betatau
            (0.01,0.05), #f0
            (0.,2.),     #betaf
            (7.75,8.75), # M*
            (0.01,0.99), #e0
            (-2.,2.)]    #zeta

            

            
    figure = corner.corner(data,
                           levels=(0.5,0.9),
                           labels=parameterLabels,
                           quantiles=(0.05,0.50,0.95),
                           show_titles=True,
                           title_fmt='.2f',
                           plot_density=True,
                           range=bounds)
                           #smooth1d=True)#,
                           #normed=True)

    corner.corner(priorSamples,
                  fig=figure,
                  no_fill_contours=True,
                  plot_datapoints=False,
                  plot_density=False,
                  levels=([1.]),
                  color='g',
                  linewidth=100,
                  range=bounds)
                  #smooth1d=True)#,
                  #normed=True)
    corner.corner(data, 
                  fig=figure,
                           levels=(0.5,0.9),
                           labels=parameterLabels,
                           quantiles=(0.05,0.50,0.95),
                           show_titles=True,
                           title_fmt='.2f',
                           plot_density=True,
                           range=bounds)
                           #smooth1d=True)#,
                           #normed=True)

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

    thin=2
    mStarBurninPrior = prior[burnin:,mStarIndex]
    mStarPrior       = mStarBurninPrior[0::thin]

    tauBurninPrior = prior[burnin:,tauIndex]
    tauPrior       = tauBurninPrior[0::thin] 

    alphaTauBurninPrior = prior[burnin:,alphaTauIndex]
    alphaTauPrior       = alphaTauBurninPrior[0::thin]
    
    betaTauBurninPrior = prior[burnin:,betaTauIndex]
    betaTauPrior       = betaTauBurninPrior[0::thin]

    priorSamples = np.atleast_1d([tauPrior,alphaTauPrior,betaTauPrior,mStarPrior]).T
    print(np.shape(priorSamples))

    '''
    plt.savefig('forPaper/gal_selectedCorner_fBins{}.pdf'.format(nFreqBins))
    plt.savefig('forPaper/gal_selectedCorner_fBins{}.png'.format(nFreqBins))
    plt.show()

    return 
    

allCorner(thin=2)    
#selectedCorner(thin=2)
