import numpy
import scipy
from numpy.random import choice
import matplotlib.pyplot as plt

import sys
#sys.path.insert(1, '/home/ADF/middlehr/repositories/PTAInference/code/galaxy_model/')
#import mergerrate as mr
sys.path.insert(1, '../models/galaxy_model/')
import mergerratemodel as mr
#import fermif
#import arctanf

green='#009E73'  #green

def schechterf(M,M0,phi0,alpha0):
	return phi0*numpy.power(M/M0,1.+alpha0)*numpy.exp(-M/M0)


def Mbh(Mstar,alpha,beta):
	return alpha+beta*numpy.log10(Mstar/1e11)


def conf_hmz(schains, directory, injection=None, data=None, priorChains=None, modelLoc=None, function=None):
    M1 = numpy.linspace(9,12,25)
    q = numpy.linspace(0.25,1,10)
    z = numpy.linspace(0.,1.5,5)
    f = numpy.linspace(-9.,-6.,15)
    if data is None:
        fbin = 1./(60.*3.154e7) # one over twice tobs in s (half bin width)
    else:
        fbin = data[0,0]/3.
    qdiff = (q[1]-q[0])/2.
    zdiff = (z[1]-z[0])/2.
    mc = numpy.linspace(5,11,30)
    mcdiff = (mc[1]-mc[0])/2.
    params = choice(len(schains)-1, 1000, replace=False)
    #params = range(len(schains))

    
    h = []
    nm = []
    nz = []
    
    for k in params:
        initpar = dict(Phi0 = schains[k,0], PhiI = schains[k,1], M0 = schains[k,2], alpha = schains[k,3], alphaI = schains[k,4], f0 = schains[k,5], beta = schains[k,6], gamma = schains[k,7], delta = schains[k,8], t0 = schains[k,9], epsilon = schains[k,10], zeta = schains[k,11], eta = schains[k,12], Ms = schains[k,13], theta = schains[k,14], sigma = schains[k,15], e0 = schains[k,16], rho = schains[k,17])
        mrate = mr.mergerrate(M1,q,z,f,**initpar)
        merger = mrate.hmodelt(fbin=None)
        out = numpy.sum(merger[1],axis=1)
        outm = numpy.multiply(out,2.*mrate.zpdiff*mcdiff)
        outm = numpy.sum(outm,axis=1)
        outz = numpy.sum(out,axis=0)
        h.append(merger[0])
        nm.append(outm)
        nz.append(outz)

    numpy.savetxt('{0}/h_lines.dat'.format(directory),h,fmt="%f")
    numpy.savetxt('{0}/m_lines.dat'.format(directory),nm,fmt="%e")
    numpy.savetxt('{0}/z_lines.dat'.format(directory),nz,fmt="%e")
    
    h = numpy.genfromtxt('{0}/h_lines.dat'.format(directory))
    nm = numpy.genfromtxt('{0}/m_lines.dat'.format(directory))
    nz = numpy.genfromtxt('{0}/z_lines.dat'.format(directory))

    down1,down5,down32,up68,up95,up99,med = numpy.percentile(h,(0.5,5,16,84,95,99.5,50),axis=0)
    mdown1,mdown5,mdown32,mup68,mup95,mup99,mmed = numpy.percentile(nm,(0.5,5,16,84,95,99.5,50),axis=0)
    zdown1,zdown5,zdown32,zup68,zup95,zup99,zmed = numpy.percentile(nz,(0.5,5,16,84,95,99.5,50),axis=0)
    mup99 = numpy.clip(mup99,1e-300,numpy.inf)
    mup95 = numpy.clip(mup95,1e-300,numpy.inf)
    mup68 = numpy.clip(mup68,1e-300,numpy.inf)
    mdown1 = numpy.clip(mdown1,1e-300,numpy.inf)
    mdown5 = numpy.clip(mdown5,1e-300,numpy.inf)
    mdown32 = numpy.clip(mdown32,1e-300,numpy.inf)
    mmed = numpy.clip(mmed,1e-300,numpy.inf)

    #sh = numpy.vstack((f,down1,down5,down32,up68,up95,up99,med)).T
    #numpy.savetxt("{0}/h.dat".format(directory),sh,fmt='%f')
    #sm = numpy.vstack((mc,mdown1,mdown5,mdown32,mup68,mup95,mup99,mmed)).T
    #numpy.savetxt("{0}/nm.dat".format(directory),sm,fmt='%e')
    #sz = numpy.vstack((z,zdown1,zdown5,zdown32,zup68,zup95,zup99,zmed)).T
    #numpy.savetxt("{0}/nz.dat".format(directory),sz,fmt='%f')
    

    # prior plotting
    if priorChains is not None:
        priorParams = choice(len(priorChains)-1, 1000, replace=False)
        #priorParams = range(len(priorChains))


        hPrior = []
        nmPrior = []
        nzPrior = []
        for k in priorParams:
            initpar = dict(Phi0 = priorChains[k,0], PhiI = priorChains[k,1], M0 = priorChains[k,2], alpha = priorChains[k,3], alphaI = priorChains[k,4], f0 = priorChains[k,5], beta = priorChains[k,6], gamma = priorChains[k,7], delta = priorChains[k,8], t0 = priorChains[k,9], epsilon = priorChains[k,10], zeta = priorChains[k,11], eta = priorChains[k,12], Ms = priorChains[k,13], theta = priorChains[k,14], sigma = priorChains[k,15], e0 = priorChains[k,16], rho = priorChains[k,17])
            mrate = mr.mergerrate(M1,q,z,f,**initpar)
            merger = mrate.hmodelt(fbin=None)
            out = numpy.sum(merger[1],axis=1)
            outm = numpy.multiply(out,2.*mrate.zpdiff*mcdiff)
            #outm = numpy.multiply(out,mrate.zpdiff/mcdiff)
            outm = numpy.sum(outm,axis=1)
            outz = numpy.sum(out,axis=0)
            hPrior.append(merger[0])
            nmPrior.append(outm)
            nzPrior.append(outz)

        numpy.savetxt('{0}/h_linesprior.dat'.format(directory),hPrior,fmt="%f")
        numpy.savetxt('{0}/m_linesprior.dat'.format(directory),nmPrior,fmt="%e")
        numpy.savetxt('{0}/z_linesprior.dat'.format(directory),nzPrior,fmt="%e")

        hPrior = numpy.genfromtxt('{0}/h_linesprior.dat'.format(directory))
        nmPrior = numpy.genfromtxt('{0}/m_linesprior.dat'.format(directory))
        nzPrior = numpy.genfromtxt('{0}/z_linesprior.dat'.format(directory))

        prior1,prior5,prior95,prior99 = numpy.percentile(hPrior,(0.5,5,95,99.5),axis=0)
        mprior1,mprior5,mprior95,mprior99 = numpy.percentile(nmPrior,(0.5,5,95,99.5),axis=0)
        zprior1,zprior5,zprior95,zprior99 = numpy.percentile(nzPrior,(0.5,5,95,99.5),axis=0)

    if injection is not None:
        mratei = mr.mergerrate(M1,q,z,f,**injection)
        inj0 = mratei.hmodelt(fbin)
        inji = mratei.hmodelt()
        h0 = inj0[0]
        hi = inji[0]
        mergeri = numpy.sum(inj0[1],axis=1)
        nm0 = numpy.multiply(mergeri,mratei.zpdiff/mcdiff)
        nm0 = numpy.sum(nm0,axis=1)
        nz0 = numpy.sum(mergeri,axis=0)


    f,mc = 10.**f,10.**mc
    down1,down5,down32,up68,up95,up99,med = 10.**down1,10.**down5,10.**down32,10.**up68,10.**up95,10.**up99,10.**med

    plt.figure(figsize=(8,5))
    #plt.locator_params(nbins=4,axis='x')
    #plt.locator_params(nbins=6,axis='y')
    plt.tick_params(axis='x', labelsize='x-large')
    plt.tick_params(axis='y', labelsize='x-large')

    if priorChains is not None:
        #plt.fill_between(f,10.**prior0,10.**prior100, hatch='.', edgecolor=(0,0,0,1.0),facecolor=(1,1,1,1))
        plt.plot(f,10.0**prior1,color=green,linewidth=2,linestyle='dotted')
        plt.plot(f,10.0**prior99,color=green,linewidth=2,linestyle='dotted')
        plt.plot(f,10.0**prior5,color=green,linewidth=2,linestyle='dashed')
        plt.plot(f,10.0**prior95,color=green,linewidth=2,linestyle='dashed')

    plt.plot(f,med,color='k',linewidth=2,label='median, 68%, 90%, 95% confidence levels')
    #plt.fill_between(f,down1,up99,facecolor='k',alpha=0.25)
    plt.fill_between(f,down5,up95,facecolor='k',alpha=0.25)
    plt.fill_between(f,down32,up68,facecolor='k',alpha=0.25)

    if injection is not None:
        plt.plot(f,10.0**h0,color='r',linewidth=2,linestyle='dashdot')
        plt.plot(f,10.0**hi,color='r',linewidth=2,linestyle='dotted')

    if data is not None:
        k = 0
        while k < len(data[:,0]):
            if data[k,3] < 1.0:
                if function == 'fermi':
                    sigma = fermif.fermis(10.0**data[k,1],0.68)
                    fermi_up = fermif.fermih(10.0**data[k,1],sigma,0.95)
                    fermi_down = fermif.fermih(10.0**data[k,1],sigma,0.5)
                    plt.errorbar(data[k,0],fermi_up,yerr=fermi_up-fermi_down,uplims=True,c='b',linewidth=2)
                elif function == 'arctan':
                    sigma = arctanf.arctans(10.0**data[k,1],0.68)
                    arctan_up = arctanf.arctanh(10.0**data[k,1],sigma,0.95)
                    arctan_down = arctanf.arctanh(10.0**data[k,1],sigma,0.5)
                    plt.errorbar(data[k,0],arctan_up,yerr=arctan_up-arctan_down,uplims=True,c='b',linewidth=2)
                elif function == 'ppta':
                    """
                    fraction = 0.8454 #0.8454 is finetuned for specific 4 1.5/T bins ppta datafile
                    sigma = 10.0**data[k,1]/1.e-15/fraction
                    fermi_up = 10.0**data[k,1]
                    fermi_down = sigma*1.e-15
                    plt.errorbar(data[k,0],fermi_up,yerr=fermi_up-fermi_down,uplims=True,c='b',linewidth=2)
                    """
                    data0 = numpy.genfromtxt('pptaData.dat',skip_footer=1)
                    plt.plot(data0[:,0],10.0**data0[:,1],c='orangered',linestyle='solid',linewidth=2)
                    #plt.axvline(6.34e-9,c='orangered',linestyle='dotted',linewidth=2)
                    plt.plot(6.34e-9,1.e-20*6.34e-9**(-2./3.),c='orangered',marker="*",ms=12)
                else: pass
            else:
                plt.errorbar(data[k,0],10.0**data[k,2],yerr=[[10.0**data[k,2]-10.0**(data[k,2]-data[k,5])],[10.0**(data[k,2]+data[k,5])-10.0**data[k,2]]],c='b',linewidth=2)
            k += 1
        #plt.plot(data[:,0],10.0**data[:,1],'k',linestyle='dotted',linewidth=2)
        plt.errorbar(3.17e-8,1.92e-15,yerr=[[0.55e-15],[0.75e-15]],c='orangered',marker="*",ms=12)
    else: pass
    #plt.plot(0.,0.,'b',label='simulated data',linewidth=2)
    plt.xlabel(r'$ f (\mathrm{Hz})$', fontsize='x-large')
    plt.ylabel(r'$ h_c$', fontsize='xx-large')
    plt.xlim(1.e-9,1.e-6)
    plt.ylim(1.e-18,1.e-13)
    plt.xscale('log')
    plt.yscale('log')
    #plt.legend(loc=3)
    plt.savefig("{0}/conf_h.pdf".format(directory),bbox_inches='tight')
    plt.savefig("{0}/conf_h.png".format(directory),bbox_inches='tight')


    plt.clf()
    plt.figure(figsize=(8,5))
    #plt.locator_params(nbins=8,axis='x')
    #plt.locator_params(nbins=12,axis='y')
    plt.tick_params(axis='x', labelsize='x-large')
    plt.tick_params(axis='y', labelsize='x-large')

    if priorChains is not None:
        #plt.fill_between(mc,mprior0,mprior100, hatch='.', edgecolor=(0,0,0,1.0),facecolor=(1,1,1,1))
        plt.plot(mc,mprior1,color=green,linewidth=2,linestyle='dotted')
        plt.plot(mc,mprior99,color=green,linewidth=2,linestyle='dotted')
        plt.plot(mc,mprior5,color=green,linewidth=2,linestyle='dashed')
        plt.plot(mc,mprior95,color=green,linewidth=2,linestyle='dashed')

    plt.plot(mc,mmed,color='k',linewidth=2,label='median, 68%, 90%, 95% confidence levels')
    #plt.fill_between(mc,mdown1,mup99,facecolor='k',alpha=0.25)
    plt.fill_between(mc,mdown5,mup95,facecolor='k',alpha=0.25)
    plt.fill_between(mc,mdown32,mup68,facecolor='k',alpha=0.25)

    if injection is not None:
        plt.plot(mc,nm0,color='r',linewidth=2,linestyle='dashdot',label='injected')

    plt.xlabel(r'$\mathcal{M} (M_\odot)$', fontsize='x-large')
    plt.ylabel(r'$ dN/dVd\log_{10} \mathcal{M} (\mathrm{Mpc}^{-3})$', fontsize='x-large')
    plt.xlim(1e6,1e11)
    #plt.ylim(1e-9,1e-1)
    plt.xscale('log')
    plt.yscale('log')
    #plt.legend(loc=3)
    plt.savefig("{0}/conf_m.pdf".format(directory),bbox_inches='tight')
    plt.savefig("{0}/conf_m.png".format(directory),bbox_inches='tight')

    if modelLoc is not None:
        plotnewmod(modelLoc,'mass')
        plt.savefig("{0}/conf_m_model.pdf".format(directory),bbox_inches='tight')
        plt.savefig("{0}/conf_m_model.png".format(directory),bbox_inches='tight')
    else: pass


    plt.clf()
    plt.figure(figsize=(8,5))
    #plt.locator_params(nbins=8,axis='x')
    #plt.locator_params(nbins=12,axis='y')
    plt.tick_params(axis='x', labelsize='x-large')
    plt.tick_params(axis='y', labelsize='x-large')

    if priorChains is not None:
        #plt.fill_between(z,zprior0,zprior100, hatch='.', edgecolor=(0,0,0,1.0),facecolor=(1,1,1,1))
        plt.plot(z,zprior1,color=green,linewidth=2,linestyle='dotted')
        plt.plot(z,zprior99,color=green,linewidth=2,linestyle='dotted')
        plt.plot(z,zprior5,color=green,linewidth=2,linestyle='dashed')
        plt.plot(z,zprior95,color=green,linewidth=2,linestyle='dashed')

    plt.plot(z,zmed,color='k',linewidth=2,label='median, 68%, 90%, 95% confidence levels')
    #plt.fill_between(z,zdown1,zup99,facecolor='k',alpha=0.25)
    plt.fill_between(z,zdown5,zup95,facecolor='k',alpha=0.25)
    plt.fill_between(z,zdown32,zup68,facecolor='k',alpha=0.25)

    if injection is not None:
        plt.plot(z,nz0,color='r',linewidth=2,linestyle='dashdot',label='injected')

    plt.xlabel(r'$z$', fontsize='xx-large')
    plt.ylabel(r'$ dN/dVdz (\mathrm{Mpc}^{-3})$', fontsize='x-large')
    plt.xlim(0,5)
    #plt.ylim(1e-5,1e-1)
    plt.yscale('log')
    #plt.legend(loc=3)
    plt.savefig("{0}/conf_z.pdf".format(directory),bbox_inches='tight')
    plt.savefig("{0}/conf_z.png".format(directory),bbox_inches='tight')

    if modelLoc is not None:
        plotnewmod(modelLoc,'z')
        plt.savefig("{0}/conf_z_model.pdf".format(directory),bbox_inches='tight')
        plt.savefig("{0}/conf_z_model.png".format(directory),bbox_inches='tight')
    else: pass


    return




