# https://github.com/probml/pyprobml/blob/master/scripts/beta_binom_approx_post_pymc3.py

# 1d approixmation to beta binomial model
# https://github.com/aloctavodia/BAP


# import superimport

import pymc3 as pm
import numpy as np
import seaborn as sns
import scipy.stats as stats
import matplotlib.pyplot as plt
import arviz as az
import math

# import pyprobml_utils as pml

# data = np.repeat([0, 1], (10, 3))
data = np.repeat([0, 1], (10, 1))
h = data.sum()
t = len(data) - h

"""
Explanation: <a href="https://colab.research.google.com/github/probml/probml-notebooks/blob/main/notebooks/beta_binom_approx_post_pymc.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Posterior for Beta-binomial distribution
We compare various approximations to the exact 1d distribution.
End of explanation
"""


plt.figure()
x = np.linspace(0, 1, 100)
xs = x  # grid
dx_exact = xs[1] - xs[0]
post_exact = stats.beta.pdf(xs, h + 1, t + 1)
post_exact = post_exact / np.sum(post_exact)
plt.plot(xs, post_exact)
plt.yticks([])
plt.title("exact posterior")
plt.savefig("bb_exact.pdf")

"""
Explanation: Exact
End of explanation
"""


# Grid
def posterior_grid(heads, tails, grid_points=100):
    grid = np.linspace(0, 1, grid_points)
    prior = np.repeat(1 / grid_points, grid_points)  # uniform prior
    likelihood = stats.binom.pmf(heads, heads + tails, grid)
    posterior = likelihood * prior
    posterior /= posterior.sum()
    # posterior = posterior * grid_points
    return grid, posterior


n = 20
grid, posterior = posterior_grid(h, t, n)
dx_grid = grid[1] - grid[0]
sf = dx_grid / dx_exact  # Jacobian scale factor
plt.figure()
# plt.stem(grid, posterior, use_line_collection=True)
plt.bar(grid, posterior, width=1 / n, alpha=0.2)
plt.plot(xs, post_exact * sf)
plt.title("grid approximation")
plt.yticks([])
plt.xlabel("θ")
plt.savefig("bb_grid.pdf")

"""
Explanation: Grid
End of explanation
"""


# Laplace
with pm.Model() as normal_aproximation:
    theta = pm.Beta("theta", 1.0, 1.0)
    y = pm.Binomial("y", n=1, p=theta, observed=data)  # Bernoulli
    mean_q = pm.find_MAP()
    std_q = ((1 / pm.find_hessian(mean_q, vars=[theta])) ** 0.5)[0]
    mu = mean_q["theta"]

print([mu, std_q])

plt.figure()
plt.plot(xs, stats.norm.pdf(xs, mu, std_q), "--", label="Laplace")
post_exact = stats.beta.pdf(xs, h + 1, t + 1)
plt.plot(xs, post_exact, label="exact")
plt.title("Quadratic approximation")
plt.xlabel("θ", fontsize=14)
plt.yticks([])
plt.legend()
plt.savefig("bb_laplace.pdf");

"""
Explanation: Laplace
End of explanation
"""


# ADVI
with pm.Model() as mf_model:
    theta = pm.Beta("theta", 1.0, 1.0)
    y = pm.Binomial("y", n=1, p=theta, observed=data)  # Bernoulli
    mean_field = pm.fit(method="advi")
    trace_mf = mean_field.sample(1000)
thetas = trace_mf["theta"]
axes = az.plot_posterior(thetas, hdi_prob=0.95)
plt.savefig("bb_mf.pdf")

plt.show()

# track mean and std
with pm.Model() as mf_model:
    theta = pm.Beta("theta", 1.0, 1.0)
    y = pm.Binomial("y", n=1, p=theta, observed=data)  # Bernoulli
    advi = pm.ADVI()
    tracker = pm.callbacks.Tracker(
        mean=advi.approx.mean.eval, std=advi.approx.std.eval  # callable that returns mean  # callable that returns std
    )
    approx = advi.fit(callbacks=[tracker])

trace_approx = approx.sample(1000)
thetas = trace_approx["theta"]

axes = az.plot_posterior(thetas, hdi_prob=0.95)
plt.savefig("bb_mf_kde_az.pdf")

plt.figure()
plt.plot(tracker["mean"])
plt.title("Mean")
plt.savefig("bb_mf_mean.pdf")

plt.figure()
plt.plot(tracker["std"])
plt.title("Std ")
plt.savefig("bb_mf_std.pdf")

plt.figure()
plt.plot(advi.hist)
plt.title("Negative ELBO")
plt.savefig("bb_mf_elbo.pdf")

plt.figure()
sns.kdeplot(thetas)
plt.title("KDE of posterior samples")
plt.savefig("bb_mf_kde.pdf")

plt.show()

"""
Explanation: ADVI
End of explanation
"""


# HMC
with pm.Model() as hmc_model:
    theta = pm.Beta("theta", 1.0, 1.0)
    y = pm.Binomial("y", n=1, p=theta, observed=data)  # Bernoulli
    trace = pm.sample(1000, random_seed=42, cores=1, chains=2)
thetas = trace["theta"]
axes = az.plot_posterior(thetas, hdi_prob=0.95)
plt.savefig("bb_hmc.pdf")

az.plot_trace(trace)
plt.savefig("bb_hmc_trace.pdf", dpi=300)

"""
Explanation: HMC
End of explanation
"""


