#!/usr/bin/env python
# coding: utf-8
import warnings
warnings.simplefilter("ignore") #Used to suppress the RuntimeWarnings. But if you are changing the code for TESTING, MAKE SURE TO COMMENT IT OUT.
'''
This script calculates the RDGW150914 constraints with ptemcee - specifically, the n=1 case varying the fundamental frequencies, tshift = 0. It produces the chain plot, corner plot, parameter constraints, and data plotted with the 1-sigma band. Since we are working specifically for the n=1 case, we also compare alpha_0 and alpha_1 in a histogram, in order to demonstrate that the two are indistinguishable with each other.

--- R^a_{yne} L^i_u, 08/09/2020

This script with the diff specifically changes the ansatz phase parameter to the phase difference parameter.

--- R^a_{yne} L^i_u, 08/18/2020
'''

import numpy as np
import corner
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib import rc
#plt.rcParams['font.family'] = 'DejaVu Sans'
#rc('text', usetex=True)
plt.rcParams.update({'font.size': 16.5})

import ptemcee
from pycbc.pool import choose_pool
import math
import h5py
import inspect
import pandas as pd
import json
import qnm
import random
from scipy.optimize import minimize

#Remember to change the following global variables
#rootpath: root path to nr data
#npoints: number of points you re using for your sampling
#nmax: tone index --> nmax = 0 if fitting the fundamental tone
#tshift: time shift after the strain peak
#vary_fund: whether you vary the fundamental frequency. Works in the model_dv function.

rootpath= "/work/rayne.liu"#"/Users/RayneLiu"
nmax=1
tshift=0
vary_fund = True

#sampler parameters
npoints=137
nwalkers = 42
ntemps=42
ndim = int(4*(nmax+1))
burnin = 42  #How many points do you burn before doing the corner plot. You need to watch the convergence of the chain plot a bit.
            #This is trivial but often forgotten: this cannot be more than npoints! Usually 1/5~1/4 npoints is what I observe.
numbins = 42 #corner plot parameter - how many bins you want
datacolor = '#105670' #'#4fa3a7'
pkcolor = '#f2c977' #'#ffb45f'
mediancolor = '#f7695c' #'#9b2814'

#Import data and necessary functions

#TimeOfMaximum
def FindTmaximum(y):
    #Determines the maximum absolute value of the complex waveform
    absval = y[:,1]*y[:,1]+y[:,2]*y[:,2]
    vmax=np.max(absval)
    index = np.argmax(absval == vmax)
    timemax=gw_sxs_bbh_0305[index,0]
    return timemax



#This loads the 22 mode data
gw = {}
gw["SXS:BBH:0305"] = h5py.File(rootpath+"/git/rdstackingproject/SXS/BBH_SKS_d14.3_q1.22_sA_0_0_0.330_sB_0_0_-0.440/Lev6/rhOverM_Asymptotic_GeometricUnits_CoM.h5", 'r')
gw_sxs_bbh_0305 = gw["SXS:BBH:0305"]["Extrapolated_N2.dir"]["Y_l2_m2.dat"]

# Remember to download metadata.json from the simulation with number: 0305. Download Lev6/metadata.json
# This postprocesses the metadata file to find the final mass and final spin
metadata = {}
with open(rootpath+"/git/rdstackingproject/SXS/BBH_SKS_d14.3_q1.22_sA_0_0_0.330_sB_0_0_-0.440/Lev6/metadata.json") as file:
    metadata["SXS:BBH:0305"] = json.load(file)

af = metadata["SXS:BBH:0305"]['remnant_dimensionless_spin'][-1]
mf = metadata["SXS:BBH:0305"]['remnant_mass']


#times --> x axis of your data
times = gw_sxs_bbh_0305[:,0]
tmax=FindTmaximum(gw_sxs_bbh_0305)
t0=tmax +tshift

#Select the data from t0 onwards
position = np.argmax(times >= (t0))
gw_sxs_bbh_0305rd=gw_sxs_bbh_0305[position:-1]
timesrd=gw_sxs_bbh_0305[position:-1][:,0]

# Depending on nmax, you load nmax number of freqs. and damping times from the qnm package
omegas = [qnm.modes_cache(s=-2,l=2,m=2,n=i)(a=af)[0] for i in range (0,nmax+1)]




