### MODULES TO IMPORT

from __future__ import division  # divide integers as if they were floating point
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import interpolate
from scipy import signal
from scipy.misc import logsumexp
import emcee
from scipy.interpolate import interp1d
import matplotlib.mlab as mlab
import sys
import time
import corner
from BBH_analysis_functions import *
# LIGO-specific stuff:
import lal as lal
import lalsimulation as lalsim
import readligo as rl

### Set up to import asd file
asd_file = "asd.txt"
asd_columns=["freq","strain"]



### Waveform generation parameters 
fmin = 20.
fmax = 1024. 
deltaF = 1./8.
scale = 20


clight = 2.99792458e8                # m/s
G = 6.67259e-11                      # m^3/kg/s^2 
MSol = 1.989e30                      # kg

fseries = np.linspace(0, fmax, int(fmax/deltaF) + 1)

params = read_bbh_file("10000_BBH_events_updated.txt",1)["0"]


### Reads in the asd file data
asd_data=read_file(asd_file,asd_columns)
dfreq=asd_data["freq"][1]-asd_data["freq"][0]
first_freqs=np.arange(asd_data["freq"][0],-dfreq,-dfreq)[::-1]
full_freqs=np.append(first_freqs,asd_data["freq"])
full_strain = np.pad(asd_data["strain"],(first_freqs.size,0),'constant', constant_values=(10.e-10, 0))
asd_func = interpolate.interp1d(full_freqs, full_strain)
psd = (asd_func(np.abs(fseries))**2)

'''
### SANITY CHECK:: GRAPHS

# Generates a waveform
htilde_L = compute_waveform_at_detector_L(fmin,fmax,deltaF,scale*params["m1_source"],scale*params["m2_source"],params["chi1L"],
                                          params["chi2L"],params["chip"],params["iota"],params["alpha"],
                                          params["luminosity_distance"],params["psi"],params["ra"],params["dec"])

#plots fft vs freq
plt.figure(figsize=(30,20))
plt.subplot(3,1,1)
plt.plot(fseries,np.absolute(htilde_L))
plt.xlabel('frequency (Hz)')
plt.ylabel('|waveform fft|')
plt.xlim(0,1200)
plt.grid()

waveform_t = np.fft.ifft(htilde_L)
fs = deltaF*waveform_t.size
t = np.linspace(0, waveform_t.size/fs,waveform_t.size)
time = shift_time(t,waveform_t, deltaF)

plt.subplot(3,1,2)
plt.plot(time,np.real(waveform_t))
plt.xlabel('time (merger at t=0)')
plt.ylabel('waveform')
plt.grid()

# plots psd data
plt.subplot(3,1,3)
plt.loglog(fseries,np.absolute(htilde_L)*np.sqrt(fseries),label="waveform*sqrt(f)")
plt.loglog(fseries,np.sqrt(psd),label="asd")
plt.xlabel('frequency (Hz)')
plt.ylabel('strain noise ASD (strain/rtHz), template h(f)*rt(f)')
plt.legend(loc='best')
plt.xlim(10,1200)
plt.grid()
plt.ylim(1e-24,1e-18)
plt.savefig("sanitycheck1.pdf")
'''


'''
### MODULATION SANITY CHECK: GRAPHS

lamda_amp = -.8
lamda_freq = .2

htilde_L = compute_waveform_at_detector_L(fmin,fmax,deltaF,scale*params["m1_source"],scale*params["m2_source"],params["chi1L"],
                                          params["chi2L"],params["chip"],params["iota"],params["alpha"],
                                          params["luminosity_distance"],params["psi"],params["ra"],params["dec"])

mod_fft = apply_modulation(htilde_L,fseries,scale*(params["m1_source"]+params["m2_source"]),lamda_amp,lamda_freq)

waveform_t = np.fft.ifft(htilde_L)
fs = deltaF*waveform_t.size
t = np.linspace(0, waveform_t.size/fs,waveform_t.size)
time = shift_time(t,waveform_t,deltaF)

plt.figure(figsize=(25,20))
plt.subplot(3,1,1)
plt.plot(time,waveform_t,label="orig")
plt.plot(time,np.fft.ifft(mod_fft),label="mod")
plt.xlabel('t')
plt.ylabel('waveform')
plt.legend(loc="upper left")
plt.grid()

plt.subplot(3,1,2)
plt.plot(time,waveform_t,label="orig")
plt.plot(time,np.fft.ifft(mod_fft),label="mod")
plt.xlabel('t')
plt.ylabel('waveform')
plt.grid()
plt.legend(loc="upper left")
plt.xlim(-.2,.1)

plt.subplot(3,1,3)
plt.plot(fseries,htilde_L,label="orig")
plt.plot(fseries,mod_fft,label="mod")
plt.xlabel('f')
plt.ylabel('waveform fft')
plt.legend(loc="upper left")
plt.grid()
plt.savefig("sanitycheck2.pdf")
'''

