import itertools

import matplotlib.pyplot as plt
import matplotlib as mpl
import pymc3 as pm
import numpy as np
from scipy import stats
import tqdm
import pandas as pd
import seaborn as sns
sns.set_style("whitegrid")

%matplotlib inline

%load_ext version_information

%version_information pymc3, scipy, numpy, seaborn



"""
Explanation: Simple examples
from https://people.duke.edu/~ccc14/sta-663/PyMC3.html
End of explanation
"""


n = 100
h = 61
alpha = 2
beta = 2

niter = 10000
with pm.Model() as model: # context management
    # define priors
    p = pm.Beta('p', alpha=alpha, beta=beta)

    # define likelihood
    y = pm.Binomial('y', n=n, p=p, observed=h)

    # inference
    start = pm.find_MAP() # Use MAP estimate (optimization) as the initial state for MCMC
    #no need to assign manually
    #     step = pm.Metropolis() # Have a choice of samplers
    trace = pm.sample(niter, start=start, random_seed=123, progressbar=True, tune=1000, njobs=5)
    
with pm.Model() as model2: # context management
    # define priors, no information this time
    p = pm.Uniform('p', 0, 1)

    # define likelihood
    y = pm.Binomial('y', n=n, p=p, observed=h)

    # inference
    start = pm.find_MAP() # Use MAP estimate (optimization) as the initial state for MCMC
    #no need to assign manually
    #     step = pm.Metropolis() # Have a choice of samplers
    trace2 = pm.sample(niter, start=start, random_seed=123, progressbar=True, tune=1000, njobs=5)
    
    

pm.traceplot(trace, combined=True)
pm.traceplot(trace2, combined=True)

fig, ax = plt.subplots(ncols=2, figsize=(15,5))
sns.distplot(trace['p'], 15, label='post', ax=ax[0]);
x = np.linspace(0, 1, 100)
ax[0].plot(x, stats.beta.pdf(x, alpha, beta), label='prior');
ax[0].legend();

sns.distplot(trace2['p'], 15, label='post', ax=ax[1]);
x = np.linspace(0, 1, 100)
ax[1].plot(x, stats.uniform.pdf(x, ), label='prior');
ax[1].legend();


"""
Explanation: Coin toss
We’ll repeat the example of determining the bias of a coin from observed coin tosses. The likelihood is binomial, and we use a beta prior.
End of explanation
"""


# generate observed data
N = 100
_mu = np.array([10])
_sigma = np.array([2])
y = np.random.normal(_mu, _sigma, N)

niter = 10000
with pm.Model() as model:
    # define priors
    mu = pm.Uniform('mu', lower=0, upper=100, shape=_mu.shape)
    sigma = pm.Uniform('sigma', lower=0, upper=10, shape=_sigma.shape)

    # define likelihood
    y_obs = pm.Normal('Y_obs', mu=mu, sd=sigma, observed=y)

    # inference
    start = pm.find_MAP()
    trace = pm.sample(niter, start=start, random_seed=123, progressbar=True, njobs=5)

pm.traceplot(trace, combined=True)

pm.summary(trace)

fig, ax = plt.subplots(ncols=2, figsize=(15,5))
sns.distplot(trace['mu'], label='post', ax=ax[0]);
x = np.linspace(np.floor(trace['mu'].min()), np.ceil(trace['mu'].max()), 100)
ax[0].plot(x, stats.uniform.pdf(x, loc=x.min(), scale=np.ptp(x)), label='prior');
ax[0].legend();

sns.distplot(trace['sigma'], label='post', ax=ax[1]);
x = np.linspace(np.floor(trace['sigma'].min()), np.ceil(trace['sigma'].max()), 100)
ax[1].plot(x, stats.uniform.pdf(x, loc=x.min(), scale=np.ptp(x)), label='prior');
ax[1].legend();
np.ptp(trace['sigma'])

ppc = pm.sample_ppc(trace, samples=500, model=model, size=100)


fig, ax = plt.subplots(ncols=2, figsize=(15,5))
sns.distplot([n.mean() for n in ppc['Y_obs']], kde=True, ax=ax[0])
ax[0].axvline(y.mean())
ax[0].set(title='Posterior predictive of the mean', xlabel='mean(x)', ylabel='Frequency');

sns.distplot([n.std() for n in ppc['Y_obs']], kde=True, ax=ax[1])
ax[1].axvline(y.std())
ax[1].set(title='Posterior predictive of the std', xlabel='std(x)', ylabel='Frequency');

"""
Explanation: Estimating mean and standard deviation of normal distribution
End of explanation
"""