#Fitting
#RD model for nmax tones. Amplitudes are in (xn*Exp[i yn]) version. But the parameters for the phase are y0, y1-y0, y2-y0, ... . Used here.
def model_dv(theta):
    #x0, y0= theta
    #Your nmax might not align with the dim of theta. Better check it here.
    assert int(len(theta)/4) == nmax + 1, 'Please recheck your n and parameters'
    dim =int(len(theta)/4)  
    w = (np.real(omegas))/mf
    tau=-1/(np.imag(omegas))*mf      
    
    avars = theta[ : (dim)]
    bvars = theta[(dim) : 2*(dim)]
    xvars = theta[2*(dim) : 3*(dim)]
    yvars = theta[3*(dim) : ] #The first variable is y0, the second y1 - y0, the third y2-y0, and so on. We can reconstruct the phase array from this.
    phase = np.zeros(dim)
    phase[0] = yvars[0]
    for p in range (1, dim):
        phase[p] = yvars[p] + yvars[0]
    
    if vary_fund == False:
        avars[0]=0
        bvars[0]=0
        
    ansatz = 0
    for i in range (dim):
        ansatz += (xvars[i]*np.exp(1j*phase[i]))*np.exp(-(timesrd-timesrd[0])/(tau[i]*(1+bvars[i]))) * (np.cos((1+avars[i])*w[i]*timesrd)-1j*np.sin((1+avars[i])*w[i]*timesrd))
    # -1j to agree with SXS convention
    return ansatz


# Logprior distribution. It defines the allowed range my variables can vary over. 
#It works for the (xn*Exp[iyn]) version. 
def log_prior(theta): 
    #Warning: we are specifically working with nmax=1 so here individual prior to the parameters are manually adjusted. This does not apply to all other nmax's.
    a_0 = theta[0]
    a_1 = theta[1]
    b_0 = theta[2]
    b_1 = theta[3]
    x_0 = theta[4]
    x_1 = theta[5]
    y_0 = theta[6]
    y_1_0 = theta[7] #This is y_1 - y_0

    if all([nmax == 1, 0 <= tshift <= 5, vary_fund == True, -1.2 <= a_0 <= 1.2, -1.5 <= a_1 <= 1.5, -9.0 <= b_0 <= 9.0, -10.0 <= b_1 <= 21.0, 0 <= x_0 <= 1.6, 0 <= x_1 <= 1.4, 0 <= y_0 <= 2*np.pi, -2*np.pi <= y_1_0 <= 2*np.pi]):        
        return 0.0
    elif all([nmax == 1, tshift == 19, vary_fund == True, -10.0 <= a_0 <= 10.0, -10.0 <= a_1 <= 10.0, -9.0 <= b_0 <= 20.0, -9.0 <= b_1 <= 25.0, 0 <= x_0 <= 0.8, 0 <= x_1 <= 0.9, 0 <= y_0 <= 2*np.pi, -2*np.pi <= y_1_0 <= 2*np.pi]):
        return 0.0
    
#PAY EXTRA ATTENTION TO THESE TWO CASES, SINCE WE SHIFTED y_0. THE CORNER PLOTS LABELS NEED TO BE TREATED WITH CARE.
    elif all([nmax == 1, 0 <= tshift <= 5, vary_fund == False, -10.0 <= a_0 <= 10.0, -1.5 <= a_1 <= 1.5, -9.0 <= b_0 <= 9.0, -10.0 <= b_1 <= 16.0, 0 <= x_0 <= 2.2, 0 <= x_1 <= 2.5, 0 <= y_0-np.pi <= 2*np.pi, -2*np.pi <= y_1_0 <= 2*np.pi,]):
        return 0.0
    elif all([nmax == 1, tshift == 19, vary_fund == False, -10.0 <= a_0 <= 10.0, -10.0 <= a_1 <= 10.0, -9.0 <= b_0 <= 9.0, -10.0 <= b_1 <= 32.0, 0 <= x_0 <= 0.6, 0 <= x_1 <= 1.0, 0 <= y_0-np.pi <= 2*np.pi, -2*np.pi <= y_1_0 <= 2*np.pi]):
        return 0.0

    return -np.inf


# LogLikelihood function. It is just a Gaussian loglikelihood based on computing the residuals^2
def log_likelihood(theta):
    modelev = model_dv(theta)
    result = -np.sum((gw_sxs_bbh_0305rd[:,1] - (modelev.real))**2+(gw_sxs_bbh_0305rd[:,2] - (modelev.imag))**2)
    if np.isnan(result):
        return -np.inf
    return result


# Logposterior distribution for the residuals case.
# The evidence is just a normalization factor
def log_probability(theta):
    lp = log_prior(theta)
    print('lp:')
    print(lp)
    if not np.isfinite(lp):
        return -np.inf
    return lp + log_likelihood(theta)