### SETTING UP THE EMCEE DEPENCENDIES

# for now, assume |lambda|<=1

# Define the posterior PDF
# Reminder: post_pdf(theta, data) = likelihood(data, theta) * prior_pdf(theta)
# We take the logarithm since emcee needs it.

# As prior, we assume an 'uniform' prior (i.e. constant prob. density)
def lnprior(theta):
    amp_guess, freq_guess = theta
    # if -10. < amp_guess < 10. and -10. < freq_guess < 10.0:
       # return 0.0
    #return -np.inf
    return -.5*np.log(amp_guess**2.+freq_guess**2.)

# As likelihood, we assume the chi-square. Note: we do not even need to normalize it.
def lnlike(theta, dataL1, psdL1, fmin, fmax, deltaF,scale,m1,m2,chi1L,chi2L,chip,iota,alpha,dist,psi,ra,dec): 
    amp_guess, freq_guess = theta
    
    template_mod_L = generate_mod_waveform_fft_at_detector(fmin,fmax,deltaF,scale*m1,scale*m2,
                                          chi1L,chi2L,chip,iota,alpha,
                                          dist,psi,amp_guess,
                                                   freq_guess,ra,dec)

    dh_L =  deltaF*4*dataL1.conjugate()*template_mod_L / psdL1
    hh_L = deltaF*4.*np.sum( template_mod_L.conjugate()*template_mod_L/psdL1 ).real
    dd_L = deltaF*4.*np.sum( dataL1.conjugate()*dataL1/psdL1 ).real

   
    scores_L = len(dataL1)*( -0.5* ( -2*np.fft.irfft(dh_L)) )
    logL = logsumexp( scores_L ).real - 0.5*( hh_L + dd_L ) + np.log(2./(2.*fmax))
    
    #dh_L =  deltaF*4*np.sum(dataL1.conjugate()*template_mod_L / psdL1).real
    #hh_L = deltaF*4.*np.sum( template_mod_L.conjugate()*template_mod_L/psdL1 ).real
    #dd_L = deltaF*4.*np.sum( dataL1.conjugate()*dataL1/psdL1 ).real
    #logL = dh_L-hh_L-dd_L
    
    return logL

def lnprob(theta, dataL1, psdL1, fmin, fmax, deltaF,scale,m1,m2,chi1L,chi2L,chip,iota,alpha,dist,psi,ra,dec):
    lp = lnprior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + lnlike(theta, dataL1, psdL1, fmin, fmax, deltaF,scale,m1,m2,chi1L,chi2L,chip,iota,alpha,dist,psi,ra,dec)


### TEST RUNS FOR KNOWN VALUES OF LAMDAS

### TEST #1:

