import numpy as np

true_mean = 2
true_sigma = 1.5
true_tau = 1 /  true_sigma**2
print(true_tau)
N = 15

data = np.random.normal(true_mean, true_sigma, N)

print(data)

%matplotlib inline
import matplotlib.pyplot as plt
plt.rc('font', family='serif') 
plt.rc('font', serif='Times New Roman')

plt.hist(data);

import scipy.stats as sp

def calc_posterior(mu, tau, x=data):
    '''
    This is basically the model. This function calculates the log 
    posterior of a given set of parameters.
    '''
    # prior for mu
    logp = sp.distributions.norm.logpdf(mu, 0, 1000)
    # prior for tau
    logp += sp.distributions.gamma.logpdf(tau, 0.01, 0.01)
    # calculate sigma from tau, needed for likelihood
    sigma = (1.0/tau)**0.5
    # data likeligood
    logp += sum(sp.distributions.norm.logpdf(x, mu, sigma))
    return logp

print(calc_posterior(mu=-2, tau=1./1**2, x=data))
print(calc_posterior(mu=-1, tau=1./0.5**2, x=data))
print(calc_posterior(mu=1, tau=1./0.1**2)) # don't HAVE to supply data as it's a KWARG
print(calc_posterior(mu=2, tau=1./2**2, x=data))

rnorm = np.random.normal
runif = np.random.rand

def metropolis(n_iterations, initial_values, prop_var=1):

    n_params = len(initial_values)
            
    # Initial proposal standard deviations
    prop_sd = [prop_var]*n_params
    
    # Initialize trace for parameters
    trace = np.empty((n_iterations+1, n_params))
    
    # Set initial values
    trace[0] = initial_values
        
    # Calculate joint posterior for initial values
    current_log_prob = calc_posterior(*trace[0], x=data)
    
    # Initialize acceptance counts
    accepted = [0]*n_params
    
    for i in range(n_iterations):
    
        if not i%1000: print 'Iteration', i
    
        # Grab current parameter values
        current_params = trace[i]
    
        for j in range(n_params):
    
            # Get current value for parameter j
            p = trace[i].copy()
    
            # Propose new value
            if j==2:
                # Ensure tau is positive
                theta = np.exp(rnorm(np.log(current_params[j]), prop_sd[j]))
            else:
                theta = rnorm(current_params[j], prop_sd[j])
            
            # Insert new value 
            p[j] = theta
    
            # Calculate log posterior with proposed value
            proposed_log_prob = calc_posterior(*p, x=data)
    
            # Log-acceptance rate
            alpha = proposed_log_prob - current_log_prob
    
            # Sample a uniform random variate
            u = runif()
    
            # Test proposed value
            if np.log(u) < alpha:
                # Accept
                trace[i+1,j] = theta
                current_log_prob = proposed_log_prob
                accepted[j] += 1
            else:
                # Reject
                trace[i+1,j] = trace[i,j]
                
    return trace, accepted

n_iter = 5000
initial_params = (-3,0.2)
trace, acc = metropolis(n_iter, initial_params, prop_var=0.1)



plt.plot(trace)
plt.axhline(y=true_tau)
plt.axhline(y=true_mean)

plt.figure(figsize=(8,8))
plt.plot(trace[:,0], trace[:,1])
plt.xlabel('$\mu$')
plt.ylabel('tau')

plt.figure(1, figsize=(12,6))

x=np.linspace(-15,15,500)
for i in range(3000):
    mu, tau = trace[i,:]
    sigma = (1.0/tau)**0.5
    plt.plot(x, sp.norm.pdf(x,mu,sigma), alpha=0.01, color='r')

plt.plot(data, np.zeros(data.shape)
         ,'+'
         , c='k'
         , markersize=20)
#plt.savefig('MH_example.png', bbox_inches='tight')



trace.shape

trace[:,0]

mu, tau = trace[1,:]

mu

tau





muv, tauv = np.meshgrid(np.linspace(-20, 20, 4), np.linspace(-20, 20, 4))





def run_chains(muv, tauv, n_iter):
    # number of chains decided by number of initial parameters provided
    n_chains = muv.size
    fig = plt.figure(figsize=(8,8))
    for chain in range(n_chains):
        print 'chain: ' , chain
        intial_params=(muv.flat[chain], tauv.flat[chain])
        print intial_params
        trace, acc = metropolis(n_iter, initial_params)
        
        plt.plot(trace[:,0], trace[:,1], alpha=0.5)
        
    plt.xlabel('$\mu$')
    plt.ylabel('tau')
    
    plt.axhline(y=true_tau, linewidth=3, c='r', ls='--')
    plt.axvline(x=true_mean, linewidth=3, c='r', ls='--')

run_chains(muv, tauv, n_iter=500)