from scipy import *
from numpy import *
from scipy import random
from scipy import stats
from scipy import linalg
import pylab

####################################################

def sampleMVN(mu, Sigma):
    Sigma_sqrt = linalg.cholesky(Sigma)
    return dot(Sigma_sqrt.T, stats.norm(0,1).rvs(len(mu))) + mu

####################################################

# Question 3a

def log_dens(x1, x2, nu):                    # log of the joint density of x1,x2
  a = (x1+x2)/(nu+2)
  return log(nu)+log(nu+1) - (nu+2) * log(exp(a)+exp(a-x1)+exp(a-x2))

nu = 1                                       # set parameter nu  

####################################################

t = linspace(-5, 5, 200)                     # define one side of the grid
mat = zeros((200,200))
for i in xrange(0,200):                      # evaluate the density over a grid defined by t
    for j in xrange(0,200):
        mat[i,j]=exp(log_dens(t[i],t[j],nu))

pylab.figure()                               # create the plot
pylab.contourf(t,t,mat)
pylab.title("Density for mu=1")              # give it a title
pylab.show()

####################################################

n = 1e4                                      # set the sample size
x = zeros((n,2))                             # create matrix to hold sample
cur_x = [0, 0]                               # set the initial value
cur_logf = log_dens(cur_x[0], cur_x[1], nu)  # evaluate density at cur_x
n_accepted = 0                               # counter of the accepted values
for i in xrange(n):                          # repeat n  times ...
    new_x = cur_x + sqrt(9) * stats.norm(0,1).rvs(2)
                                             # create proposed value
    new_logf = log_dens(new_x[0], new_x[1], nu)     
                                             # evaluate density at proposed value
    prob_acceptance = exp(new_logf-cur_logf) # probability of acceptance
    if random.random_sample(1)<prob_acceptance:
                                             # if we accept the proposed value ...
        n_accepted = n_accepted+1            # ... increment the counter of accepted values
        cur_x = new_x                        # ... accept the new value
        cur_logf = new_logf                  # ... retain the density at the accepted value
    x[i,:] = cur_x                           # store current value

####################################################

pylab.figure()
pylab.subplot(131)                           # split the screen into three
pylab.plot(x[:,0], x[:,1], '.')              # plot the sample
pylab.title("Joint distribution")            # add a title
pylab.subplot(132)
pylab.plot(x[:,0])                           # plot the sample path of x1
pylab.title("Sample path of X1")             # add a title
pylab.subplot(133)
pylab.plot(x[:,1])                           # plot the sample path of x2
pylab.title("Sample path of X2")             # add a title
pylab.show()

####################################################

mean(x[:,0])                                 # Monte Carlo estimate of E(X1)
mean(x[:,1])                                 # Monte Carlo estimate of E(X2)

####################################################

n_accepted/n                                 # proportion of accepted values

####################################################

cor_x1 = corrcoef(x[1:,0],x[:-1,0])[0,1]     # autocorrelation of x1
ess_x1 = n * (1-cor_x1) / (1+cor_x1)         # corresponding effective sample size
cor_x2 = corrcoef(x[1:,1],x[:-1,1])[0,1]     # autocorrelation of x2
ess_x2 = n * (1.-cor_x2) / (1.+cor_x2)       # corresponding effective sample size

####################################################

# Question 3b

def sample_cond_dist(x2, nu):                # function that samples from x1 given x2
    u = random.random_sample(1)              # generate uniform u
    v = u**(1./(nu+1))                       # compute intermediate result
    return float(-log((1.+exp(-x2))*(1-v)/v))# return inverse CDF at u

####################################################

x = zeros((n,2))                             # create matrix to hold sample
cur_x = [0, 0]                               # set the initial value
for i in xrange(n):                          # repeat n  times ...
    cur_x[0] = sample_cond_dist(cur_x[1],nu) # update x1 by sampling from x1 given x2         
    cur_x[1] = sample_cond_dist(cur_x[0],nu) # update x2 by sampling from x2 given x1 
    x[i,:] = cur_x                           # store current value

####################################################

# Question 3c

nu = 0.05                                    # set parameter nu  
t = linspace(-50, 10, 200)                   # define one side of the grid
mat = zeros((200,200))
for i in xrange(0,200):                      # evaluate the density over a grid defined by t
    for j in xrange(0,200):
        mat[i,j]=exp(log_dens(t[i],t[j],nu))

pylab.figure()                               # create the plot
pylab.contourf(t,t,mat)
pylab.title("Density for mu=0.05")           # give it a title
pylab.show()

####################################################

