import numpy as np
import matplotlib.pyplot as plt
import pymc3 as pm
import seaborn as sns
import time
#print('Running on PyMC3 v{}'.format(pm.__version__))
plt.style.use('seaborn-darkgrid')

# Initialize random number generator
np.random.seed(123)

# True parameter values
alpha, sigma = 1, 1
beta = [1, 2.5]

# Size of dataset
size = 100

# Predictor variable
X1 = np.random.randn(size)
X2 = np.random.randn(size) * 0.2

# Simulate outcome variable
Y = alpha + beta[0]*X1 + beta[1]*X2 + np.random.randn(size)*sigma

fig, axes = plt.subplots(1, 2, sharex=True, figsize=(10,4))
axes[0].scatter(X1, Y)
axes[1].scatter(X2, Y)
axes[0].set_ylabel('Y'); axes[0].set_xlabel('X1'); axes[1].set_xlabel('X2');
plt.show()

basic_model = pm.Model()
with basic_model:

    # Priors for unknown model parameters
    alpha = pm.Normal('alpha', mu=0, sd=10)
    beta = pm.Normal('beta', mu=0, sd=10, shape=2)
    sigma = pm.HalfNormal('sigma', sd=1)
    #sigma = pm.HalfNormal('sigma', sd=2)

    # Expected value of outcome
    mu = alpha + beta[0]*X1 + beta[1]*X2

    # Likelihood (sampling distribution) of observations
    Y_obs = pm.Normal('Y_obs', mu=mu, sd=sigma, observed=Y)

map_estimate = pm.find_MAP(model=basic_model, method='powell')

#print "Map estimate = " , map_estimate
print("Map estimate = " , map_estimate)


time1 = time.clock()

with basic_model:
    # draw 1000 posterior samples
    trace = pm.sample(1000, njobs=1)
    #mean_field = pm.fit(method='advi')

time2 = time.clock()

#print 'sample time = ', time2 - time1, ' seconds'
print('sample time = ', time2 - time1, ' seconds')

pm.traceplot(trace);
#pm.plot_posterior(mean_field.sample(1000), color='LightSeaGreen');
plt.show()