#
# Model
#
def rvmodel(x,ampl,offset,phase):
    return offset+ampl*np.cos(2*np.pi*x+phase)

#
# The Bayesian formula:  P_posterior(model|data) =
#         P_prior(model) * P(data|model) / someconstant
#
def lnprob(theta, x, y, yerr):
    lp = lnprior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + lnlike(theta, x, y, yerr)

#
# Define the prior
#
bounds = [(10.,100.),(-50.,50.),(0.,2*np.pi)]
def lnprior(theta):
    global bounds
    ampl, offset, phase = theta
    if bounds[0][0] < ampl < bounds[0][1] and \
       bounds[1][0] < offset < bounds[1][1] and \
       bounds[2][0] < phase < bounds[2][1]:
        return 0.0
    return -np.inf

#
# Likelihood function
#
def lnlike(theta, x, y, yerr):
    ampl, offset, phase = theta
    model = rvmodel(x,ampl,offset,phase)
    sigma2 = yerr**2
    return -0.5*(np.sum((y-model)**2*/sigma2 + np.log(2*np.pi*sigma2)))

#
# Initialize the walkers of MCMC 
#
ndim, nwalkers = 3, 100
pos = []
for i in range(nwalkers):
    w = []
    for k in range(ndim):
        w.append(bounds[k][0]+(bounds[k][1]-bounds[k][0])*np.random.rand())
    pos.append(np.array(w))

#
# Set up the Emcee 
#
import emcee
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(x, y, yerr))

#
# Now run the MCMC chain
#
nchain = 500
sampler.run_mcmc(pos, nchain)

#
# See how the population evolves the first 10 steps
#
figevol = plt.figure()
ism = 50
plt.plot(sampler.chain[:,ism,0],sampler.chain[:,ism,1],'.')
for ism in range(0,50,10):
    plt.plot(sampler.chain[:,ism,0],sampler.chain[:,ism,1],'.')
plt.xlabel('ampl')
plt.ylabel('offset')

#
# Triangle plot, discarding first 50 samples
#
import corner
samples = sampler.chain[:, 50:, :].reshape((-1, ndim))
fig3 = corner.corner(samples, labels=["ampl", "offset", "phase"])

plt.show()