import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde
from scipy.ndimage.filters import gaussian_filter

# Load posterior samples
source = "./Figure_2_posterior.dat"
OmgT,aT = np.loadtxt(source,unpack=True,usecols=(0,1),skiprows=2)

# Save copies of input data
OmgT_original = np.copy(OmgT)
aT_original = np.copy(aT)

# Mirror data in the x-direction (if specified)
# When obtaining 68% and 95% contours below, the KDE interpolator cannot correctly
# interpret a hard boundary in our data. For instance, our Omega0 posterior has a hard cut-off
# at log-Omega0 = -13. The KDE interpolator sees that our posterior is zero for smaller amplitudes,
# and will therefore try to smoothly interpolate between zero and whatever our finite posterior
# value is at log-Omega0 = -13. By mirroring the data across log-Omega0 = -13, we impose a symmetry
# condition on the KDE results, preventing the above anomalous behavior.
mirror_OmgT = -13.
mirror_aT = -8.

# Instantiate arrays and loop across data
tmp_OmgT = np.zeros(len(OmgT))
tmp_aT = np.zeros(len(aT))
for i,val in enumerate(OmgT):

    # Get height above lower bound
    height = val-(mirror_OmgT)

    # Save a value reflected across lower bound
    tmp_OmgT[i] = mirror_OmgT-height
    tmp_aT[i] = aT[i]

tmp2_OmgT = np.zeros(len(OmgT))
tmp2_aT = np.zeros(len(aT))
for i,val in enumerate(aT):

    # Get height above lower bound
    height = val-(mirror_aT)

    # Save a value reflected about lower bound
    tmp2_OmgT[i] = OmgT[i]
    tmp2_aT[i] = mirror_aT-height

# Join data
OmgT = np.concatenate((OmgT,tmp_OmgT,tmp2_OmgT))
aT = np.concatenate((aT,tmp_aT,tmp2_aT))

# Set up KDE
n = len(OmgT_original)
kdeWidth = n**(-1./(2.+4.))*(0.2) # Scott's rule (decreased by a factor of 2/3 to improve resolution)
TensorKDE = gaussian_kde([aT,OmgT],bw_method=kdeWidth)

# Define grid over which to evaluate KDE
logOmgMin=-13.
logOmgMax=-5.
aMin=-8.
aMax=8.
OmegaTs = np.arange(logOmgMin,logOmgMax,(logOmgMax-logOmgMin)/300.)
aTs = np.arange(aMin,aMax,(aMax-aMin)/300.)
x_grid,y_grid = np.meshgrid(OmegaTs,aTs)

# Evaluate KDE over grid
dx = len(OmegaTs)
dy = len(aTs)
kdeVals = np.zeros((dx,dy))
for i in range(dx):
    for j in range(dy):
        kdeVals[i,j] = TensorKDE([aTs[j],OmegaTs[i]])

# Sum up pdf values, from loudest pixel to quietest pixel
sortedVals = np.sort(kdeVals.flatten())[::-1]
cdfVals = np.cumsum(sortedVals)/np.sum(sortedVals)

# Find where cdf crosses 0.68 and 0.95 to get central 68% and 95% regions
i68 = np.argmin(np.abs(cdfVals - 0.68))
i95 = np.argmin(np.abs(cdfVals - 0.95))
val68 = sortedVals[i68]
val95 = sortedVals[i95]

# Set up figure and plot contours
fig,ax = plt.subplots()
kdeValsSmoothed = gaussian_filter(kdeVals,7.)
CS = ax.contour(x_grid.T,y_grid.T,kdeValsSmoothed,levels=(val95,val68))

# Extract coordinates from 95% credible contour
xData95 = np.array([])
yData95 = np.array([])
for seg in CS.allsegs[0]:
    xData95 = np.append(xData95,seg[:,0])
    yData95 = np.append(yData95,seg[:,1])

# Extract coordinates from 68% credible contour
xData68 = np.array([])
yData68 = np.array([])
for seg in CS.allsegs[1]:
    xData68 = np.append(xData68,seg[:,0])
    yData68 = np.append(yData68,seg[:,1])

# Save contour data
np.savetxt('contourData_68.dat',np.transpose([xData68,yData68]),delimiter='\t')
np.savetxt('contourData_95.dat',np.transpose([xData95,yData95]),delimiter='\t')