from scipy import *
from scipy import stats
import pylab
import sys
sys.path.append("/home/ludger/lib/python2.6/site-packages/")
import statistics

def gammaPDF(x, alpha, beta):
    return stats.gamma(alpha).pdf(x*beta)*beta

def gammaSample(n, alpha, beta):
    return stats.gamma(alpha).rvs(n)/beta

def kde(x, t, weights):
    weights = weights / sum(weights) * len(weights)
    return statistics.pdf(x, t, weights)

# Task 1

y = array([3, 6, 3, 5, 9, 14, 12, 11, 19, 18, 
           15, 4, 1, 6, 11, 21, 11, 3, 7, 18])

# Task 2

alpha = 0.1
beta = 0.1

alpha_post_1 = alpha + sum(y[0:5])
alpha_post_2 = alpha + sum(y[5:10])
beta_post = beta + 5
print(alpha_post_1)
print(alpha_post_2)
print(beta_post)

# Task 3

post_mean_1_labelled = alpha_post_1 / beta_post
post_mean_2_labelled = alpha_post_2 / beta_post
print(post_mean_1_labelled)
print(post_mean_2_labelled)

# Task 4

n = 10000
lambda_1 = gammaSample(n, alpha_post_1, beta_post)
lambda_2 = gammaSample(n, alpha_post_2, beta_post)
weights = zeros(n)
for i in xrange(n):
    weights[i] = prod(0.5 * stats.poisson(lambda_1[i]).pmf(y[10:20]) + 0.5 * stats.poisson(lambda_2[i]).pmf(y[10:20]))

# Task 5

post_mean_1_all = sum(lambda_1*weights) / sum(weights)
post_mean_2_all = sum(lambda_2*weights) / sum(weights)
print(post_mean_1_all)
print(post_mean_2_all)

# Task 6

t = linspace(0,20,1000)
lambda_1_density = kde(lambda_1, t, weights)
pylab.figure(1)
pylab.plot(t, lambda_1_density,'r-')
pylab.plot(t, gammaPDF(t, alpha_post_1, beta_post), 'g:')
pylab.show()

lambda_2_density = kde(lambda_2, t, weights)
pylab.figure(2)
pylab.plot(t, lambda_2_density,'r-')
pylab.plot(t, gammaPDF(t, alpha_post_2, beta_post), 'g:')
pylab.show()

# Task 7

allocations = zeros((n,10))
for i in xrange(n):
    allocations[i,:] = stats.poisson(lambda_1[i]).pmf(y[10:20]) / (stats.poisson(lambda_1[i]).pmf(y[10:20]) + stats.poisson(lambda_2[i]).pmf(y[10:20]))

prob_of_group_1 = zeros(10)
for i in range(10):
   prob_of_group_1[i] = sum(allocations[:,i]*weights) / sum(weights)

print prob_of_group_1