sigma_proposal = array([[1000, 998], [998, 1000]])    # set covariance of the proposal
x = zeros((n,2))                             # create matrix to hold sample
cur_x = [0, 0]                               # set the initial value
cur_logf = log_dens(cur_x[0], cur_x[1], nu)  # evaluate density at cur.x
n_accepted = 0                               # counter of the accepted values
for i in xrange(n):                          # repeat n  times ...
    new_x = cur_x + sampleMVN([0,0], sigma_proposal)
                                             # create proposed value with the above covariance
    new_logf = log_dens(new_x[0], new_x[1], nu)     
                                             # evaluate density at proposed value 
    prob_acceptance = exp(new_logf-cur_logf) # probability of acceptance
    if random.random_sample(1)<prob_acceptance:
                                             # if we accept the proposed value ...
        n_accepted = n_accepted+1            # ... increment the counter of accepted values
        cur_x = new_x                        # ... accept the new value
        cur_logf = new_logf                        # ... retain the density at the accepted value
    x[i,:] = cur_x                           # store current value

####################################################

# Question 4

x = load("/home/ioana/public_html/mixgamma.npy")

####################################################

def log_posterior(lam, x):                       # evaluates the log-posterior
    return sum(log(0.5*gammaPDF(x,4.,lam[0])+0.5*gammaPDF(x,5.,lam[1]))) + sum(log(gammaPDF(lam,4.,0.5)))         

####################################################

sd_proposal = 1                              # use large enough a standard deviation of the proposal
n = 1e4                                      # set the sample size
lam = zeros((n,2))                           # create matrix to hold sample
cur_lambda = [2., 8.]                        # set the initial value
cur_logf = log_posterior(cur_lambda, x) + sum(log(cur_lambda))  # evaluate log-posterior at cur_lambda
n_accepted = 0                               # counter of the accepted values
for i in xrange(n):                          # repeat n  times ...
    new_lambda = cur_lambda * exp(sd_proposal*stats.norm(0,1).rvs(2))
                                             # create proposed value using a multiplicative perturbation}
    new_logf = log_posterior(new_lambda, x) + sum(log(new_lambda))
                                             # evaluate log-posterior at proposed value
    prob_acceptance = exp(new_logf-cur_logf) # probability of acceptance
    if random.random_sample(1)<prob_acceptance:
                                             # if we accept the proposed value ...  
        n_accepted = n_accepted+1            # ... increment the counter of accepted values
        cur_lambda = new_lambda              # ... accept the new value
        cur_logf = new_logf                  # ... retain the density at the accepted value
    lam[i,:] = cur_lambda                    # store current value

####################################################

pylab.figure()
pylab.subplot(131)                           # split the screen into three
pylab.plot(lam[:,0], lam[:,1], '.')          # plot the sample
pylab.title("Joint distribution")            # add a title
pylab.subplot(132)
pylab.plot(lam[:,0])                         # plot the sample path of lambda1
pylab.title("Sample path of X1")             # add a title
pylab.subplot(133)
pylab.plot(lam[:,1])                         # plot the sample path of lambda2
pylab.title("Sample path of X2")             # add a title
pylab.show()

####################################################

# Question 5

def log_dens (x1, x2, lam=5):                # define the density as R function
    return -lam*(abs(x1**2+x2**2-1))

####################################################

lam = 5.
t = linspace(-1.5, 1.5, 200)                # define one side of the grid
mat = zeros((200,200))
for i in xrange(0,200):                      # evaluate the density over a grid defined by t
    for j in xrange(0,200):
        mat[i,j]=exp(log_dens(t[i], t[j], lam))

pylab.figure()                               # create the plot
pylab.contourf(t,t,mat)
pylab.title("Density")                       # give it a title
pylab.show()

####################################################

sd_proposal = 0.2                            # set standard deviation of the proposal
n = 1e4                                      # set the sample size
x = zeros((n,2))                             # create matrix to hold sample
cur_x = [0 , 0]                              # set the initial value
cur_logf = log_dens(cur_x[0],cur_x[1])       # evaluate log-density at cur.x
n_accepted = 0                               # counter of the accepted values
for i in xrange(n):                          # repeat n  times ...             
    new_x = cur_x + sd_proposal * stats.norm(0,1).rvs(2)
                                             # create proposed value
    new_logf = log_dens(new_x[0],new_x[1])   # evaluate log-density at proposed value}
    prob_acceptance = exp(new_logf-cur_logf) # probability of acceptance
    if random.random_sample(1)<prob_acceptance:
                                             # if we accept the proposed value ...  
        n_accepted = n_accepted+1            # ... increment the counter of accepted values
        cur_x = new_x                        # ... accept the new value
        cur_logf = new_logf                  # ... retain the density at the accepted value 
    x[i,:] = cur_x                           # store current value
print(n_accepted/n)                          # proportion of accepted values
print corrcoef(x[1:,0],x[:-1,0])[0,1])       # autocorrelation of x1
pylab.figure()
pylab.plot(x[:,0],x[:,1],'.')                # plot the sample
pylab.show()

####################################################

