import numpy as np
import seaborn
import matplotlib.pyplot as plt
from math import sqrt
from scipy.stats import multivariate_normal

seaborn.set_style("white")

x = np.array([-1.87, -1.76, -1.67, -1.22, -0.07, 0.11, 0.67, 1.60, 2.22, 2.51])
y = np.array([0.06, 1.67, 0.54, -1.45, -0.18, -0.67, 0.92, 2.95, 5.13, 5.18])

D = 20
eps = 1
base = np.linspace(-2, 3, num=D)


def map_features(x, base=base):
    x = x.reshape(len(x), 1)
    phi = np.exp(-(x - base) ** 2 / (2 * eps ** 2))
    phi = np.concatenate((phi, np.ones((phi.shape[0], 1))), 1)
    return phi


def draw_prior(x):
    phi = map_features(x)
    D = phi.shape[1]
    w = np.random.normal(0, 1, D)
    return np.dot(phi, w)


def posterior():
    phi = map_features(x)
    D = phi.shape[1]
    invsigma = np.identity(D) + np.dot(phi.T, phi)
    sigma = np.linalg.inv(invsigma)
    mu = np.dot(sigma, np.dot(phi.T, y))
    return mu, sigma


def plot_posteriors():
    plt.figure(figsize=(9, 6))
    mu, sigma = posterior()
    px = np.linspace(-2, 3, num=200)
    phi = map_features(px)
    plt.plot(x, y, "o", label="data")
    for i in xrange(5):
        w = np.random.multivariate_normal(mu, sigma)
        plt.plot(px, np.dot(phi, w), label="post. " + str(i))
    plt.legend(loc="lower right")
    plt.savefig("posterior.pdf")


def plot_tube():
    fig = plt.figure(figsize=(9, 6))
    mu, sigma = posterior()
    px = np.linspace(-2, 3, num=200)
    phi = map_features(px)
    plt.plot(px, np.dot(phi, mu), label="post. pred. mean")
    ax = fig.get_axes()[0]
    inf = [np.dot(row, mu) - 1.96 * sqrt(1 + np.dot(row, np.dot(sigma, row)))
           for row in phi]
    sup = [np.dot(row, mu) + 1.96 * sqrt(1 + np.dot(row, np.dot(sigma, row)))
           for row in phi]
    ax.fill_between(px, inf, sup, alpha=0.2)
    plt.plot(x, y, "o", label="data")
    plt.legend(loc="lower right")
    plt.savefig("predictive.pdf")


def plot_priors():
    plt.figure(figsize=(9, 6))
    px = np.linspace(-2, 3, num=200)
    for i in xrange(5):
        plt.plot(px, draw_prior(px), label="prior " + str(i))
    plt.legend(loc="lower right")
    plt.savefig("prior.pdf")


def marginal_likelihood():
    eps = 1
    ml = []
    px = np.linspace(1, 50, num=10)
    for D in px:
        base = np.linspace(-2, 3, num=D)
        phi = map_features(x, base)
        z = np.zeros(phi.shape[0])
        rv = multivariate_normal(mean=z, cov=np.identity(phi.shape[0])
                                 + np.dot(phi, phi.T))
        ml.append(rv.pdf(y))
    plt.figure(figsize=(9, 6))
    plt.bar(px, ml)
    plt.savefig("ml.pdf")



if __name__ == "__main__":
    plot_priors()
    plot_posteriors()
    plot_tube()
    marginal_likelihood()
#plt.plot(x, y, "o")
#plt.show()