#Fit with ptemcee
#Set the number of cores of your processors
pool = choose_pool(24)
pool.size = 24
vary_param = float(vary_fund)
np.random.seed(42)
pos = np.array([[random.uniform(-0.1,0.1), random.uniform(-0.1,0.1), 4.28313743e-01, random.uniform(2.5, 2.6) + (1-vary_param) * np.pi]])
for i in range (1,nmax+1):
    pos_aux = np.array([[random.uniform(-0.1,0.1), random.uniform(-0.1,0.1), random.uniform(0.3,0.4), random.uniform(0.01, 0.02)]])
    pos = np.concatenate((pos, pos_aux), axis = 0)
    
pos = pos.T.flatten()
pos = list(pos)
pos += 1e-5 * np.random.randn(ntemps, nwalkers, ndim)
#print(pos)
sampler = ptemcee.Sampler(nwalkers, ndim, log_likelihood, log_prior, ntemps=ntemps, pool=pool)
sampler.run_mcmc(pos,npoints)


#Define labels and start plotting
paramlabels_a = [r'$\alpha_'+str(i)+'$' for i in range (nmax+1)]
paramlabels_b = [r'$\beta_'+str(i)+'$' for i in range (nmax+1)]
paramlabels_x = [r'$x_'+str(i)+'$' for i in range (nmax+1)]
paramlabels_y = [r'$y_0$']+[r'$y_'+str(i)+'-y_0$' for i in range (1, nmax+1)] if vary_fund == True else [r'$y_0 - \pi$']+[r'$y_'+str(i)+'-y_0$' for i in range (1, nmax+1)]
paramlabels = paramlabels_a + paramlabels_b + paramlabels_x + paramlabels_y
#Need to delete alpha_0 and alpha_1 for the corner plot
paramlabels_corner = paramlabels_a + paramlabels_b + paramlabels_x + paramlabels_y
if vary_fund == False:
    del paramlabels_corner[0]
    del paramlabels_corner[1]



#Chain plot
fig, axes = plt.subplots(4*(nmax+1), 1, sharex=True, figsize=(12, 9*(nmax+1)))
for i in range(4*(nmax+1)):
    axes[i].plot(sampler.chain[0,:, :, i].T, color="k", alpha=0.4)
    axes[i].yaxis.set_major_locator(MaxNLocator(5))
    axes[i].set_ylabel(paramlabels[i])
axes[-1].set_xlabel('Iterations')
#plt.show()
fig.savefig(rootpath+'/git/rdstackingproject/plotsmc/Diff_vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_chain.png', format = 'png', dpi = 384, bbox_inches = 'tight')



#Burn samples, calculate peak likelihood value (not necessarily so in atlas) and make corner plot
samples = sampler.chain[0,:, burnin:, :].reshape((-1, ndim))
#samples for corner plot
samples_corn = samples if vary_fund == True else np.delete(samples, np.s_[0,2], 1)
#print('Values with peak likelihood:')
lglk = np.array([log_likelihood(samples[i]) for i in range(len(samples))])
pk = samples[np.argmax(lglk)]
#print('pk:')
#print(pk)
pk_corn = pk if vary_fund == True else np.delete(pk, [0,2])
#print('pkFalse:')
#print(pk)
 
#Now calculate median (50-percentile) value
median = np.median(samples_corn, axis=0)
#print(samples)
#print(samples_corn)

figcorn = corner.corner(samples_corn, bins = numbins, hist_bin_factor = 5, color = datacolor, truths=pk_corn, truth_color = pkcolor, plot_contours = True, labels = paramlabels_corner, quantiles=(0.05, 0.16, 0.5, 0.84, 0.95), levels=[1-np.exp(-0.5), 1-np.exp(-1.64 ** 2/2)], show_titles=True)

#Extract the axes in order to add more important line plots
naxes = len(pk_corn)
axes = np.array(figcorn.axes).reshape((naxes, naxes))

# Loop over the diagonal
for i in range(naxes):
    ax = axes[i, i]
    ax.axvline(median[i], color=mediancolor)

# Loop over the histograms
for yi in range(naxes):
    for xi in range(yi):
        ax = axes[yi, xi]
        ax.axvline(median[xi], color=mediancolor)
        ax.axhline(median[yi], color=mediancolor)
        ax.plot(median[xi], median[yi], color = mediancolor, marker = 's')
figcorn.savefig(rootpath+'/git/rdstackingproject/plotsmc/Diff_vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_corner.pdf', format = 'pdf')


