# -*- coding: utf-8 -*-
"""
Code to create all figures and results for Tutorial 8 -- Additional Practice

Tested with Python 3.7.4
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal, norm, poisson, uniform

# set seed for reproducibility
# (plots in the solution were created with a different seed though)
np.random.seed(1)

# -------------------------------------------------------------------
# Intro-text: multivariate Normal
# -------------------------------------------------------------------

plt.figure(figsize=(6.2, 5))
ax = 4   
x, y = np.mgrid[-ax:ax:.005, -ax:ax:.005]
pos = np.empty(x.shape + (2,))
pos[:, :, 0] = x; pos[:, :, 1] = y

rv = multivariate_normal(mean=[0,0])
density = rv.pdf(pos)
plt.contourf(x, y, density, cmap='viridis_r', antialiased=True)
plt.xlim((-ax, ax))
plt.ylim((-ax, ax))
plt.xlabel('x')
plt.ylabel('y')
cbar = plt.colorbar()

plt.savefig("figures/density.pdf", bbox_inches='tight', pad_inches=0)

# -------------------------------------------------------------------
# Exercise 1
# -------------------------------------------------------------------

def mh(p_star, param_init, num_samples=3000, vari=1., warmup=0):
    x = []
    x_current = param_init
    for n in range(num_samples+warmup):
        x_proposed = multivariate_normal.rvs(mean=x_current, cov=vari)
        a = multivariate_normal.pdf(x_current, mean=x_proposed, cov=vari) * p_star(x_proposed)
        a = a / (multivariate_normal.pdf(x_proposed, mean=x_current, cov=vari) * p_star(x_current))
        if a >= 1:
            x_next = np.copy(x_proposed)
        elif uniform.rvs(0, 1) < a:
            x_next = np.copy(x_proposed)
        else:
            x_next = np.copy(x_current)

        if n >= warmup:
            x.append(x_next)
        x_current = x_next

    return x


def target(x):
    return norm.pdf(x[0], loc=0, scale=1) * norm.pdf(x[1], loc=0, scale=1)


# starting from (0,0)
arr = np.array(mh(target, [0,0], num_samples=5000, vari=1., warmup=0))
x, y = arr[:,0], arr[:,1]

plt.figure(figsize=(5, 5))
plt.scatter(x, y, alpha=0.1)
plt.scatter(x[:20], y[:20], c='r',label="First 20")
plt.legend(loc='upper left')
plt.xlim((-ax, ax+3))
plt.ylim((-ax, ax+3))
plt.xlabel('x')
plt.ylabel('y')

plt.savefig("figures/1_00.pdf", bbox_inches='tight', pad_inches=0)

# starting from (7,7)
arr = np.array(mh(target, [7,7], num_samples=5000, vari=1., warmup=0))
x, y = arr[:,0], arr[:,1]

plt.figure(figsize=(5, 5))
plt.scatter(x, y, alpha=0.1)
plt.scatter(x[:20], y[:20], c='r',label="First 20")
plt.legend(loc='upper left')
plt.xlim((-ax, ax+3))
plt.ylim((-ax, ax+3))
plt.xlabel('x')
plt.ylabel('y')

plt.savefig("figures/1_77.pdf", bbox_inches='tight', pad_inches=0)

# -----------------------------------------------------------------
# Exercise 2
# -------------------------------------------------------------------

# Data

xx1 = np.array([-0.5051905265552105, -0.17185719322187715, 0.16147614011145617, 0.49480947344478954, 0.8150985069051909])
yy1 = np.array([1, 0, 2, 1, 2])

N1 = len(xx1)


def poisson_regression(params):
    a = params[0]
    b = params[1]
    # mean zero, standard deviation 10 == variance 100
    p = norm.pdf(a, loc=0, scale=10) * norm.pdf(b, loc=0, scale=10) 
    for n in range(N1):
        p = p * poisson.pmf(yy1[n], np.exp(a * xx1[n] + b)) 

    return p


# sample
S = 5000
samples = np.array(mh(poisson_regression, np.array([0, 0]), num_samples=S, vari=1, warmup=1000))

# plot
aa, bb = samples[:,0], samples[:,1]

print("Posterior mean for alpha is {0:5.2f}".format(np.mean(aa)))
print("Posterior mean for beta is {0:5.2f}".format(np.mean(bb)))
print("Correlation coefficient under the posterior is {0:5.2f}".format(np.corrcoef(aa,bb)[0,1]))

fig = plt.figure(figsize=(5,5))
plt.scatter(aa,bb,alpha=0.1)
plt.ylim([-3, 1.5])
plt.xlim([-2.5, 4.5])
plt.xlabel(r'$\alpha$')
plt.ylabel(r'$\beta$')

plt.savefig("figures/2_samples.pdf", bbox_inches='tight', pad_inches=0)


# -----------------------------------------------------------------
# Exercise 3
# -------------------------------------------------------------------

S = 5000

stepsi = 1
samples1 = np.array(mh(poisson_regression, np.array([0,0]), num_samples=S, vari=stepsi, warmup=1000))
samples2 = np.array(mh(poisson_regression, np.array([2,2]), num_samples=S, vari=stepsi, warmup=1000))
samples3 = np.array(mh(poisson_regression, np.array([3,-1]), num_samples=S, vari=stepsi, warmup=1000))

aa1, aa2, aa3 = samples1[:,0], samples2[:,0], samples3[:,0]
bb1, bb2, bb3 = samples1[:,1], samples2[:,1], samples3[:,1]

fig = plt.figure(figsize=(5,5))
plt.scatter(aa1, bb1, alpha=0.1)
plt.ylim([-4.5, 4])
plt.xlim([-2.5, 6])
plt.xlabel(r"$\alpha$")
plt.ylabel(r"$\beta$")

fig.savefig("figures/3_vari1.pdf", bbox_inches='tight')

# trace plot
fig = plt.figure(figsize=(10,4))
plt.xlim([0, S])
plt.xlabel("Iteration")
plt.ylabel("Sample")
plt.plot(np.arange(S), bb1, alpha=0.5, antialiased=True, label="Init @ (0,0)")
plt.plot(np.arange(S), bb2, alpha=0.5, antialiased=True, label="Init @ (2,2)")
plt.plot(np.arange(S), bb3, alpha=0.5, antialiased=True, label="Init @ (3,-1)")
plt.legend(loc='lower right')

fig.savefig("figures/3_traceplot_vari1.pdf", bbox_inches='tight')


# --------------


stepsi = 0.001
samples1 = np.array(mh(poisson_regression, np.array([0,0]), num_samples=S, vari=stepsi, warmup=1000))
samples2 = np.array(mh(poisson_regression, np.array([2,2]), num_samples=S, vari=stepsi, warmup=1000))
samples3 = np.array(mh(poisson_regression, np.array([3,-1]), num_samples=S, vari=stepsi, warmup=1000))

aa1, aa2, aa3 = samples1[:,0], samples2[:,0], samples3[:,0]
bb1, bb2, bb3 = samples1[:,1], samples2[:,1], samples3[:,1]

fig = plt.figure(figsize=(5,5))
plt.scatter(aa1, bb1, alpha=0.1)
plt.ylim([-4.5, 4])
plt.xlim([-2.5, 6])
plt.xlabel(r"$\alpha$")
plt.ylabel(r"$\beta$")

fig.savefig("figures/3_vari001.pdf", bbox_inches='tight')

# trace plot
fig = plt.figure(figsize=(10,4))
plt.xlim([0, S])
plt.xlabel("Iteration")
plt.ylabel("Sample")
plt.plot(np.arange(S), bb1, alpha=0.5, antialiased=True, label="Init @ (0,0)")
plt.plot(np.arange(S), bb2, alpha=0.5, antialiased=True, label="Init @ (2,2)")
plt.plot(np.arange(S), bb3, alpha=0.5, antialiased=True, label="Init @ (3,-1)")
plt.legend(loc='lower right')

fig.savefig("figures/3_traceplot_vari001.pdf", bbox_inches='tight')

# --------------

stepsi = 50
samples1 = np.array(mh(poisson_regression, np.array([0,0]), num_samples=S, vari=stepsi, warmup=1000))
samples2 = np.array(mh(poisson_regression, np.array([2,2]), num_samples=S, vari=stepsi, warmup=1000))
samples3 = np.array(mh(poisson_regression, np.array([3,-1]), num_samples=S, vari=stepsi, warmup=1000))

aa1, aa2, aa3 = samples1[:,0], samples2[:,0], samples3[:,0]
bb1, bb2, bb3 = samples1[:,1], samples2[:,1], samples3[:,1]

fig = plt.figure(figsize=(5,5))
plt.scatter(aa1, bb1, alpha=0.1)
plt.ylim([-4.5, 4])
plt.xlim([-2.5, 6])
plt.xlabel(r"$\alpha$")
plt.ylabel(r"$\beta$")

fig.savefig("figures/3_vari50.pdf", bbox_inches='tight')

# trace plot
fig = plt.figure(figsize=(10,4))
plt.xlim([0, S])
plt.xlabel("Iteration")
plt.ylabel("Sample")
plt.plot(np.arange(S), bb1, alpha=0.5, antialiased=True, label="Init @ (0,0)")
plt.plot(np.arange(S), bb2, alpha=0.5, antialiased=True, label="Init @ (2,2)")
plt.plot(np.arange(S), bb3, alpha=0.5, antialiased=True, label="Init @ (3,-1)")
plt.legend(loc='lower right')

fig.savefig("figures/3_traceplot_vari50.pdf", bbox_inches='tight')