def plotnewmod(dataFile, z_or_m):
    newmoddata = numpy.loadtxt('{0}'.format(dataFile,z_or_m)) 
    xvalues=(newmoddata[:,0])
    con997_low=(newmoddata[:,1])
    con997_high=(newmoddata[:,6])
    con2s_low=(newmoddata[:,2])
    con2s_high=(newmoddata[:,5])
    con1s_low=(newmoddata[:,3])
    con1s_high=(newmoddata[:,4])

    xnonzero=[]
    con997_high_nonzero=[]
    con997_low_nonzero=[]

    for i in range(len(xvalues)):
        if con997_low[i]!=0.0 and con997_high[i]!=0.0:
        #if con997_high[i]!=0.0:
            xnonzero.append(xvalues[i])
            con997_high_nonzero.append(con997_high[i])
            con997_low_nonzero.append(con997_low[i])
            itop=i

    checklow=True

    for j in range(len(xvalues)):
        if con997_low[j]==0.0 and checklow==True:
            con997_low[j]=con997_high[0]
        else:
            checklow=False
        if con997_low[j]==0.0 and checklow==False:
            con997_low[j]=con997_high[len(con997_high)-1]
    
    plt.fill_between(xnonzero,con997_low_nonzero,con997_high_nonzero,facecolor='k',edgecolor='k',hatch='XX',alpha=0.5)
    
    return 



