import numpy
import matplotlib.pyplot as plt
from fgivenx import plot_contours, plot_lines, plot_dkl


# Model definitions
# =================
# Define a simple straight line function, parameters theta=(m,c)
def f(x, theta):
    m, c = theta
    return m * x + c


numpy.random.seed(1)

# Posterior samples
nsamples = 1000
ms = numpy.random.normal(loc=-5, scale=1, size=nsamples)
cs = numpy.random.normal(loc=2, scale=1, size=nsamples)
samples = numpy.array([(m, c) for m, c in zip(ms, cs)]).copy()

# Prior samples
ms = numpy.random.normal(loc=0, scale=5, size=nsamples)
cs = numpy.random.normal(loc=0, scale=5, size=nsamples)
prior_samples = numpy.array([(m, c) for m, c in zip(ms, cs)]).copy()

# Set the x range to plot on
xmin, xmax = -2, 2
nx = 100
x = numpy.linspace(xmin, xmax, nx)

# Set the cache
cache = 'cache/test'
prior_cache = cache + '_prior'

# Plotting
# ========
fig, axes = plt.subplots(2, 2)

# Sample plot
# -----------
ax_samples = axes[0, 0]
ax_samples.set_ylabel(r'$c$')
ax_samples.set_xlabel(r'$m$')
ax_samples.plot(prior_samples.T[0], prior_samples.T[1], 'b.')
ax_samples.plot(samples.T[0], samples.T[1], 'r.')

# Line plot
# ---------
ax_lines = axes[0, 1]
ax_lines.set_ylabel(r'$y = m x + c$')
ax_lines.set_xlabel(r'$x$')
plot_lines(f, x, prior_samples, ax_lines, color='b', cache=prior_cache)
plot_lines(f, x, samples, ax_lines, color='r', cache=cache)

# Predictive posterior plot
# -------------------------
ax_fgivenx = axes[1, 1]
ax_fgivenx.set_ylabel(r'$P(y|x)$')
ax_fgivenx.set_xlabel(r'$x$')
cbar = plot_contours(f, x, prior_samples, ax_fgivenx,
                     colors=plt.cm.Blues_r, lines=False,
                     cache=prior_cache)
cbar = plot_contours(f, x, samples, ax_fgivenx, cache=cache)

# DKL plot
# --------
ax_dkl = axes[1, 0]
ax_dkl.set_ylabel(r'$D_\mathrm{KL}$')
ax_dkl.set_xlabel(r'$x$')
ax_dkl.set_ylim(bottom=0, top=2.0)
plot_dkl(f, x, samples, prior_samples, ax_dkl,
         cache=cache, prior_cache=prior_cache)

ax_lines.get_shared_x_axes().join(ax_lines, ax_fgivenx, ax_samples)

fig.tight_layout()
fig.savefig('plot.png')

import numpy
import matplotlib.pyplot as plt
from fgivenx import plot_contours, samples_from_getdist_chains

file_root = './plik_HM_TT_lowl/base_plikHM_TT_lowl'
samples, weights = samples_from_getdist_chains(['logA', 'ns'], file_root)

def PPS(k, theta):
    logA, ns = theta
    return logA + (ns - 1) * numpy.log(k)

k = numpy.logspace(-4,1,100)
cbar = plot_contours(PPS, k, samples, weights=weights)
cbar = plt.colorbar(cbar,ticks=[0,1,2,3])
cbar.set_ticklabels(['',r'$1\sigma$',r'$2\sigma$',r'$3\sigma$'])

plt.xscale('log')
plt.ylim(2,4)
plt.ylabel(r'$\ln\left(10^{10}\mathcal{P}_\mathcal{R}\right)$')
plt.xlabel(r'$k / {\rm Mpc}^{-1}$')
plt.tight_layout()
plt.savefig('planck.png')