#!/usr/bin/python3
# Animation for 1D MCMC

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation, lines

n_chain = 20000
n_burnin = 200
d = 0.5

def f(x):  # auxiliary function
    return (x - 3)**2 * np.sin(x) + x
    
def df(x): # its derivative
    return 2 * (x - 3) * np.sin(x) + (x - 3)**2 * np.cos(x) + 1

def P(x):  # some interesting non-normalized probability density function
    if x < 0:
        return 0.
    if x > 5:
        return f(5) * np.exp(df(5) / f(5) * (x-5))
    return f(x)
  
def newpoint(x):
    return x + d * np.random.choice([-1, 1])
    
def mcmcstep(x):
    xnew = newpoint(x)
    if P(xnew) > P(x):
        return xnew
    if np.random.random() < P(xnew) / P(x):
        return xnew
    return x
    
mcmcpoints = np.empty(n_chain)
burninpoints = np.empty(n_burnin)

x = 15.01 # choose some starting value
for i in range(n_burnin):
    burninpoints[i] = x
    x = mcmcstep(x)
for i in range(n_chain):
    mcmcpoints[i] = x
    x = mcmcstep(x)

X = np.linspace(-1, 15, 100)     
intP = 15.9  # this happens to be the normalization of P
for npoints in [200, 2000, 20000]:
    plt.figure()
    plt.plot(X, [ d * npoints / intP * P(x) for x in X], 'k')
    plt.hist(burninpoints, bins=np.arange(0,16,d) - d/2, alpha=0.5, label="200 points du 'burn-in'")
    plt.hist(mcmcpoints[:npoints], bins=np.arange(0,16,d)- d/2, alpha=0.5, label="N points de la chaîne")
    plt.title(f"N = {npoints}")
    plt.legend()
plt.show()



  
  

fig = plt.figure()
ax = plt.axes(xlim=(-1,16), ylim=(0, 30))

def init():
    ax.plot(X, [5* P(x) for x in X], 'k')

def animate(i):
    histo = ax.hist(burninpoints[:i], bins=np.arange(0,16,d), color='b')

anim = animation.FuncAnimation(fig, animate, init_func=init, frames=200, interval=10)
                               
plt.show()  


  

fig = plt.figure()
ax = plt.axes(xlim=(-1,16), ylim=(0, 30))

def init():
    ax.plot(X, [5* P(x) for x in X], 'k')

def animate(i):
    histo = ax.hist(mcmcpoints[:i], bins=np.arange(0,16,d), color='r')

anim = animation.FuncAnimation(fig, animate, init_func=init, frames=200, interval=10)
                               
plt.show()  
  
    
'''    
def newpoint(x, y):
    phi = np.random.random() * 2 * np.pi
    return x + d*np.cos(phi), y + d*np.sin(phi)
    
def mcmcstep(x, y):
    xnew, ynew = newpoint(x, y)
    if np.random.random() > P(x, y) / P(xnew, ynew):
        return xnew, ynew
    else:
        return x, y
    
x = np.random.random() * 10 - 5
y = np.random.random() * 10 - 5

for i in range(n_burnin):
    x, y = mcmcstep(x, y)
    
chain = np.empty([2, n_chain])
for i in range(n_chain):
    x, y = mcmcstep(x, y)
    chain[:, i] = [x, y]
    
plt.scatter(chain[0, :], chain[1, :])
plt.show()
'''