def conf_g(schains, directory, injection=None, priorChains=None):

    M1 = numpy.logspace(9,12,25)

    #params = choice(len(schains)-1, 100, replace=False)
    params = range(len(schains))

    n = []
    for k in params:
        gsmf = schechterf(M1,10.**schains[k,2],10.**schains[k,0],schains[k,3])
        n.append(gsmf)

    down1,down5,down32,up68,up95,up99,med = numpy.percentile((n),(0.5,5,16,84,95,99.5,50),axis=0)
    up99 = numpy.clip(up99,1e-300,numpy.inf)
    up95 = numpy.clip(up95,1e-300,numpy.inf)
    up68 = numpy.clip(up68,1e-300,numpy.inf)
    down1 = numpy.clip(down1,1e-300,numpy.inf)
    down5 = numpy.clip(down5,1e-300,numpy.inf)
    down32 = numpy.clip(down32,1e-300,numpy.inf)
    med = numpy.clip(med,1e-300,numpy.inf)

    s = numpy.vstack((M1,down1,down5,down32,up68,up95,up99,med)).T
    numpy.savetxt("{0}/ng.dat".format(directory),s,fmt='%.3e')


    # if plotting prior
    priorn = []

    if priorChains is not None:
        priorparams = choice(len(priorChains)-1, 1000, replace=False)
        #priorparams = range(len(priorChains))
        for k in priorparams:
            gsmf = schechterf(M1,10.**priorChains[k,2],10.**priorChains[k,0],priorChains[k,3])
            priorn.append(gsmf)

        prior0,prior100 = numpy.percentile((priorn),(5.0,95.0),axis=0)
    else: pass

    if injection is not None:
        gsmfi = schechterf(M1,10.**injection['M0'],10.**injection['Phi0'],injection['alpha'])


    plt.figure(figsize=(8,5))
    #plt.locator_params(nbins=12,axis='x')
    #plt.locator_params(nbins=12,axis='y')
    plt.tick_params(axis='x', labelsize='x-large')
    plt.tick_params(axis='y', labelsize='x-large')

    if priorChains is not None:
        #plt.fill_between(M1,prior0,prior100, hatch='.', edgecolor=(0,0,0,1.0),facecolor=(1,1,1,1))
        plt.plot(M1,prior0,color=green,linewidth=2,linestyle='dashed')
        plt.plot(M1,prior100,color=green,linewidth=2,linestyle='dashed')

    plt.plot(M1,med,color='k',linewidth=2,label='median, 68%, 90%, 95% confidence levels')
    #plt.fill_between(M1,down1,up99,facecolor='k',alpha=0.25)
    plt.fill_between(M1,down5,up95,facecolor='k',alpha=0.25)
    plt.fill_between(M1,down32,up68,facecolor='k',alpha=0.25)
    if injection is not None:
        plt.plot(M1,gsmfi,color='r',linewidth=2,linestyle='dashdot',label='injected')
    plt.xlabel(r'$M_G (M_\odot)$', fontsize='x-large')
    plt.ylabel(r'$\Phi (\mathrm{Mpc}^{-3})$', fontsize='x-large')
    plt.xlim(1e9,1e12)
    plt.ylim(1e-7,1e-1)
    plt.xscale('log')
    plt.yscale('log')
    #plt.legend(loc=3)
    plt.savefig("{0}/conf_g.pdf".format(directory),bbox_inches='tight')
    plt.savefig("{0}/conf_g.png".format(directory),bbox_inches='tight')

    return