sd_proposal = 0.1                            # set standard deviation of the proposal
n = 1e4                                      # set the sample size
x = zeros((n,2))                             # create matrix to hold sample
cur_x = [0, 0]                               # set the initial value
cur_logf = log_dens(cur_x[0],cur_x[1])       # evaluate density at cur_x
n_accepted = 0                               # counter of the accepted values
for i in xrange(n):                          # repeat n  times ...              
    new_x = cur_x + sd_proposal * stats.norm(0,1).rvs(2)     
                                             # MH MOVE 1: create proposed value
    new_logf = log_dens(new_x[0],new_x[1])   # MH MOVE 1: evaluate density at proposed value
    prob_acceptance = exp(new_logf-cur_logf) # MH MOVE 1: probability of acceptance
    if random.random_sample()<prob_acceptance: 
                                             # MH MOVE 1: if we accept the proposed value ...
        n_accepted = n_accepted+1            # ... increment the counter of accepted values
        cur_x = new_x                        # ... accept the new value
        cur_logf = new_logf                  # ... retain the density at the accepted value
    angle = float(2*pi*random.random_sample(1))
                                             # MH MOVE 2: draw an angle at random
    rotation_matrix = array([[cos(angle), sin(angle)], [-sin(angle), cos(angle)]])
                                             # MH MOVE 2: Rotation matrix
    cur_x = dot(rotation_matrix,cur_x)       # MH MOVE 2: perform the rotation
    x[i,:] = cur_x                           # store current value
print(n_accepted/n)                          # proportion of accepted values
print corrcoef(x[1:,0],x[:-1,0])[0,1])       # autocorrelation of x1
pylab.figure()
pylab.plot(x[:,0],x[:,1],'.')                # plot the sample
pylab.show()

####################################################

n = 1e4                                      # set the sample size
x = zeros((n,2))                             # create matrix to hold sample
right_tail_prob = 1/(2-exp(-5))              # compute probability of r greater than 1 
for i in xrange(n):                          # sample r (n times)
    u = float(random.random_sample(1))       # inversion method: draw u from uniform
    if random.random_sample(1)<right_tail_prob:
        r = 1-log(u)/5                       # inversion method: case R > 1
    else:
        r = 1+log(exp(-5)+(1-exp(-5))*u)/5   # inversion method: case 0 < R < 1
    theta = float(2*pi*random.random_sample(1)) 
                                             # sample angle theta (n times)
    x[i,:] = [sqrt(r)*cos(theta),sqrt(r)*sin(theta)]
                                             # transform to x1 and x2
pylab.figure()
pylab.plot(x[:,0],x[:,1],'.')                # plot the sample
pylab.show()

####################################################

# Question 6c

leaf = array([[3.28, 3.09, 3.03, 3.03], 
              [3.52, 3.48, 3.38, 3.38],
              [2.88, 2.80, 2.81, 2.76],
              [3.34, 3.38, 3.24, 3.26]])
y_total = mean(leaf)                    # compute mean of all observations
y_bar = leaf.mean(1)                    # compute mean of rows in observation matrix
n = 1e5                                 # set sample size
alphas = zeros((n,4))                   # create matrix and vectors to hold samples
mu = zeros(n)
sigma2_alpha = zeros(n)
sigma2 = zeros(n)
cur_alpha = 0.01 * ones(4)              # set starting values
cur_sigma2 = 0.005
for i in xrange(n):                     # Gibbs sampler algorithm
    cur_mu = y_total - mean(cur_alpha) + sqrt(cur_sigma2/16) * stats.norm(0,1).rvs(1)
                                        # update mu
    sigma2_alpha_rate = 0.5 * sum(cur_alpha**2)
    cur_sigma2_alpha = 1/gammaSample(1, 1, sigma2_alpha_rate)
                                        # update sigma2.alpha
    sigma2_rate = 0.5 * sum( (leaf - cur_alpha - cur_mu)**2 )
    cur_sigma2 = 1/gammaSample(1, 7, sigma2_rate) 
                                        # update sigma2
    alpha_ratio = cur_sigma2_alpha / (4 * cur_sigma2_alpha + cur_sigma2)
    cur_alpha = 4*alpha_ratio*(y_bar-cur_mu) + sqrt(cur_sigma2 * alpha_ratio) * stats.norm(0,1).rvs(4)
                                        # update alphas
    alphas[i,:] = cur_alpha             # Store current values
    mu[i] = cur_mu
    sigma2[i] = cur_sigma2
    sigma2_alpha[i] = cur_sigma2_alpha

####################################################

s = range(60001, 100000, 2)

pylab.figure()
pylab.subplot(121)
pylab.plot(sigma2[s])
pylab.title("sigma^2")
pylab.subplot(122)
pylab.plot(sigma2_alpha[s])
pylab.title("sigma_alpha^2")

pylab.figure()
pylab.subplot(121)
pylab.hist(sigma2[s],normed=True)
pylab.title("sigma^2")
pylab.subplot(122)
pylab.hist(sigma2_alpha[s],normed=True)
pylab.title("sigma_alpha^2")

####################################################

s1 = range(60002, 100000, 2)
s2 = range(60000, 99997, 2)
print(corrcoef(sigma2[s1], sigma2[s2])[0,1])
print(corrcoef(sigma2_alpha[s1], sigma2_alpha[s2])[0,1])
s_all = range(60000,100000)
print(median(sigma2[s_all]))
print(median(sigma2_alpha[s_all]))