def run_emcee_test(lamda_amp_target,lamda_freq_target,fmin,fmax,deltaF,scale,m1,m2,chi1L,chi2L,chip,iota,alpha,dist,psi,ra,dec,psd,
					nwalkers,iterations,chain_name,corner_name,thin,bins,centers):

     dataL1 = generate_mod_waveform_fft_at_detector(fmin,fmax,deltaF,scale*m1,scale*m2,
                                          chi1L,chi2L,chip,iota,alpha,
                                          dist,psi,lamda_amp_target,
                                                   lamda_freq_target,ra,dec)
     
     ndim = 2
     p0 = np.random.uniform(-1.,1.,size=ndim*nwalkers).reshape(nwalkers,ndim)

     sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(dataL1, psd, fmin, fmax, deltaF,scale,m1,m2,chi1L,chi2L,chip,iota,alpha,dist,psi,ra,dec))

     # sampling w/ mcmc
     time0 = time.time()
     width = 50
     for i, result in enumerate(sampler.sample(p0, iterations=iterations, thin=thin)):
        n = int((width+1) * float(i) / iterations)
        sys.stdout.write("\r[{0}{1}]".format('#' * n, ' ' * (width - n)))
     time1=time.time()
     print time1-time0

     # graphing the chains
     amp = []
     freq = []
     for i in sampler.chain:
        amp.append([x[0] for x in i])
        freq.append([x[1] for x in i])
     plt.figure(figsize=(20,15))
     plt.subplot(2,1,1)
     for chain in amp:
        plt.plot(chain)
     plt.xlabel("step")
     plt.ylabel("lamda_amp")
     plt.subplot(2,1,2)
     for chain in freq:
        plt.plot(chain)
     plt.xlabel("step")
     plt.ylabel("lamda_freq")
     plt.savefig(chain_name)

     # graphing the corner
     samples = sampler.chain[:, :, :].reshape(-1, ndim)
     fig = corner.corner(samples,labels=["amp","freq"],
                     truths=[lamda_amp_target,lamda_freq_target] )
     fig.savefig(corner_name)

     # returns graph of the probability distribution for lambda
     amp = [x[0] for x in samples]
     amp_vals,edges = np.histogram(amp,bins=bins, normed=True)
     freq = [x[1] for x in samples]
     freq_vals, n = np.histogram(freq,bins=bins, normed=True)
     
     return amp_vals,freq_vals





bins = np.linspace(-1.,1.,41) 
centers=[]
for i in range(len(bins)-1):
	centers.append((bins[i]+bins[i+1])/2.)


x=run_emcee_test(.6,.4,fmin,fmax,deltaF,scale,params["m1_source"],params["m2_source"],params["chi1L"],params["chi2L"],params["chip"],params["iota"],params["alpha"],params["luminosity_distance"],params["psi"],params["ra"],params["dec"],asd_func(np.abs(fseries))**2,150,1000,"chain_83_jeff_64","corner_83_jeff_64",10,bins,centers)

'''
pdf_file = open("pdf_data.txt","w")


nslist=[bins1,bins2]
for j in binslist:
    line = ""
    for i in j:
        line+=str(i)+","
        
    line=line[:-1]+"\n"
    print line
    myfile.write(line)
myfile.close()

print bins1
print bins2")

plt.show()
plt.figure()
plt.subplot(2,1,1)
plt.plot(x[2],x[0])
plt.subplot(2,1,2)
plt.plot(x[2],x[1])
plt.savefig("test")

x=run_emcee_test(-.2,-.8,fmin,fmax,deltaF,scale,params["m1_source"],params["m2_source"],params["chi1L"],params["chi2L"],params["chip"],params["iota"],params["alpha"],params["luminosity_distance"],params["psi"],params["ra"],params["dec"],asd_func(np.abs(fseries))**2,150,1000,"testingchain-2-8","testingcorner-2-8",10)


x=run_emcee_test(-1.,1.,fmin,fmax,deltaF,scale,params["m1_source"],params["m2_source"],params["chi1L"],params["chi2L"],params["chip"],params["iota"],params["alpha"],params["luminosity_distance"],params["psi"],params["ra"],params["dec"],asd_func(np.abs(fseries))**2,150,1000,"testingchain-11","testingcorner-11",10)

#run_emcee_test(-1.,1.,fmin,fmax,deltaF,scale,params["m1_det"],params["m2_det"],params["chi1L"],params["chi2L"],params["chip"],params["iota"],params["alpha"],params["luminosity_distance"],params["psi"],params["ra"],params["dec"],asd_func(np.abs(fseries))**2,100,10000,"chain_100w_10000it_50th_-11","corn_100w_10000it_50th_-11",50)


#run_emcee_test(-.2,-.8,fmin,fmax,deltaF,scale,params["m1_det"],params["m2_det"],params["chi1L"],params["chi2L"],params["chip"],params["iota"],params["alpha"],params["luminosity_distance"],params["psi"],params["ra"],params["dec"],asd_func(np.abs(fseries))**2,100,10000,"chain_100w_10000it_50th_-2-8","corn_100w_10000it_50th_-2-8",50)


'''