def conf_s(schains, directory, injection=None, priorChains=None):

    M1 = numpy.logspace(9,12,25)

    #params = choice(len(schains)-1, 100, replace=False)
    params = range(len(schains))

    n = []
    for k in params:
        Mb = 10.**Mbh(M1,schains[k,13],schains[k,14])
        n.append(Mb)

    down1,down5,down32,up68,up95,up99,med = numpy.percentile((n),(0.5,5,16,84,95,99.5,50),axis=0)

    s = numpy.vstack((M1,down1,down5,down32,up68,up95,up99,med)).T
    numpy.savetxt("{0}/ns.dat".format(directory),s,fmt='%.3e')


    # if plotting prior
    priorn = []

    if priorChains is not None:
        priorparams = choice(len(priorChains)-1, 1000, replace=False)
        #priorparams = range(len(priorChains))
        for k in priorparams:
            Mb = 10.**Mbh(M1,priorChains[k,13],priorChains[k,14])
            priorn.append(Mb)

        prior0,prior100 = numpy.percentile((priorn),(5.0,95.0),axis=0)
    else: pass

    if injection is not None:
        Mbi = 10.**Mbh(M1,injection['Ms'],injection['theta'])


    plt.figure(figsize=(8,5))
    #plt.locator_params(nbins=12,axis='x')
    #plt.locator_params(nbins=12,axis='y')
    plt.tick_params(axis='x', labelsize='x-large')
    plt.tick_params(axis='y', labelsize='x-large')

    if priorChains is not None:
        #plt.fill_between(M1,prior0,prior100, hatch='.', edgecolor=(0,0,0,1.0),facecolor=(1,1,1,1))
        plt.plot(M1,prior0,color=green,linewidth=2,linestyle='dashed')
        plt.plot(M1,prior100,color=green,linewidth=2,linestyle='dashed')

    plt.plot(M1,med,color='k',linewidth=2,label='median, 68%, 90%, 95% confidence levels')
    #plt.fill_between(M1,down1,up99,facecolor='k',alpha=0.25)
    plt.fill_between(M1,down5,up95,facecolor='k',alpha=0.25)
    plt.fill_between(M1,down32,up68,facecolor='k',alpha=0.25)
    if injection is not None:
        plt.plot(M1,Mbi,color='r',linewidth=2,linestyle='dashdot',label='injected')
    plt.xlabel(r'$M_{bulge} (M_\odot)$', fontsize='x-large')
    plt.ylabel(r'$M_{BH} (M_\odot)$', fontsize='x-large')
    plt.xlim(1e9,1e12)
    plt.ylim(1e5,1e10)
    plt.xscale('log')
    plt.yscale('log')
    #plt.legend(loc=3)
    plt.savefig("{0}/conf_s.pdf".format(directory),bbox_inches='tight')
    plt.savefig("{0}/conf_s.png".format(directory),bbox_inches='tight')

    return