#Now export the covariance and correlation matrices, as well as the logevidence
#The covariant matrix
samplesT = samples_corn.T
cov_matr=np.cov(samplesT)
#print('Covariant matrix')
df1 = pd.DataFrame(cov_matr, index = paramlabels_corner, columns=paramlabels_corner)
df1.to_hdf(rootpath+'/git/rdstackingproject/plotsmc/Diff_vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_covariance.h5', key='df')

#The correlation matrix
corr_matr=np.corrcoef(samplesT)
#print('Correlation matrix')
df2 = pd.DataFrame(corr_matr,index = paramlabels_corner, columns=paramlabels_corner)
df2.to_hdf(rootpath+'/git/rdstackingproject/plotsmc/Diff_vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_correlation.h5', key='df')

#The logevidence (with error) and save it
logevidence = sampler.log_evidence_estimate(fburnin=0.2)
df3 = pd.DataFrame(logevidence,index = [r'logevidence', r'error'], columns=[r'vary='+str(vary_fund)+', nmax='+str(nmax)+', tshift='+str(tshift)+', '+str(npoints)+'pts'])
df3.to_hdf(rootpath+'/git/rdstackingproject/plotsmc/Diff_vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_logevidence.h5', key='df')

#The autocorrelation time, this time I don't intend to delete alpha0 and beta0, just keep the nan's there (if they are nan's) since we don't really use this data, only look at it. It's good to keep it for sanity check's sake.
atcrr = sampler.get_autocorr_time()[0]
df4 = pd.DataFrame(atcrr,index = paramlabels, columns=[r'Autocorrelation time, vary='+str(vary_fund)+', nmax='+str(nmax)+', tshift='+str(tshift)+', '+str(npoints)+'pts'])
df4.to_hdf(rootpath+'/git/rdstackingproject/plotsmc/Diff_vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_autocrrtime.h5', key='df')    




#Now plot alpha_0 on top of alpha_1 - only on top, not stacked, and only valid for vary_fund == True
if vary_fund == True:
    samplea0 = samplesT[0]
    samplea1 = samplesT[1]
    fighist1 = plt.figure(figsize = (8, 6))
    n0, bins0, patches0 = plt.hist(samplea0, bins = numbins * 6, alpha = 0.5, label = r'$\alpha_0$')
    n1, bins1, patches1 = plt.hist(samplea1, bins = numbins * 10, alpha = 0.5, label = r'$\alpha_1$')
    #n01, bins01, patches01 = plt.hist([samplea0, samplea1], numbins * 10, rwidth = 3, alpha = 0.5, label = [r'$\alpha_0$', r'$\alpha_1$'])
    plt.legend()
    fighist1.savefig(rootpath+'/git/rdstackingproject/plotsmc/Diff_vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_histtop.pdf', format = 'pdf')

    #This plot is stacked
    fighist1 = plt.figure(figsize = (8, 6))
    sn01, sbins01, spatches01 = plt.hist([samplea0, samplea1], numbins * 10, rwidth = 1, stacked = True, alpha = 0.5, label = [r'$\alpha_0$', r'$\alpha_1$'])
    plt.legend()
    fighist1.savefig(rootpath+'/git/rdstackingproject/plotsmc/Diff_vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt_histstack.pdf', format = 'pdf')


    
#Now plot the NR data against the mcmc fit data, together with the 1-sigma varying error data
onesig_bounds = np.array([np.percentile(samples[:, i], [16, 84]) for i in range(len(samples[0]))]).T
modelfitpk = model_dv(pk)
figband = plt.figure(figsize = (12, 9))
#Plot the 1-sigma_percentile
for j in range(len(samples)):
    sample = samples[j]
    if np.all(onesig_bounds[0] <= sample) and np.all(sample <= onesig_bounds[1]):
        plt.plot(timesrd, model_dv(sample).real, "#79CAF2", alpha=0.3)
    
plt.plot(timesrd, gw_sxs_bbh_0305rd[:,1], "k", alpha=0.7, lw=2, label=r'NR_re')
plt.plot(timesrd, modelfitpk.real, "r", alpha=0.7, lw=2, label=r'FitMCmax_re')
plt.title(r'Comparison of the MC fit data and the $1-\sigma$ error band')
plt.legend()
plt.xlim(timesrd[0], timesrd[0]+80)
plt.xlabel("t")
plt.ylabel("h")

figband.savefig(rootpath+'/git/rdstackingproject/plotsmc/Diff_vary'+str(vary_fund)+'nmax='+str(nmax)+'_tshift='+str(tshift)+'_'+str(npoints)+'pt _band.png', format = 'png', dpi = 384, bbox_inches = 'tight')

#import os
#os.system('afplay /System/Library/Sounds/Submarine.aiff')