import random, math
from math import *
import numpy as np
import pylab as p
import matplotlib.pyplot as plt
from matplotlib.ticker import LinearLocator, FormatStrFormatter

%matplotlib inline

def gibbs(N=5000, thin=1000, outequilibrium=500):
    
    x0 = 0.
    y0 = 0.
        
    for i in range(outequilibrium):
            x0 = random.gammavariate(3, 1.0/( y0*y0 + 4 ) )
            y0 = random.gauss(1.0/( x0 + 1 ), 1.0/math.sqrt( 2*x0 + 2 ))

    x = np.array([x0])
    y = np.array([y0])
    
    xaux = x0
    yaux = y0
    
    for i in range(N):
        for j in range(thin):
            xaux = random.gammavariate(3, 1.0/( yaux*yaux + 4 ) )
            yaux = random.gauss(1.0/( xaux + 1 ), 1.0/math.sqrt( 2*xaux + 2 ))
        
        # Storing uncorrelated data
        x = np.append( x, xaux )
        y = np.append( y, yaux )
            
    return x,y

x,y = gibbs(N=1000,thin=10)

m = np.mean(x)
print '<X> = %f' %(m)

m = np.mean(y)
print '<Y> = %f' %(m)

n = x.shape[0]
m = 0
for j in range(n):
    m += x[j]*math.cos(y[j])
m /= float(n)
print '< X cos(Y) > = %f' %(m)

Nmedias = 100
Nj = 100

avgX = np.array([])
avgY = np.array([])
x = np.array([])
y = np.array([])

for j in range(Nmedias):
    
    xn,yn = gibbs(N = Nj, thin = 50)
        
    x = np.append( x, xn )
    avgX = np.append( avgX, np.mean(x) )
    y = np.append( y, yn )
    avgY = np.append( avgY, np.mean(y) )
    

plt.plot(avgX, '-r')
plt.plot(avgY, '-b')

Nbins = 30
H, xedges, yedges = np.histogram2d(x, y, bins=( Nbins, Nbins) )
extent = [yedges[0], yedges[-1], xedges[-1], xedges[0]]

X, Y = np.meshgrid(xedges[0:Nbins], yedges[0:Nbins])

fig = plt.figure()
ax = fig.gca(projection='3d')
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(X, Y, H, rstride=1, cstride=1, cmap=cm.coolwarm,
        linewidth=0, antialiased=False)

#ax.set_zlim(-1.01, 1.01)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

fig.colorbar(surf, shrink=0.5, aspect=5)

plt.xlabel("x")
plt.ylabel("y")
plt.show()


%reset

import numpy as np
import matplotlib.pyplot as plt
import pylab

%matplotlib inline

print 'Reading the data...'
s = np.loadtxt('data/dados_nonmixture.data')


print 'We have %d observations' %(s.shape[0])

print 'Building histogram...'
hist, bin_edges = np.histogram(s, density=True, bins=15)
plt.bar(bin_edges[:-1], hist, width=0.2)
plt.show()

from pymc import *

# Defining the model
def mymodel():
    s = np.loadtxt('data/dados_nonmixture.data')
    mean = Uniform( 'mean', lower = s.min(), upper = s.max() )
    precision = Uniform('precision', lower = 0.0001, upper = 1.5)
    process = Normal('process', mu = mean, tau = precision, observed = True, value = s)
    return locals()


# Model object
model = MCMC( mymodel() )

# Sampling procedure
model.sample( iter = 2000, thin = 10 , burn = 100)

model.stats()

# Checking the model variables
Matplot.plot(model)

graph = pymc.graph.graph(model)
graph.write_png("graph_nomix.png")

print 'Reading the data...'
s = np.loadtxt('data/dados_mixture.data')

print 'We have %d observations' %(s.shape[0])

print 'Building histogram...'
hist, bin_edges = np.histogram(s, density=True, bins=15)
plt.bar(bin_edges[:-1], hist, width=0.2)
plt.show()

# Definition of the model
def hip():
    
    # This is the data
    s = np.loadtxt('data/dados_mixture.data')
    
    # The Bernoulli parameter
    theta = Uniform("theta", lower = 0., upper = 1. )
    
    # Bernoulli process to mix the two distributions
    bern = Bernoulli("bern", p = theta, size = s.shape[0])
    
    # Means and variances for the two gaussian distributions
    mean1 = Uniform('mean1', lower = s.min(), upper = s.max() )
    mean2 = Uniform('mean2', lower = s.min(), upper = s.max() )
    std_dev = Uniform('std_dev', lower = 0., upper = 5.)
    
    # Constructing the mean value for the mixed process
    @deterministic(plot = False)
    def mean(bern = bern, mean1 = mean1, mean2 = mean2):
        return bern * mean1 + (1 - bern) * mean2
    
    @deterministic(plot = False)
    def precision(std_dev = std_dev):
        return 1.0 / ( std_dev * std_dev )
    
    # The process itself    
    process = Normal('process', mu = mean, tau = precision, value = s, \
                        observed = True)    
        
    return locals()

# Defining the model
model = MCMC( hip() )

# Steps
MCstep = 200
Nburns = 10000
Nsteps = MCstep*200 + Nburns

# Iteration process
model.sample( iter = Nsteps, burn = Nburns, thin = MCstep )

model.stats( variables = ['mean1', 'mean2'] )

# Checking the dynamics of mean1 and mean2
Matplot.plot( model.trace('mean1') )
Matplot.plot( model.trace('mean2') )

# Trying again with larger burn interval
model.sample( iter = 800*MCstep, burn = 300*MCstep, thin = MCstep )

Matplot.plot( model.trace('mean1') )
Matplot.plot( model.trace('mean2') )

Matplot.plot( model.trace('theta') )
Matplot.plot( model.trace('std_dev') )

model.stats(start = 300)

graph = pymc.graph.graph(model)
graph.write_png("graph_mix.png")